Polars中对List[f64]类型列执行分组后逐元素聚合的最优实现方法

刚好遇到过类似的需求，Polars其实提供了非常优雅且高效的方式来实现这种列表元素级的聚合，完全不用手动逐个索引处理，尤其适合你这种256长度的列表场景，下面是两种最优实现方案：

方法1：使用pl.element()结合聚合函数 #

这是最直接的方式，pl.element()是Polars专门用来指代列表中对应位置元素的占位符，在agg中使用时，它会自动对齐分组内所有列表的相同索引位置，然后对这些元素应用聚合函数（sum、mean等）：

import polars as pl

df = pl.DataFrame(
    {
        "Case": ["case1", "case1"],
        "List": [[1, 2, 3], [4, 5, 6]],
    }
)

# 分组后对列表执行元素级求和
sum_result = df.group_by("Case").agg(
    pl.col("List").agg(pl.element().sum())
)

print(sum_result)

执行后会得到你期望的结果：

shape: (1, 2)
┌───────┬────────────┐
│ Case  ┆ List       │
│ ---   ┆ ---        │
│ str   ┆ list[i64]  │
╞═══════╪════════════╡
│ case1 ┆ [5, 7, 9]  │
└───────┴────────────┘

如果需要计算均值，只需要把sum()换成mean()即可：

mean_result = df.group_by("Case").agg(
    pl.col("List").agg(pl.element().mean())
)
# 输出结果为 [[2.5, 3.5, 4.5]]

方法2：使用pl.list.eval() #

另一种等价的方式是用pl.list.eval()，它专门用于对列表列进行元素级的计算，同样可以结合pl.element()实现聚合：

sum_result = df.group_by("Case").agg(
    pl.col("List").list.eval(pl.element().sum())
)

这个方法和方法1的效果完全一致，底层都是Polars的向量化操作，性能上没有差异，你可以根据自己的习惯选择。

重要注意事项 #

• 确保分组内的所有列表长度一致：Polars无法对长度不同的列表执行元素级聚合，会直接抛出错误，你的场景中每个列表长度都是256，刚好符合要求。
• 性能优势：这两种方法都是完全向量化的，底层由Polars的优化引擎处理，比手动遍历索引的Python循环效率高几个数量级，即使处理大规模数据也能保持高效。

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