如何在Polars中实现按名称的滚动日期窗口索引计数？

需求说明 #

按Name分组，为每条记录生成过去X天内（含当前日期）该名称的出现顺序索引：

• 遇到新名称时，索引重置为1
• 当前日期超出窗口范围（与窗口内最早日期差超过X天）时，索引从1开始重新计数

以2天窗口为例，示例中John的2023-01-03条目，窗口包含2023-01-02和01-03的记录，因此索引从5开始（承接01-02的4条记录）。

解决方案代码 #

import polars as pl

WINDOW_DAYS = 2

df = pl.from_repr("""
┌─────────┬─────────────────────┬─────────┐
│ Name    ┆ Date                ┆ Counter │
│ ---     ┆ ---                 ┆ ---     │
│ str     ┆ datetime[ns]        ┆ i64     │
╞═════════╪═════════════════════╪═════════╡
│ John    ┆ 2023-01-01 00:00:00 ┆ 1       │
│ John    ┆ 2023-01-01 00:00:00 ┆ 2       │
│ John    ┆ 2023-01-01 00:00:00 ┆ 3       │
│ John    ┆ 2023-01-01 00:00:00 ┆ 4       │
│ John    ┆ 2023-01-02 00:00:00 ┆ 5       │
│ John    ┆ 2023-01-02 00:00:00 ┆ 6       │
│ John    ┆ 2023-01-02 00:00:00 ┆ 7       │
│ John    ┆ 2023-01-02 00:00:00 ┆ 8       │
│ John    ┆ 2023-01-03 00:00:00 ┆ 5       │
│ John    ┆ 2023-01-03 00:00:00 ┆ 6       │
│ New Guy ┆ 2023-01-01 00:00:00 ┆ 1       │
└─────────┴─────────────────────┴─────────┘
""")

# 生成滚动索引
result = df.with_columns(
    row_idx=pl.int_range(0, pl.count())
).group_by("Name").map_groups(
    lambda group: group.with_columns(
        Counter=pl.col("row_idx")
        .rank(method="ordinal")
        .over(
            pl.col("Date").filter(
                pl.col("Date") >= pl.col("Date") - pl.duration(days=WINDOW_DAYS)
            )
        )
    )
).drop("row_idx")

print(result)

代码说明 #

1. 生成行索引：先为每条记录添加唯一的row_idx，确保排序的准确性
2. 按名称分组处理：通过group_by.map_groups对每个名称组单独计算，避免跨名称干扰
3. 动态窗口筛选：在over窗口内用filter筛选出当前日期前X天内的所有记录，再用rank(method="ordinal")计算当前记录在窗口内的顺序索引

对之前尝试方法的问题分析 #

• rolling聚合方法：rolling.agg返回的是每个窗口的聚合列表，无法直接映射到每条记录的索引，不符合需求
• rolling_sum_by方法：未按Name分组导致跨名称计数错误，且计算的是累计总和而非顺序索引
• over(Name, Date)窗口：仅限定同一名称和同一日期的记录，无法覆盖过去X天的动态窗口范围

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