如何在Polars DataFrame中传递整行至Python函数生成新列

问题根源 #

你原代码用了map_batches处理结构体列，但这个方法是批量处理整个列的结构体集合，而非逐行传递单个行。你在lambda里把整个结构体Series传给test2，导致函数内部取row["Group"]拿到的是整个Group列的Series，返回的字符串结果和Polars预期的类型不匹配，触发报错。

正确解决方案（逐行传递整行） #

使用map_elements替代map_batches，这个方法会对结构体列的每个元素（即单行数据）单独处理，完全符合你传递整行的需求：

import polars as pl

def test2(auth, row):
    c = row["Group"]
    d = row["Val"]
    return "{}-{}-{}".format(c, str(d), auth)

df = pl.DataFrame({
    'Group': ['A', 'B', 'C', 'D', 'E'],
    'Val': [1001, 1002, 1003, 1004, 1005]
})

auth_token = "xxxxxxxxx"

df = df.with_columns(
    pl.struct(pl.all())
    .map_elements(lambda row: test2(auth_token, row), return_dtype=pl.String)
    .alias("response")
)

print(df)

运行后会得到预期结果：

shape: (5, 3)
┌───────┬──────┬───────────────────┐
│ Group ┆ Val  ┆ response          │
│ ---   ┆ ---  ┆ ---               │
│ str   ┆ i64  ┆ str               │
╞═══════╪══════╪═══════════════════╡
│ A     ┆ 1001 ┆ A-1001-xxxxxxxxx  │
│ B     ┆ 1002 ┆ B-1002-xxxxxxxxx  │
│ C     ┆ 1003 ┆ C-1003-xxxxxxxxx  │
│ D     ┆ 1004 ┆ D-1004-xxxxxxxxx  │
│ E     ┆ 1005 ┆ E-1005-xxxxxxxxx  │
└───────┴──────┴───────────────────┘

关于strict=False的说明 #

• 这个参数是设置在map_elements或map_batches方法中的，用于允许Polars接受类型不严格匹配的返回值。
• 但你这个场景根本不需要它，因为用对map_elements后类型是完全匹配的，报错的核心原因是处理方法选错了，而非类型不严格。

批量处理的替代方案 #

如果想保留map_batches（适合处理大数据集时批量优化），需要手动遍历结构体列中的每个单行元素：

df = df.with_columns(
    pl.struct(pl.all())
    .map_batches(lambda struct_series: pl.Series(
        [test2(auth_token, row) for row in struct_series]
    ), return_dtype=pl.String)
    .alias("response")
)

对比指定列的实现 #

你之前用指定列的方式需要在lambda里拆分结构体字段，而整行传递的方案不需要修改函数参数，后续如果DataFrame列名或数量变化，函数和处理逻辑都不需要调整，灵活性更高。

内容的提问来源于stack exchange，提问作者configure.ng