本文详解如何在 Polars 中高效实现类似 SQL COUNT_IF(...) OVER (PARTITION BY ... ORDER BY t RANGE BETWEEN CURRENT ROW AND '5 MINUTES' FOLLOWING) 的逻辑——即保留那些在同一用户未来 5 分钟内（含/不含当前行）至少发生一次事件 A 的原始记录，并提供可复用的模式化写法与关键注意事项。
本文详解如何在 polars 中高效实现类似 sql `count_if(...) over (partition by ... order by t range between current row and '5 minutes' following)` 的逻辑——即保留那些在**同一用户未来 5 分钟内（含/不含当前行）至少发生一次事件 a** 的原始记录，并提供可复用的模式化写法与关键注意事项。
在 Polars 中实现“某条记录是否被后续特定事件在指定时间窗口内覆盖”这类条件过滤，不能直接套用 SQL 的 QUALIFY + RANGE 窗口函数语法，因为 Polars 的表达式层（Expr）不支持基于时间范围的滚动窗口聚合与行级布尔判断的混合操作（例如 .over().rolling() 会报错 window expression not allowed in aggregation）。但通过合理组合 DataFrame.rolling()、agg() 和 list 操作，我们能精准还原该语义，且保持原始行粒度。
核心思路：以滚动窗口聚合为桥梁，重建行级布尔标签
Polars 的 rolling(..., group_by=...) 是唯一支持按时间范围（period, offset, closed）分组聚合的 API。虽然它天然产生聚合结果，但我们可通过以下三步将其“映射回原始行”：
-
添加行索引：使用
.with_row_index("i")记录原始顺序； -
时间窗口内聚合：对每个
(user_id, t)滚动窗口，聚合出该窗口内所有cat值（pl.col("cat").alias("_filt")），同时用pl.all().first()保留该窗口对应的第一行原始字段（即当前行的快照）； -
行级过滤 + 恢复顺序：基于
_filt列（类型为list[str]）判断是否含"A"，再按i排序还原原始位置。
以下是完整可运行代码（基于示例数据）：
import polars as pl
import datetime
df = pl.DataFrame({
"t": [
datetime.datetime(2025, 1, 1, 0, 0),
datetime.datetime(2025, 1, 1, 0, 0),
datetime.datetime(2025, 1, 1, 0, 0),
datetime.datetime(2025, 1, 1, 0, 1),
datetime.datetime(2025, 1, 1, 0, 3),
datetime.datetime(2025, 1, 1, 0, 3),
datetime.datetime(2025, 1, 1, 0, 4),
datetime.datetime(2025, 1, 1, 0, 5),
datetime.datetime(2025, 1, 1, 0, 5),
datetime.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"],
})
# ✅ 包含当前行的逻辑：检查 [t, t+5m] 内是否有 'A'
result_inclusive = (
df
.with_row_index("i")
.rolling(
index_column="t",
period="5m",
offset="0m",
group_by="user_id",
closed="both" # 关键：包含当前行和未来5分钟内所有行
)
.agg(
pl.all().first(), # 保留该窗口对应原始行的全部字段（取首行）
pl.col("cat").alias("_filt") # 聚合出窗口内所有 cat 值列表
)
.filter(pl.col("_filt").list.contains("A")) # 行级布尔过滤
.sort("i") # 恢复原始顺序
.drop(["_filt", "i"]) # 清理辅助列
)
print("✅ 包含当前行（[t, t+5m]）的结果：")
print(result_inclusive)
⚠️ 关键注意事项与变体处理
-
closed="both"vsclosed="right"：-
"both"（默认）包含当前行及未来窗口内所有行 → 对应 SQLBETWEEN CURRENT ROW AND ...； - 若需严格排除当前行（仅看后续事件），则必须用
closed="both"（因closed="right"在 Polars 中不支持offset="0m"的精确右开区间），然后对_filt列手动切片：.filter(pl.col("_filt").list.slice(1, None).list.contains("A"))list.slice(1,)跳过第一个元素（即当前行），确保只检查后续事件。
-
性能提示：
rolling(..., group_by=...)底层会先按group_by分组再排序，因此要求index_column（如t）为数值或时间类型。若数据量极大，建议提前对(user_id, t)排序并设置maintain_order=False（但此处需索引恢复顺序，故不可省略with_row_index）。-
泛化心智模型：
SQL 的OVER (PARTITION BY X ORDER BY Y RANGE BETWEEN A AND B)在 Polars 中没有直接等价物，但可统一转化为：*
df.rolling(Y, period=..., offset=..., group_by=X, closed=...)→ 聚合生成中间 list 列 → 用 `list.方法（contains,any,all,slice`）做行级判断**。
这是处理“时间范围依赖型行过滤”的标准范式。
最终，该方案完全避免了 .sql() 方法，在纯 Polars 表达式中实现了语义等价、性能可控、逻辑清晰的时序事件过滤，是构建用户行为分析、会话识别、异常检测等场景的可靠基础。
