本文详解如何在 Polars 中高效实现类似 SQL COUNT_IF(...) OVER (PARTITION BY user_id ORDER BY t RANGE BETWEEN CURRENT ROW AND '5 MINUTES' FOLLOWING) 的逻辑，即筛选出「在相同用户下、当前记录之后（含/不含自身）5 分钟内至少发生一次事件 A」的所有原始记录，并提供可复用的模式化解决方案。
本文详解如何在 polars 中高效实现类似 sql `count_if(...) over (partition by user_id order by t range between current row and '5 minutes' following)` 的逻辑，即筛选出「在相同用户下、当前记录之后（含/不含自身）5 分钟内至少发生一次事件 a」的所有原始记录，并提供可复用的模式化解决方案。
在 Polars 中模拟 SQL 的 RANGE 窗口函数（如“未来 5 分钟内是否存在某事件”）不能直接通过 .over() + 表达式链完成——因为 Polars 的 rolling 表达式不支持嵌套在聚合上下文（如 .filter() 或 .over()）中；而 DataFrame 级的 .rolling() 又默认执行聚合，会丢失原始行粒度。解决这一矛盾的关键在于：先用 rolling(..., group_by=...) 构建带时间窗口的分组聚合结果，再通过 pl.all().first() 保留原始行信息，最后基于窗口内事件列表做逻辑判断并还原原始顺序。
以下是完整、可运行的实现方案（以筛选“5 分钟内后续出现 'A' 事件”的记录为例）：
import polars as pl
from datetime import datetime
# 示例数据（注意：t 列必须为 Datetime 类型）
df = pl.DataFrame({
"t": [
datetime(2025, 1, 1, 0, 0),
datetime(2025, 1, 1, 0, 0),
datetime(2025, 1, 1, 0, 0),
datetime(2025, 1, 1, 0, 1),
datetime(2025, 1, 1, 0, 3),
datetime(2025, 1, 1, 0, 3),
datetime(2025, 1, 1, 0, 4),
datetime(2025, 1, 1, 0, 5),
datetime(2025, 1, 1, 0, 5),
datetime(2025, 1, 1, 0, 6),
],
"user_id": ["Bart", "Aaron", "Charles", "Aaron", "Aaron", "Charles", "Bart", "Bart", "Aaron", "Aaron"],
"cat": ["C", "C", "B", "A", "C", "B", "A", "A", "B", "A"],
})
# ✅ 核心方案：保留原始行 + 时间窗口内事件检查
result = (
df
.with_row_index("i") # 步骤1：添加原始索引，用于最终排序还原顺序
.rolling(
index_column="t",
period="5m",
offset="0m",
group_by="user_id",
closed="both" # 包含当前行；若仅需“后续”（不含当前），请见下方说明
)
.agg(
pl.all().first(), # 步骤2：对每列取窗口内首行值 → 保留原始行内容
pl.col("cat").alias("_filt") # 步骤3：单独聚合 cat 列为列表，用于条件判断
)
.filter(pl.col("_filt").list.contains("A")) # 步骤4：窗口内存在 'A' 则保留该行
.sort("i") # 步骤5：按原始索引排序，恢复输入顺序
.drop(["_filt", "i"]) # 步骤6：清理临时列
)
print(result)⚠️ 关键注意事项：
closed="both"表示窗口包含当前行及其后 5 分钟内的所有行；若业务要求严格“后续”（即排除当前行本身），需改用closed="right"或保持"both"并手动切片：.filter(pl.col("_filt").list.slice(1, None).list.contains("A"))pl.all().first()是核心技巧：它确保每个窗口输出一行，且该行内容完全对应窗口起始位置的原始记录（因 Polarsrolling默认按index_column升序处理，且first取窗口左端点行）。group_by必须显式指定（如"user_id"），否则跨用户混算将导致逻辑错误。- 时间列
t必须是Datetime类型（推荐time_unit='us'），否则period解析可能失败。
延伸：SQL → Polars 窗口思维转换指南
| SQL 概念 | Polars 等效实现 | 说明 |
|----------|----------------|------|
| PARTITION BY X | group_by="X" in .rolling() | 分组维度必须由 DataFrame 级滚动指定 |
| ORDER BY Y | index_column="Y" | 时间/序号列必须作为 index_column，不可用 .sort() 替代 |
| RANGE BETWEEN CURRENT ROW AND INTERVAL '5m' FOLLOWING | period="5m", offset="0m", closed="both" | offset="0m" 表示从当前时刻开始，closed="both" 包含两端 |
| COUNT_IF(cond) > 0 | .agg(pl.col(...).alias("_tmp")).filter(pl.col("_tmp").list.contains(...)) | 将布尔条件转化为 list.contains() 检查 |
此模式不仅适用于事件存在性判断，还可扩展至：计算窗口内某事件首次出现时间、统计后续事件频次、提取后续首条匹配记录等场景，是 Polars 处理复杂时序依赖逻辑的基石范式。
