Python 如何在Polars中处理用户自定义函数的多行结果

该错误可能是由于UDF返回的复杂返回类型造成的.Polars可以快速推断出内部词典的数据类型.但是，稍后判断的词典可能具有不同数量的字段，因此会引发错误.

下面是一个更简单、可重现的例子.

import random

import polars as pl

def ufunc(x: int):
    return [
        {f"word_{i}": "elephant" for i in range(random.randint(1, 4))}
        for _ in range(random.randint(1, 4))
    ]

pl.DataFrame({"id": [1, 2]}).with_columns(pl.col("id").map_elements(ufunc))

最初，我认为正确地将returned_dtype参数设置为pl.Expr.map_elements就可以解决问题.然而，我没有找到一种方法来指定具有不同数量的字段的 struct 数据类型.

尽管如此，当将返回值包装到另一个列表中时，Polars似乎可以正确推断返回的数据类型.然后，我们可以使用.list.first获得第一个(也是唯一的)列表元素，使用pl.DataFrame.explode将行分解成行，最后取消内部 struct 的嵌套.

(
    pl.DataFrame({"id": [1, 2]})
    .with_columns(
        pl.col("id").map_elements(lambda x: [ufunc(x)]).list.first()
    )
    .explode("id")
    .unnest("id")
)

shape: (7, 4)
┌──────────┬──────────┬──────────┬──────────┐
│ word_0   ┆ word_1   ┆ word_2   ┆ word_3   │
│ ---      ┆ ---      ┆ ---      ┆ ---      │
│ str      ┆ str      ┆ str      ┆ str      │
╞══════════╪══════════╪══════════╪══════════╡
│ elephant ┆ elephant ┆ null     ┆ null     │
│ elephant ┆ elephant ┆ elephant ┆ elephant │
│ elephant ┆ null     ┆ null     ┆ null     │
│ elephant ┆ null     ┆ null     ┆ null     │
│ elephant ┆ elephant ┆ null     ┆ null     │
│ elephant ┆ elephant ┆ elephant ┆ null     │
│ elephant ┆ null     ┆ null     ┆ null     │
└──────────┴──────────┴──────────┴──────────┘