如何使用Polars计算同一ID截至当前记录的历史值聚合？

需求说明 #

现有事件数据集，同一id对应多条记录，需新增agg列，存储当前记录之前同id的value字段总和（首条记录的agg值为0）。已知PySpark可通过Window函数实现，需Polars的实现方案。

实现方案 #

在Polars中，可通过窗口分组+累加和+偏移的组合操作实现该需求，无需复杂自定义函数。核心思路是：先按id分组并按datetime排序，计算分组内的累加和后偏移一位，首条记录的空值填充为0。

代码示例 #

import polars as pl

# 构造示例数据集
df = pl.DataFrame({
    "id": [1, 1, 2, 2, 2],
    "datetime": ["2023-01-01", "2023-01-02", "2023-01-01", "2023-01-02", "2023-01-03"],
    "value": [10, 20, 5, 15, 25]
})

# 新增agg列
result_df = df.with_columns(
    pl.col("value")
    .cumsum()
    .shift(1)
    .fill_null(0)
    .over(pl.window("id", order_by="datetime"))
    .alias("agg")
)

print(result_df)

代码解释 #

• over(pl.window("id", order_by="datetime"))：指定按id分组，每组内按datetime排序，确保累加顺序符合时间逻辑
• cumsum()：计算分组内截至当前记录的value累加和
• shift(1)：将累加结果向后偏移一位，使当前记录对应之前所有记录的总和
• fill_null(0)：将每组首条记录的空值（偏移后无前置记录）填充为0，匹配需求要求

补充说明 #

目前Polars官方文档对这类带偏移的窗口聚合操作没有明确的针对性指引，Polars仓库的开放Issue #8976正在讨论该场景的优化及文档补充事宜。

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