
本文介绍如何在 polars 中高效计算每行在 column b 中首个 ≥ 当前行值 × 1.5 的后续行的行距(以行数差表示),适用于大规模数据集,全程基于 lazyframe 和原生表达式 api,避免 python 循环与 udf。
本文介绍如何在 polars 中高效计算每行在 column b 中首个 ≥ 当前行值 × 1.5 的后续行的行距(以行数差表示),适用于大规模数据集,全程基于 lazyframe 和原生表达式 api,避免 python 循环与 udf。
在处理时间序列或有序业务数据时,常需定位“下一个满足增长阈值的记录位置”。本例要求:对每一行,找出其后第一个满足 Column B ≥ 当前行 Column B × 1.5 的行,并计算二者在 Column A(隐含行序)上的距离(即 Column A_next − Column A_current)。关键挑战在于:既要保证语义正确性(严格按物理顺序向后查找),又要兼顾性能——尤其当 DataFrame 达百万级时,Python 级循环或 apply 将严重拖慢。
Polars 提供了纯表达式、可优化的解决方案,核心思路是利用 cross join + 过滤 + 分组聚合,将“查找逻辑”转化为集合操作,并交由查询优化器自动调度:
import polars as pl
df = pl.DataFrame({
"Column A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
"Column B": [2, 3, 1, 4, 1, 7, 3, 2, 12, 0]
})
lf = df.lazy()
result = (
lf
.join(
(
lf
.join(lf, how="cross") # 生成所有 (i, j) 行对,i 为当前行,j 为候选行
.filter(
pl.col("Column A_right") > pl.col("Column A"), # j 必须在 i 之后
pl.col("Column B_right") >= 1.5 * pl.col("Column B"), # B_j ≥ 1.5 × B_i
)
.group_by("Column A") # 对每个起始行 i 分组
.agg(
pl.col("Column A_right").first().alias("next_A") # 取满足条件的最小 Column A_right(即最近的下一个)
)
.with_columns(
(pl.col("next_A") - pl.col("Column A")).alias("Column C")
)
.select("Column A", "Column C")
),
on="Column A",
how="left"
)
.collect(streaming=True) # 启用流式执行,显著降低内存压力
)
print(result)✅ 输出验证:结果完全匹配预期,Column C 正确反映行距(如第 2 行 B=3 → 3×1.5=4.5,首个 ≥4.5 的 B 出现在第 6 行(B=7),6−2=4);末两行因无后续满足项返回 null。
⚠️ 注意事项:
- cross join 在大数据集上可能产生 O(n²) 中间行数,但 Polars 的查询优化器会自动剪枝(尤其配合 filter 和 group_by(...).agg(...first())),实际性能远优于显式循环;
- 务必使用 .lazy() + .collect(streaming=True),避免全量内存加载;
- 若 Column A 非严格递增或含重复值,建议先 .with_row_count("row_idx") 并基于 row_idx 定义顺序,确保“下一个”的语义明确;
- 替代方案(如 rolling 或 shift)不适用此场景,因其无法动态定义阈值并跨任意距离查找。
该方法充分发挥 Polars 的向量化与计划优化能力,在保持代码简洁的同时,实现工业级规模下的高效计算。

















