Python Polars中按列分组并使用多参数自定义函数的实现问题

你的错误根源在于用struct.map_batches()结合over的方式不适合分组处理多列自定义函数，而且struct列在map_batches中无法直接用字符串索引访问字段。Polars中分组处理多列自定义函数更适合用map_groups（或group_by().apply()），直接操作每组的DataFrame。

针对你的需求，正确实现方式如下：

import polars as pl
import statsmodels.api as sm

lowess = sm.nonparametric.lowess

df = pl.DataFrame([
    pl.Series('x', ['a', 'a', 'a', 'b','b', 'b']),
    pl.Series('y', [1, 2, 3, 1, 2, 3]),
    pl.Series('z', [.2, .3, .5, .1, .3, .7])
])

# 用map_groups处理分组后的自定义函数
result = df.group_by('x').map_groups(
    lambda group: group.with_columns(
        # lowess返回二维数组，第二列是拟合值
        pl.Series('lowess_z', lowess(group['z'], group['y'], frac=0.5)[:, 1])
    )
)

print(result)

关键说明： #

• group_by('x').map_groups()：对每个分组的子DataFrame直接操作，可直接访问组内的y和z列，避免struct索引的问题。
• lowess函数返回(n, 2)格式的数组，第一列是原始y值，第二列是拟合后的z值，因此取[:,1]提取拟合结果，转为Polars Series作为新列。
• 若需多列分组，只需将group_by('x')改为group_by(['x', '其他列名'])即可。
• 注意：你原代码中frac=0.1在这个小数据集里不满足lowess的参数要求（需frac*n >= 2），示例中调整为frac=0.5避免警告。

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