
本文介绍一种纯 Polars 表达式 API 的高效方案,通过交叉连接(cross join)与分组聚合,在大规模 DataFrame 中快速计算每行到首个满足 B ≥ 当前行 B × 1.5 的后续行的行距(以索引差表示),避免 Python 循环,支持流式执行。
本文介绍一种纯 polars 表达式 api 的高效方案,通过交叉连接(cross join)与分组聚合,在大规模 dataframe 中快速计算每行到首个满足 `b ≥ 当前行 b × 1.5` 的后续行的行距(以索引差表示),避免 python 循环,支持流式执行。
在处理大规模时序或顺序敏感数据时,常需定位“下一个满足阈值条件的元素位置”。本例中,目标是为每行计算:从当前行开始向后查找,首个满足 Column B ≥ 当前行 Column B × 1.5 的行与当前行在 Column A(作为逻辑索引)上的距离(即 Column A_right - Column A)。直接遍历效率低且难以并行,而 Polars 提供了基于关系代数的高性能替代方案。
核心思路是将问题转化为「自连接 + 条件过滤 + 每组取最早匹配」:
- 使用 how="cross" 构建所有 (当前行, 候选行) 组合;
- 筛选满足 候选行.Column A > 当前行.Column A(确保向后查找)且 候选行.Column B >= 1.5 * 当前行.Column B 的配对;
- 对每个当前行(按 Column A 分组),取 Column A_right 最小的匹配项(即最近的满足条件的后续行);
- 计算距离并左连接回原表,缺失匹配则自动为 null。
以下是完整可运行代码(推荐始终在 LazyFrame 上操作以启用查询优化和流式处理):
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")
.filter(
pl.col("Column A_right") > pl.col("Column A"),
pl.col("Column B_right") >= 1.5 * pl.col("Column B"),
)
.group_by("Column A")
.agg(
pl.col("Column A_right").first().alias("next_A") # 取最小 Column A_right(即最近匹配)
)
.select(
"Column A",
(pl.col("next_A") - pl.col("Column A")).alias("Column C")
)
),
on="Column A",
how="left"
)
.collect(streaming=True)
)
print(result)✅ 关键优势说明:
- ✅ 纯表达式驱动:全程使用 Polars 原生表达式,无 Python UDF 或 .apply(),可被查询优化器深度优化;
- ✅ 流式执行(streaming=True):对超大数据集友好,避免内存峰值;
- ✅ 逻辑清晰可扩展:如需改为“首个 ≥ 当前值 + 50”(绝对增量),仅需将 1.5 * pl.col("Column B") 替换为 pl.col("Column B") + 50;
- ⚠️ 注意事项:交叉连接复杂度为 O(n²),虽经 filter 和 group_by(...).agg(...first()) 剪枝,但对千万级行仍建议确认硬件资源;若 Column A 是严格递增整数索引,亦可考虑用 pl.int_range() + arr.eval() 等窗口技巧进一步优化,但本解法通用性更强、语义更直观。
最终输出与预期完全一致:第 9、10 行因后续无满足 B ≥ 12×1.5=18 或 0×1.5=0(注意 0 ≥ 0 成立,但无后续行)的记录,故 Column C 为 null。

















