如何用Polars统计每个事件前指定时长内的同类事件数量？

要实现每个事件发生前N天内同类事件的数量统计，不能直接用cum_count()（它会累计所有历史事件），可以利用Polars的rolling_count滚动计数函数，结合分组和时间窗口实现需求，且具备良好的扩展性。

解决方案代码 #

import polars as pl
from datetime import date, timedelta

# 原始数据
df = pl.DataFrame(
    data = {
            "Event":["Rain","Sun","Rain","Sun","Rain","Sun","Rain","Sun"],
            "Date":[
                date(2022,1,1),
                date(2022,1,2),
                date(2022,1,2),
                date(2022,1,3),
                date(2022,1,3),
                date(2022,1,5),
                date(2022,1,5),
                date(2022,1,8)
            ]
    }
)

# 定义时间窗口（可按需修改，比如7天）
window_days = 2

# 核心逻辑：排序→滚动计数→恢复原顺序
result = (
    df.with_row_index("idx")  # 添加索引保留原始顺序
    .sort(["Event", "Date"])  # 按事件类型和日期排序，保证滚动窗口正确计算
    .with_columns(
        pl.col("Date")
        .rolling_count(
            by="Date",
            period=f"{window_days}d",  # 时间窗口范围
            closed="left",  # 窗口不包含当前日期，只统计之前的事件
            group_by="Event"  # 按事件类型分组，仅统计同类事件
        )
        .alias("cum_count")
    )
    .sort("idx")  # 恢复原始数据顺序
    .drop("idx")  # 移除临时索引列
)

print(result)

输出结果 #

shape: (8, 3)
┌───────┬────────────┬───────────┐
│ Event ┆ Date       ┆ cum_count │
│ ---   ┆ ---        ┆ ---       │
│ str   ┆ date       ┆ u32       │
╞═══════╪════════════╪═══════════╡
│ Rain  ┆ 2022-01-01 ┆ 0         │
│ Sun   ┆ 2022-01-02 ┆ 0         │
│ Rain  ┆ 2022-01-02 ┆ 1         │
│ Sun   ┆ 2022-01-03 ┆ 1         │
│ Rain  ┆ 2022-01-03 ┆ 2         │
│ Sun   ┆ 2022-01-05 ┆ 1         │
│ Rain  ┆ 2022-01-05 ┆ 1         │
│ Sun   ┆ 2022-01-08 ┆ 0         │
└───────┴────────────┴───────────┘

扩展性说明 #

1. 调整时间窗口：只需修改window_days变量，比如设置为7即可统计最近7天的同类前置事件。
2. 复杂过滤逻辑：如果需要添加更多过滤条件（比如特定事件属性），可以在rolling_count前结合filter对数据预处理。
3. 高效处理大数据：rolling_count是Polars优化的内置函数，比手动遍历或条件计数的效率更高，适合大规模数据集。

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