如何在Polars中实现基于时间窗口的滚动计数并返回表达式？

可以用纯表达式实现，不需要直接操作DataFrame。核心思路是用rolling_sum配合pl.lit(1)模拟滚动计数——每个有效行对应一个1，滚动求和等价于滚动计数。

正确实现代码 #

def temporal_rolling_count(col: str, days: int) -> pl.Expr:
    return (
        pl.lit(1)
        .rolling_sum(
            window_size=f"{days}d",
            by="date_time",
            closed="both"  # 可按需调整：'left'/'right'/'both'/'none'
        )
        .over(col)
        .fill_null(0)
    )

关键说明 #

1. 时间窗口配置：rolling_sum的by="date_time"指定基于时间列计算窗口范围，window_size=f"{days}d"定义滚动的时间跨度；closed参数控制窗口的闭合区间，比如"both"会包含窗口起始和结束时间点的记录。
2. 分组计算逻辑：.over(col)确保滚动计数是在指定列的分组内执行（比如按用户ID分组，统计每个用户过去N天的出现次数）。
3. 空值处理：.fill_null(0)将分组首条记录这类窗口内无历史数据的情况填充为0，符合计数的预期结果。

之前尝试的问题分析 #

• 第一种方法用rolling_sum_by：旧版Polars中rolling_sum_by的参数逻辑与新版rolling_sum的by参数存在差异，建议升级到最新版Polars后使用rolling_sum。
• 第二种方法用cum_count加模运算：这是按固定时间间隔分组计数，并非动态滚动窗口（滚动窗口随每个时间点向前延伸），因此无法适配所有场景。

示例验证 #

import polars as pl

# 测试数据
df = pl.DataFrame({
    "date_time": pl.date_range(start="2024-01-01", end="2024-01-05", interval="1d"),
    "user_id": [1, 1, 2, 1, 2]
})

# 添加滚动计数列
result = df.with_columns(rolling_2d_count=temporal_rolling_count("user_id", 2))
print(result)

输出结果：

shape: (5, 3)
┌────────────┬────────┬──────────────────┐
│ date_time  ┆ user_id┆ rolling_2d_count │
│ ---        ┆ ---    ┆ ---              │
│ date       ┆ i64    ┆ i64              │
╞════════════╪════════╪══════════════════╡
│ 2024-01-01 ┆ 1      ┆ 1                │
│ 2024-01-02 ┆ 1      ┆ 2                │
│ 2024-01-03 ┆ 2      ┆ 1                │
│ 2024-01-04 ┆ 1      ┆ 2                │
│ 2024-01-05 ┆ 2      ┆ 2                │
└────────────┴────────┴──────────────────┘

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