如何在Polars中将apply/map_rows的结果添加为新列？

我有一个宽DataFrame，需要对部分列应用自定义逻辑生成新列，目前已经得到包含目标值的单列DataFrame。尝试过多种.with_columns写法但没成功，而且没有行标识符时拼接操作不可靠。要求实现通用UDF，而非Polars表达式形式。

示例代码 #

import polars as pl

df = pl.DataFrame({
    "foo": [16, 28, 0 ],
    "bar": [None, 4,17 ],
    "yat": [41, 174,15 ],
    "tar": [None, 4,0 ],
})

def udf(row: tuple[float])->str: 
    '''This code is illustrative - meant to be a row-wise UDF'''
    return ' + '.join([f'{x}x{"^"+str(i) if i>0 else ""}' for i,x in enumerate(row) if x!=0 and x is not None])

# 当前生成单列DataFrame的写法
df.select('foo','bar', 'yat', 'tar').map_rows(udf).rename({'map':'poly'})

当前输出 #

shape: (3, 1)
┌────────────────────────────┐
│ poly                       │
│ ---                        │
│ str                        │
╞════════════════════════════╡
│ 16x + 41x^2                │
│ 28x + 4x^1 + 174x^2 + 4x^3 │
│ 17x^1 + 15x^2              │
└────────────────────────────┘

期望输出 #

shape: (3, 5)
┌─────┬──────┬─────┬──────┬────────────────────────────┐
│ foo ┆ bar  ┆ yat ┆ tar  ┆ poly                       │
│ --- ┆ ---  ┆ --- ┆ ---  ┆ ---                        │
│ i64 ┆ i64  ┆ i64 ┆ i64  ┆ str                        │
╞═════╪══════╪═════╪══════╪════════════════════════════╡
│ 16  ┆ null ┆ 41  ┆ null ┆ 16x + 41x^2                │
│ 28  ┆ 4    ┆ 174 ┆ 4    ┆ 28x + 4x^1 + 174x^2 + 4x^3 │
│ 0   ┆ 17   ┆ 15  ┆ 0    ┆ 17x^1 + 15x^2              │
└─────┴──────┴─────┴──────┴────────────────────────────┘

Pandas中的实现方式 #

df.assign(poly = lambda x: x.apply(udf, axis=1))

解决方案 #

可以通过pl.struct将目标列打包为结构体，再结合map_elements调用自定义UDF，最后用with_columns直接将新列添加到原DataFrame中，确保行对齐：

方式1：指定目标列 #

(df.with_columns(
   pl.struct('foo','bar','yat','tar')
     .map_elements(lambda x: udf(x.values()))
     .alias('poly'))
)

方式2：使用所有列 #

如果需要对DataFrame的所有列应用UDF，可以用pl.all()代替指定列名：

(df.with_columns(
   pl.struct(pl.all())
     .map_elements(lambda x: udf(x.values()))
     .alias('poly'))
)

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