本文介绍使用 Polars 的 join_where（推荐）或交叉连接 + 过滤两种方式，为每行查找其后首个满足“price ≥ 当前行 limit”条件的最小 index 值，并安全填充 null，兼顾性能与可读性。
本文介绍使用 polars 的 join_where（推荐）或交叉连接 + 过滤两种方式，为每行查找其后首个满足“price ≥ 当前行 limit”条件的最小 index 值，并安全填充 null，兼顾性能与可读性。
在数据处理中，常需对每一行执行“向后查找”逻辑：即在当前行之后的所有行中，找到第一个满足某数值条件（如 price >= limit）且具有最小索引值的记录。传统循环或 explode + group_by 方式易导致内存膨胀与代码冗长。Polars 提供了更声明式、向量化且高效的解决方案。
✅ 推荐方案：join_where（实验性但高性能）
join_where 是 Polars 0.20+ 引入的专用操作，专为“带条件的自连接”设计，避免全量笛卡尔积，显著提升性能和内存效率：
import polars as pl
df_1 = pl.DataFrame({
'name': ['Alpha', 'Alpha', 'Alpha', 'Alpha', 'Alpha'],
'index': [0, 3, 4, 7, 9],
'limit': [12, 18, 11, 5, 9],
'price': [10, 15, 12, 8, 11]
})
out = (
df_1.join(
df_1
.join_where(
df_1.select('index', 'price'), # 右表仅需 index 和 price
pl.col('index_right') > pl.col('index'), # 索引必须严格大于当前行
pl.col('price_right') >= pl.col('limit') # 价格不低于当前 limit
)
.group_by('index')
.agg(pl.col('index_right').min().alias('min_index')),
on='index',
how='left'
)
)
该方案分三步完成：
-
条件自连接：将原表与自身子集（含
index,price）按index_right > index和price_right >= limit关联； -
聚合取最小：按原始
index分组，取所有匹配项中最小的index_right； -
左连接回填：将结果以
index为键合并回原表，自动处理无匹配时的null。
⚠️ 注意：
join_where目前标记为实验性（experimental），生产环境建议关注 Polars 官方文档更新；若需稳定 API，可选用下方替代方案。
? 替代方案：交叉连接 + 显式过滤（兼容性强）
当 join_where 不可用时，可用 how='cross' 搭配 filter 实现等效逻辑（注意：时间复杂度为 O(n²)，小数据集适用）：
out_fallback = (
df_1.join(
df_1
.join(df_1.select('index', 'price'), how='cross')
.filter(
pl.col('index_right') > pl.col('index'),
pl.col('price_right') >= pl.col('limit')
)
.group_by('index')
.agg(pl.col('index_right').min().alias('min_index')),
on='index',
how='left'
)
)
此写法语义清晰，兼容所有 Polars 版本，但需警惕数据规模——若原始 DataFrame 行数达万级，交叉连接可能引发内存压力。
? 关键要点与最佳实践
-
索引语义明确：示例中的
index列是业务索引（非 Polars 默认行号），因此必须显式参与比较，不可依赖.row_number()。 -
null 处理自然：未匹配行在
group_by(...).agg(...)后自动缺失，经left join即得null，无需额外fill_null()。 -
避免重复计算：切勿在
filter中重复引用未 select 的列（如limit不在右表中，故左表提供）。 -
性能对比提示：对 10k 行数据，
join_where通常比 cross + filter 快 3–5 倍，且内存占用更低。
最终输出完全符合预期：
shape: (5, 5) ┌───────┬───────┬───────┬───────┬───────────┐ │ name ┆ index ┆ limit ┆ price ┆ min_index │ │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ i64 ┆ i64 ┆ i64 ┆ i64 │ ╞═══════╪═══════╪═══════╪═══════╪═══════════╡ │ Alpha ┆ 0 ┆ 12 ┆ 10 ┆ 3 │ │ Alpha ┆ 3 ┆ 18 ┆ 15 ┆ null │ │ Alpha ┆ 4 ┆ 11 ┆ 12 ┆ 9 │ │ Alpha ┆ 7 ┆ 5 ┆ 8 ┆ 9 │ │ Alpha ┆ 9 ┆ 9 ┆ 11 ┆ null │ └───────┴───────┴───────┴───────┴───────────┘
掌握这两种模式，你便能优雅、高效地解决各类“向前/向后滚动条件查找”问题，真正发挥 Polars 声明式查询的强大表达力。
