如何基于匹配键与范围重叠高效连接两个Polars DataFrame？

问题背景 #

现有两个Polars DataFrame：

• df：包含ID与Reference Value列
• df2：包含ID、上下参考范围（Lower/Upper）及Map Value列

需要实现：当ID匹配且Reference Value处于对应上下范围内时，为df添加对应的Map Value。原循环when-then的方法在数据量大时会触发栈溢出，需替换为可扩展的高效方案。

核心原因：循环when-then的弊端 #

循环生成when-then条件分支会让Polars构建过于庞大的表达式树，大数据量下直接触发栈溢出，且这种写法完全没利用Polars的向量化优势，性能极低。

推荐高效方案 #

方案一：直接Join+过滤（最优选择） #

利用Polars的向量化Join操作，先按ID关联两个DataFrame，再过滤出符合范围条件的行，全程无循环，性能拉满。

示例代码

import polars as pl

# 构造示例数据
df = pl.DataFrame({
    "ID": ["A", "A", "B", "C", "C"],
    "Reference Value": [15, 25, 30, 5, 12]
})

df2 = pl.DataFrame({
    "ID": ["A", "A", "B", "C"],
    "Lower": [10, 20, 25, 0],
    "Upper": [20, 30, 35, 10],
    "Map Value": ["Low_A", "High_A", "Match_B", "Low_C"]
})

# 执行匹配
result = (
    df.join(df2, on="ID", how="left")
    .filter(pl.col("Reference Value").is_between(pl.col("Lower"), pl.col("Upper")))
    .select(df.columns + ["Map Value"])
)

print(result)

说明

• 用left join可以保留df中所有行，无匹配的行Map Value会显示为null；如果只需要保留有匹配的行，换成inner join即可。
• Polars的Join操作经过高度优化，支持超大数据量，不会出现栈溢出问题。

方案二：分组广播（适合特殊数据分布） #

如果df2中每个ID对应的范围规则数量极少，可以按ID分组后，将对应规则广播到组内处理，减少单次Join的数据量。

示例代码

result = (
    df.group_by("ID", maintain_order=True)
    .map_groups(
        lambda group: group.join(
            df2.filter(pl.col("ID") == group["ID"][0]),
            on="ID",
            how="left"
        ).filter(pl.col("Reference Value").is_between(pl.col("Lower"), pl.col("Upper")))
        .select(df.columns + ["Map Value"])
    )
)

print(result)

特殊场景处理 #

如果同一个ID+Reference Value匹配多个Map Value（比如范围重叠），会返回多行，可根据需求处理：

# 取每个匹配组合的第一个Map Value
result = (
    df.join(df2, on="ID", how="left")
    .filter(pl.col("Reference Value").is_between(pl.col("Lower"), pl.col("Upper")))
    .group_by(df.columns)
    .agg(pl.col("Map Value").first())
)

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