Python Polars定制函数返回多列

发布于03月13日

_func用于返回两列:

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

def _func(x: IntoExpr):
    x1 = x+1
    x2 = x+2
    return pl.struct([x1, x2])
    
df = pl.DataFrame({"test": np.arange(1, 11)})
df.with_columns(
    _func(pl.col("test")).alias(["test1", "test2"])
)

我try 使用pl.struct来包装返回值，但不起作用.

预期输出:

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列，这样我就不必专门为每列alias.

在您的例子中，问题是在最后您有一个同名的Expr序列，所以您需要在执行之前的某个时刻重命名它们.

解决方案可能取决于您计划为列使用哪种名称. 您可以这样做:

def _func(x: IntoExpr):
    x1 = x+1
    x2 = x+2
    return x1, x2

df.with_columns(
    x.alias(f"test{i+1}") for i,x in enumerate(_func(pl.col("test")))
)

# alternatevely
# df.with_columns(
#    x.name.suffix(f"{i+1}") for i,x in enumerate(_func(pl.col("test")))
#)

┌──────┬───────┬───────┐
│ 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    │
└──────┴───────┴───────┘

对于自定义名称，您可以使用zip:

f.with_columns(
    x.alias(n) for x, n in zip(_func(pl.col('test')), ["a","b"])
)

┌──────┬─────┬─────┐
│ test ┆ a   ┆ b   │
│ ---  ┆ --- ┆ --- │
│ 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  │
└──────┴─────┴─────┘

或者，您可以使用with_columns()接受**named_exprs作为参数的事实，将表达式列表转换为DICT并解压缩:

df.with_columns(
    **dict(zip(['a','b'], _func(pl.col('test'))))
)

┌──────┬─────┬─────┐
│ test ┆ a   ┆ b   │
│ ---  ┆ --- ┆ --- │
│ 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  │
└──────┴─────┴─────┘