Polars：按分组计算窗口大小为2的滚动求和

原数据 #

import polars as pl

df = pl.DataFrame({'group': [1, 1, 1, 3, 3, 3, 4, 4], 'value': [1, 4, 2, 5, 3, 4, 2, 3]})

解决方案 #

要实现按分组计算窗口大小为2的滚动求和，只需结合group_by与rolling_sum方法，同时设置min_periods=1以保留每个分组的首个元素值：

result = df.group_by('group', maintain_order=True).agg(
    pl.col('value').rolling_sum(window_size=2, min_periods=1)
).explode('value')

print(result)

代码说明 #

• group_by('group', maintain_order=True)：按group列分组，且保持原数据的行顺序，保证输出结果与预期顺序一致。
• rolling_sum(window_size=2, min_periods=1)：对每个分组内的value列执行窗口大小为2的滚动求和，min_periods=1确保当窗口元素不足2个时（如分组的第一个元素），直接返回当前元素，而非空值。
• explode('value')：将分组聚合得到的列表型结果展开为单行数据，还原成原数据的行数结构。

输出结果 #

┌───────┐
│ value │
│ ---   │
│ i64   │
╞═══════╡
│ 1     │
│ 5     │
│ 6     │
│ 5     │
│ 8     │
│ 7     │
│ 2     │
│ 5     │
└───────┘

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