如何让Polars自定义函数返回多列？

方法1：让函数直接返回表达式列表 #

这是最直接的实现方式，自定义函数返回多个Polars表达式组成的列表，with_columns会自动将每个表达式作为新列添加到DataFrame中。

修改后的代码：

from polars.type_aliases import IntoExpr, IntoExprColumn
import polars as pl
import numpy as np

def _func(x: IntoExpr):
    x1 = x + 1
    x2 = x + 2
    # 返回带别名的表达式列表，直接传入with_columns即可生成新列
    return [x1.alias("test1"), x2.alias("test2")]
    
df = pl.DataFrame({"test": np.arange(1, 11)})
result = df.with_columns(_func(pl.col("test")))
print(result)

输出与预期一致：

shape: (10, 3)
┌──────┬───────┬───────┐
│ test ┆ test1 ┆ test2 │
│ ---  ┆ ---   ┆ ---   │
│ i32  ┆ i32   ┆ i32   │
╞══════╪═══════╪═══════╡
│ 1    ┆ 2     ┆ 3     │
│ 2    ┆ 3     ┆ 4     │
│ 3    ┆ 4     ┆ 5     │
│ 4    ┆ 5     ┆ 6     │
│ 5    ┆ 6     ┆ 7     │
│ 6    ┆ 7     ┆ 8     │
│ 7    ┆ 8     ┆ 9     │
│ 8    ┆ 9     ┆ 10    │
│ 9    ┆ 10    ┆ 11    │
│ 10   ┆ 11    ┆ 12    │
└──────┴───────┴───────┘

方法2：使用结构体并展开列 #

如果要保留返回结构体的写法，可以通过.struct.unnest()将结构体列拆分为多个独立列：

方式2.1：给结构体字段命名后展开 #

from polars.type_aliases import IntoExpr, IntoExprColumn
import polars as pl
import numpy as np

def _func(x: IntoExpr):
    x1 = x + 1
    x2 = x + 2
    # 创建带字段名的结构体
    return pl.struct({"test1": x1, "test2": x2})
    
df = pl.DataFrame({"test": np.arange(1, 11)})
# 直接展开结构体列，字段名会自动作为新列名
result = df.with_columns(_func(pl.col("test")).struct.unnest())
print(result)

方式2.2：先返回匿名结构体，再重命名展开 #

from polars.type_aliases import IntoExpr, IntoExprColumn
import polars as pl
import numpy as np

def _func(x: IntoExpr):
    x1 = x + 1
    x2 = x + 2
    return pl.struct([x1, x2])
    
df = pl.DataFrame({"test": np.arange(1, 11)})
# 展开后给列重命名
result = df.with_columns(
    _func(pl.col("test")).struct.unnest().alias(["test1", "test2"])
)
print(result)

原代码问题解析 #

你之前的写法错误在于：pl.struct()返回的是单个结构体列，而.alias(["test1", "test2"])试图给单个列设置多个名称，这不符合Polars的API逻辑。Polars中alias()一次只能给一个列重命名，拆分结构体列需要用.struct.unnest()方法。

内容的提问来源于stack exchange，提问作者Kevin Li