
本文详解如何在 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 表达式中实现了语义等价、性能可控、逻辑清晰的时序事件过滤,是构建用户行为分析、会话识别、异常检测等场景的可靠基础。

















