如何用Polars原生表达式实现跨DataFrame的区间过滤与聚合？

核心思路 #

抛弃逐行迭代的map_rows，改用Polars原生向量化表达式API，利用引擎的批量处理能力提升性能。以下提供两种适用不同场景的方案：

方案一：向量化广播匹配（中小数据量） #

直接通过广播比较实现区间匹配，再求和，代码简洁且高效：

import polars as pl

# 示例数据
df1 = pl.DataFrame({
    "start": [1, 5, 9],
    "end": [4, 8, 12]
})

df2 = pl.DataFrame({
    "idx": [2, 3, 6, 7, 10, 11],
    "values": [10, 20, 30, 40, 50, 60]
})

# 生成结果
result = df1.with_columns(
    sum_values=pl.sum(
        pl.when((df2["idx"] >= pl.col("start")) & (df2["idx"] <= pl.col("end")))
        .then(df2["values"])
        .otherwise(0)
    )
)

print(result)

输出结果：

shape: (3, 3)
┌───────┬─────┬────────────┐
│ start ┆ end ┆ sum_values │
│ ---   ┆ --- ┆ ---        │
│ i64   ┆ i64 ┆ i64        │
╞═══════╪═════╪════════════╡
│ 1     ┆ 4   ┆ 30         │
│ 5     ┆ 8   ┆ 70         │
│ 9     ┆ 12  ┆ 110        │
└───────┴─────┴────────────┘

方案二：二分查找+前缀和（大数据量） #

当df2数据量极大时，广播会占用过多内存，此时先对df2排序并计算前缀和，再通过二分查找定位区间边界，用前缀和差值得到求和结果，内存占用更低：

# 先对df2按idx排序并计算前缀和
df2_sorted = df2.sort("idx").with_columns(cumsum=pl.col("values").cumsum())

result = df1.with_columns(
    # 定位区间左右边界的位置
    start_pos=pl.col("start").search_sorted(df2_sorted["idx"], side="left"),
    end_pos=pl.col("end").search_sorted(df2_sorted["idx"], side="right")
).with_columns(
    # 通过前缀和差值计算区间和
    sum_values=pl.when(pl.col("end_pos") > 0)
    .then(df2_sorted["cumsum"][pl.col("end_pos")-1])
    .otherwise(0) - pl.when(pl.col("start_pos") > 0)
    .then(df2_sorted["cumsum"][pl.col("start_pos")-1])
    .otherwise(0)
).drop("start_pos", "end_pos")

print(result)

性能优势说明 #

• map_rows是Python级别的逐行迭代，每一行都要单独执行查询，性能瓶颈明显。
• 原生表达式API利用Polars的向量化引擎，一次性完成所有行的条件判断与计算，避免了循环开销，性能提升数倍至数十倍。

内容的提问来源于stack exchange，提问作者ignoring_gravity