如何用Polars优雅获取目标值上下最接近的两行数据

方法一：基于值匹配的单次过滤 #

先计算diff列，再直接定位到需要的两个strike值，通过isin一次筛选出目标行，避免多次构建DataFrame和合并操作：

import polars as pl

data = {
    "strike": [5,10,15,20,25,30],
    "target": [16] * 6,
}

df = pl.DataFrame(data).with_columns(diff=pl.col("strike") - pl.col("target"))

# 获取小于target的最大strike和大于target的最小strike
lower_strike = df.filter(pl.col("diff") < 0).select(pl.col("strike").max()).item()
upper_strike = df.filter(pl.col("diff") > 0).select(pl.col("strike").min()).item()

# 一次筛选出两行结果
result = df.filter(pl.col("strike").is_in([lower_strike, upper_strike]))
print(result)

输出结果：

shape: (2, 3)
┌────────┬────────┬──────┐
│ strike ┆ target ┆ diff │
│ ---    ┆ ---    ┆ ---  │
│ i64    ┆ i64    ┆ i64  │
╞════════╪════════╪══════╡
│ 15     ┆ 16     ┆ -1   │
│ 20     ┆ 16     ┆ 4    │
└────────┴────────┴──────┘

方法二：利用分组聚合（适配多目标场景） #

如果需要处理多组不同的target值，可以通过分组聚合一次性获取每组的上下边界行：

result = (
    df
    # 标记行属于"低于目标"或"高于目标"组
    .with_columns(group=pl.when(pl.col("diff") < 0).then("lower").when(pl.col("diff") > 0).then("upper"))
    # 排除diff等于0的行（如果存在）
    .filter(pl.col("group").is_not_null())
    # 按分组聚合，取lower组的最大值、upper组的最小值
    .group_by("group", "target")
    .agg(
        strike=pl.col("strike").max().when(pl.col("group") == "lower").otherwise(pl.col("strike").min()),
        diff=pl.col("diff").max().when(pl.col("group") == "lower").otherwise(pl.col("diff").min())
    )
    .drop("group")
    .sort("strike")
)
print(result)

方法三：基于位置切片（已排序数据最优解） #

如果strike列已经按升序排列，可以直接定位到第一个大于目标值的行索引，然后切片取前一行和当前行，这种方法无需多次过滤，效率最高：

# 找到第一个strike大于target的行索引
first_upper_idx = df.filter(pl.col("strike") > pl.col("target")).select(pl.int_range(0, pl.count()).min()).item()
# 切片取前一行和当前行
result = df.slice(first_upper_idx - 1, 2)
print(result)

对比原实现的优势 #

原方法需要重复构建DataFrame并计算diff列，最终还要用vstack合并结果，上述方法：

• 只计算一次diff列，避免重复计算
• 无需合并操作，通过单次过滤/聚合/切片直接得到结果
• 更符合Polars向量化操作的风格，代码简洁易读

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