Python Polars基于值范围的条件连接实现及无限上限处理

常规区间的实现方式 #

首先初始化示例数据：

import polars as pl

# 第一个DataFrame
df1 = pl.DataFrame({
    "date": ["2022-01-01", "2022-01-02", "2022-01-03"],
    "value": [3, 7, 12]
})

# 第二个区间分类DataFrame
df2 = pl.DataFrame({
    "category": ["A", "B", "C"],
    "lower": [0, 5, 10],
    "upper": [5, 10, 15]
})

利用Polars的cut函数可以高效实现区间匹配，该函数专门用于数值分箱并映射标签：

# 提取区间边界和分类标签
bins = df2["lower"].to_list() + [df2["upper"].max()]
labels = df2["category"].to_list()

# 新增category列
result = df1.with_columns(
    category=pl.col("value").cut(
        bins=bins,
        labels=labels,
        include_lowest=True  # 确保左区间闭合（如0<=value<5）
    )
)

print(result)

执行后将得到目标结果：

┌────────────┬───────┬──────────┐
│ date       ┆ value ┆ category │
│ ---        ┆ ---   ┆ ---      │
│ str        ┆ i64   ┆ str      │
╞════════════╪═══════╪══════════╡
│ 2022-01-01 ┆ 3     ┆ A        │
│ 2022-01-02 ┆ 7     ┆ B        │
│ 2022-01-03 ┆ 12    ┆ C        │
└────────────┴───────┴──────────┘

处理上限为无限大的情况 #

如果分类C的上限是无限大（即匹配所有value >=10的行），只需修改分箱边界，将最后一个边界设为float('inf')即可：

# 修改区间分类DataFrame，C的upper设为无限大
df2_infinite = pl.DataFrame({
    "category": ["A", "B", "C"],
    "lower": [0, 5, 10],
    "upper": [5, 10, float('inf')]
})

# 提取边界和标签
bins_infinite = df2_infinite["lower"].to_list() + [float('inf')]
labels_infinite = df2_infinite["category"].to_list()

# 生成结果
result_infinite = df1.with_columns(
    category=pl.col("value").cut(
        bins=bins_infinite,
        labels=labels_infinite,
        include_lowest=True
    )
)

print(result_infinite)

这种方式同样能正确匹配所有区间，包括大于等于10的数值。

另一种高效方法：不等连接（join_asof） #

如果区间数量较多，join_asof也是高效的选择，适合大规模数据：

# 先对区间DataFrame按lower排序（join_asof要求右侧DataFrame排序）
df2_sorted = df2.sort("lower")

# 执行不等连接，匹配value >= lower且value < upper的区间
result_join = df1.join_asof(
    df2_sorted,
    left_on="value",
    right_on="lower",
    strategy="forward"
).filter(pl.col("value") < pl.col("upper")).drop("lower", "upper")

print(result_join)

处理无限大上限时，只需去掉filter中的value < upper判断，或者将df2的C类upper设为极大值（如10**18）即可。

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