如何在Polars中高效使用over实现分组聚合为列表？

核心高效实现 #

要通过over窗口函数按分组收集列表，关键是在窗口内使用agg()聚合操作——这是最高效的方式，无需依赖性能低下的mapping_strategy="join"参数：

import polars as pl

df = pl.DataFrame(
    [pl.Series("id", ["a", "b", "a"]), pl.Series("x", [0, 1, 2])]
)

# 用over+agg生成分组列表并广播到每一行
result = df.with_columns(
    group_x_list=pl.col("x").agg().over("id")
)
print(result)

输出结果：

shape: (3, 3)
┌─────┬─────┬─────────────┐
│ id  ┆ x   ┆ group_x_list │
│ --- ┆ --- ┆ ---         │
│ str ┆ i64 ┆ list[i64]   │
╞═════╪═════╪═════════════╡
│ a   ┆ 0   ┆ [0, 2]      │
│ b   ┆ 1   ┆ [1]         │
│ a   ┆ 2   ┆ [0, 2]      │
└─────┴─────┴─────────────┘

为什么直接用pl.col("x").over("id")不生效？ #

直接调用pl.col("x").over("id")只会返回窗口内的逐行原值，不会触发聚合——over默认是窗口内的逐行映射逻辑，只有当你在窗口中使用聚合函数（如agg()、sum()、max()）时，才会对分组内的所有值计算并将结果广播到该分组的每一行。

不同实现的效率对比 #

1. over+agg()（推荐）：

   • 底层基于Polars的窗口优化机制，避免显式Join操作，性能与groupby+join持平。
   • 保留原DataFrame的行数，适合需要将分组列表关联到每一行的场景。

2. groupby+join替代方案：

   • 代码相对繁琐，但性能和over+agg()接近，适合仅需分组聚合结果（行数等于分组数）的场景。

3. mapping_strategy="join"的over用法：

   • 底层会执行隐式Join，性能损耗较大，不建议用于大数据集。

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