Polars 中使用 join_asof 实现分段映射与差值计算

本文介绍如何在 polars 中高效生成新列 `baz`：对 dataframe 的 `bar` 列中每个值，从预排序列表 `points` 中查找**不超过该值的最大元素**，再计算二者之差。核心方案是利用 `join_asof` 的 `backward` 策略，全程基于表达式链式操作，无需 python 循环或 `apply`。

在 Polars 中处理此类“查找最近下界（floor lookup）+ 差值”任务时，最高效、最符合声明式风格的方式是使用 join_asof —— 它专为有序键的近似连接设计，性能远超逐行 apply 或 map_dict。

✅ 解决方案：join_asof + backward 策略

假设原始 DataFrame 为 df，且 points = [0, 1500, 3000, 4500, 6000, 7500, 9000, 10500, 12000]（已严格升序），执行以下链式操作：

points = [0, 1500, 3000, 4500, 6000, 7500, 9000, 10500, 12000]
df_points = pl.DataFrame({"point": points}).set_sorted("point")

result = (
    df.sort("bar")  # 关键：left_on 列需升序（join_asof 要求）
    .join_asof(df_points, left_on="bar", right_on="point")  # 默认 strategy="backward"
    .with_columns(baz=pl.col("bar") - pl.col("point"))
    .drop("point")
    .sort("foo")  # 恢复原始行序（可选）
)

? 原理说明： join_asof(..., strategy="backward") 对左表每行 bar=y，在右表 point 中查找满足 point ≤ y 的最大 point 值（即 floor）。 因 df_points 已通过 .set_sorted("point") 标记为有序，Polars 可跳过排序开销，直接二分查找，时间复杂度为 O(n log m)。 最终 baz = bar - point 即为所求差值。

⚠️ 注意事项

• 必须保证 points 已升序，并调用 .set_sorted("point") 显式声明；否则 join_asof 行为未定义或报错。
• join_asof 要求左表 left_on 列（此处为 "bar"）在连接前已排序（.sort("bar") 不可省略），否则结果不正确。
• 若 bar 中存在小于 points[0]（即 0）的值，point 将为 null，导致 baz 也为 null。如需兜底（例如设为 bar 自身），可改用：

  .with_columns(baz=pl.col("bar") - pl.col("point").fill_null(0))

✅ 输出验证

结果与预期一致（部分截取）：

┌─────┬───────┬──────┐
│ foo ┆ bar   ┆ baz  │
│ --- ┆ ---   ┆ ---  │
│ i64 ┆ i64   ┆ i64  │
╞═════╪═══════╪══════╡
│ 86  ┆ 11592 ┆ 1092 │  ← 11592 - 10500 = 1092
│ 109 ┆ 2765  ┆ 1265 │  ← 2765 - 1500 = 1265
│ 160 ┆ 1134  ┆ 1134 │  ← 1134 - 0 = 1134
└─────┴───────┴──────┘

此方法兼具高性能、可读性与可扩展性，是 Polars 生态中处理分段基准映射类问题的标准范式。