如何在Polars中识别连续日期并为每组分配标识ID

解决方案代码 #

import polars as pl

df = pl.DataFrame({
    "date": ["2024-01-01", "2024-01-02", "2024-01-03", "2024-01-05", "2024-01-06", "2024-01-07"]
})
df = df.with_columns(pl.col("date").str.to_date())

# 生成连续日期组标识
result = df.with_columns(
    # 计算日期偏移与行索引的差值，连续日期该值固定
    _diff=(pl.col("date") - pl.col("date").min()).dt.days() - pl.int_range(0, pl.count())
).with_columns(
    # 用rle_id对差值分组，再加1让ID从1开始
    id=pl.col("_diff").rle_id() + 1
).drop("_diff")  # 移除临时计算列

print(result)

输出结果 #

shape: (6, 2)
┌────────────┬─────┐
│ date       ┆ id  │
│ ---        ┆ --- │
│ date       ┆ i64 │
╞════════════╪═════╡
│ 2024-01-01 ┆ 1   │
│ 2024-01-02 ┆ 1   │
│ 2024-01-03 ┆ 1   │
│ 2024-01-05 ┆ 2   │
│ 2024-01-06 ┆ 2   │
│ 2024-01-07 ┆ 2   │
└────────────┴─────┘

核心逻辑拆解 #

1. 计算日期偏移量：将每个日期与数据集最早日期做差，得到天数差。连续日期的天数差会按1递增。
2. 对比行索引：生成从0开始的行索引序列，用日期偏移量减去行索引。连续日期的差值会保持不变，出现日期中断时，差值会跳变（比如示例中2024-01-05的偏移量是4，行索引是3，差值为1，和前一组的0形成区分）。
3. 分组标识生成：用rle_id()方法对差值列分组，相同差值会被分配同一个基础ID，最后加1让标识从1开始，符合需求。

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