Polars：按分组计算每年首个时间点的扩展窗口累计和

可以通过排序+全局累计和+按年分组取首个日期的组合方式实现需求，步骤清晰且符合Polars的高效风格：

import polars as pl

df = pl.DataFrame(
    {
        "date": [
            "2020-03-01",
            "2020-05-01",
            "2020-11-01",
            "2021-01-01",
            "2021-02-03",
            "2021-06-08",
            "2022-01-05",
            "2020-07-01",
            "2020-09-01",
            "2022-01-05",
            "2023-02-04",
        ],
        "group": [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
        "value": [1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4],
    },
).with_columns(pl.col("date").str.strptime(pl.Date))

# 核心处理逻辑
result = (
    df.sort(["group", "date"])
    .with_columns(
        # 计算每个group的全局累计和
        pl.col("value").cumsum().over("group").alias("total_cumsum"),
        # 提取日期对应的年份
        pl.col("date").dt.year().alias("year")
    )
    .group_by(["group", "year"])
    .agg(
        # 取每年的首个日期
        pl.col("date").min().alias("date"),
        # 匹配首个日期对应的累计和
        pl.col("total_cumsum")
        .filter(pl.col("date") == pl.col("date").min())
        .first()
        .alias("value")
    )
    .sort(["group", "date"])
    .drop("year")  # 移除中间列
)

print(result)

输出结果 #

┌────────────┬───────┬───────┐
│ date       ┆ group ┆ value │
│ ---        ┆ ---   ┆ ---   │
│ date       ┆ i64   ┆ i64   │
╞════════════╪═══════╪═══════╡
│ 2020-03-01 ┆ 1     ┆ 1     │
│ 2021-01-01 ┆ 1     ┆ 10    │
│ 2022-01-05 ┆ 1     ┆ 28    │
│ 2020-07-01 ┆ 2     ┆ 1     │
│ 2022-01-05 ┆ 2     ┆ 6     │
│ 2023-02-04 ┆ 2     ┆ 10    │
└────────────┴───────┴───────┘

步骤解释 #

1. 排序：先按group和date升序排列，确保每个分组内的时间顺序正确，累计和计算逻辑才成立。
2. 全局累计和：用cumsum().over("group")计算每个分组从第一条记录到当前记录的累计值，得到全局的累计序列。
3. 按年分组提取首个日期：
   • 提取每条记录的年份，按group+year分组；
   • 每组内取最小日期（即该年的首个时间点），并匹配该日期对应的累计和；
4. 整理结果：排序后移除中间生成的year列，得到最终格式。

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