如何使用Polars基于现有列创建整数列：计算与列表最大不大于值的差值

因为points是已排序的整数列表，我们可以利用Polars的search_sorted方法高效完成需求：

1. 使用search_sorted找到每个bar元素在points中的插入位置（该位置是第一个大于bar元素的索引）
2. 插入位置减1，得到满足x <= y的最大x的索引
3. 通过索引取出对应的x，计算bar - x得到baz列

具体代码如下：

import polars as pl

# 构建示例DataFrame
df = pl.DataFrame({
    "foo": [86, 109, 109, 153, 153, 153, 160, 222, 225, 239],
    "bar": [11592, 2765, 4228, 4214, 7217, 11095, 1134, 5509, 10150, 4151]
})

points = [0, 1500, 3000, 4500, 6000, 7500, 9000, 10500, 12000]

# 添加baz列
df = df.with_columns(
    pl.col("bar").search_sorted(points, side="right").alias("idx") - 1,
).with_columns(
    (pl.col("bar") - pl.lit(points).take(pl.col("idx"))).alias("baz")
).drop("idx")

print(df)

运行后输出结果：

shape: (10, 3)
┌──────┬───────┬──────┐
│ foo  ┆ bar   ┆ baz  │
│ ---  ┆ ---   ┆ ---  │
│ i64  ┆ i64   ┆ i64  │
╞══════╪═══════╪══════╡
│ 86   ┆ 11592 ┆ 1092 │
│ 109  ┆ 2765  ┆ 1265 │
│ 109  ┆ 4228  ┆ 1228 │
│ 153  ┆ 4214  ┆ 1214 │
│ 153  ┆ 7217  ┆ 1217 │
│ 153  ┆ 11095 ┆ 595  │
│ 160  ┆ 1134  ┆ 1134 │
│ 222  ┆ 5509  ┆ 1509 │
│ 225  ┆ 10150 ┆ 1150 │
│ 239  ┆ 4151  ┆ 1151 │
└──────┴───────┴──────┘

补充说明 #

• search_sorted的side="right"参数确保返回第一个大于y的元素索引，减1后即可得到最大的x <= y的索引
• 若bar中存在小于points首个元素的场景，可额外添加边界判断逻辑避免索引越界

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