Polars中用DataFrame列作shift()参数触发ComputeError的解决咨询

Polars的shift()方法仅支持固定的整数偏移量，无法直接传入列作为逐行变化的偏移值，这就是你遇到ComputeError: n must be a single value.报错的原因。

解决方案：通过索引计算实现逐行偏移 #

可以通过计算每行的目标索引，再使用take()方法提取对应位置的值，同时可选择处理索引越界的情况：

示例代码

import polars as pl

df = pl.DataFrame({
    "a": [1, 2, 3, 4],
    "offset": [1, -1, 1, -1]
})

# 实现带越界处理的逐行shift
result = df.with_columns(
    # 计算目标索引：当前行号 + 偏移量
    target_idx=pl.int_range(0, pl.count()) + pl.col("offset"),
    # 将索引限制在有效范围内（0到行数-1）
    valid_idx=pl.col("target_idx").clip(0, pl.count()-1)
).with_columns(
    # 仅当目标索引有效时取值，否则返回null
    shift=pl.when(pl.col("target_idx").is_between(0, pl.count()-1))
           .then(pl.col("a").take(pl.col("valid_idx")))
           .otherwise(None)
).drop("target_idx", "valid_idx")

print(result)

输出结果

shape: (4, 3)
┌─────┬────────┬───────┐
│ a   ┆ offset ┆ shift │
│ --- ┆ ---    ┆ ---   │
│ i64 ┆ i64    ┆ i64   │
╞═════╪════════╪═══════╡
│ 1   ┆ 1      ┆ 2     │
│ 2   ┆ -1     ┆ 1     │
│ 3   ┆ 1      ┆ 4     │
│ 4   ┆ -1     ┆ 3     │
└─────┴────────┴───────┘

简化版（无越界处理）

如果可以确保所有偏移量都不会导致索引越界，可省略越界判断步骤：

result = df.with_columns(
    shift=pl.col("a").take(pl.int_range(0, pl.count()) + pl.col("offset"))
)

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