如何在Polars中实现按连续日期重置的累计计数？

示例数据 #

import polars as pl

df = pl.DataFrame(
    {
        "date": [
            "2024-08-01",
            "2024-08-02",
            "2024-08-03",
            "2024-08-04",
            "2024-08-04",
            "2024-08-05",
            "2024-08-06",
            "2024-08-08",
            "2024-08-09",
        ],
        "type": ["A", "A", "A", "A", "B", "B", "B", "A", "A"],
    }
).with_columns(pl.col("date").str.to_date())

期望输出 #

shape: (9, 3)
┌────────────┬──────┬───────────────┐
│ date       ┆ type ┆ days_in_a_row │
│ ---        ┆ ---  ┆ ---           │
│ date       ┆ str  ┆ i64           │
╞════════════╪══════╪═══════════════╡
│ 2024-08-01 ┆ A    ┆ 1             │
│ 2024-08-02 ┆ A    ┆ 2             │
│ 2024-08-03 ┆ A    ┆ 3             │
│ 2024-08-04 ┆ A    ┆ 4             │
│ 2024-08-04 ┆ B    ┆ 1             │
│ 2024-08-05 ┆ B    ┆ 2             │
│ 2024-08-06 ┆ B    ┆ 3             │
│ 2024-08-08 ┆ A    ┆ 1             │
│ 2024-08-09 ┆ A    ┆ 2             │
└────────────┴──────┴───────────────┘

其中days_in_a_row计数器会在同类型日期间隔超过1天或类型切换时重置。

解决方案代码 #

result = df.with_columns(
    # 同类型内计算当前日期与前一行的差值
    date_diff=pl.col("date").diff().over("type"),
    # 标记需要重置计数器的位置：首次出现或间隔超1天
    reset_flag=pl.when(
        (pl.col("date_diff") > pl.duration(days=1)) | pl.col("date_diff").is_null()
    ).then(1).otherwise(0),
    # 生成连续日期段的分组ID
    group_id=pl.col("reset_flag").cumsum().over("type")
).with_columns(
    # 在(type+group_id)分组内计算连续天数
    days_in_a_row=pl.cum_count().over(["type", "group_id"])
).drop(["date_diff", "reset_flag", "group_id"])

print(result)

步骤解释 #

1. 计算日期差值：用diff().over("type")在每个类型分组内，得到当前日期与前一行日期的间隔，第一条记录为null。
2. 标记重置点：当日期间隔超过1天，或者是该类型的第一条记录时，标记为1，其余为0。这个标记用来识别计数器需要重置的位置。
3. 生成子分组ID：在每个类型分组内，对重置标记做累加，将连续的日期段划分为独立子分组（比如类型A的2024-08-01至04是一个子分组，2024-08-08至09是另一个）。
4. 计算连续天数：在type和group_id的组合分组内，用cum_count()生成从1开始递增的连续计数器。
5. 清理辅助列：删除中间生成的date_diff、reset_flag、group_id列，得到最终结果。

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