循环移动 Polars 数据框中的列

Question

如何获得 Polars 数据框中列的循环移位？

例如，假设我从这个数据框开始，

df = pl.DataFrame(dict(a=range(1, 10)))

我想添加一个列

b

，它是列

a

的循环移位，其中移位的位数是3。那么预期的输出是这样的：

shape: (9, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 7   │
│ 2   ┆ 8   │
│ 3   ┆ 9   │
│ 4   ┆ 1   │
│ 5   ┆ 2   │
│ 6   ┆ 3   │
│ 7   ┆ 4   │
│ 8   ┆ 5   │
│ 9   ┆ 6   │
└─────┴─────┘

Answer 1

这是一种使用

when

/

then

/

otherwise

表达式组合前向移位和后向移位以获得一个循环移位的方法：

(
    df
    .with_columns(
        b = pl.when(
            pl.col("a").shift(3).is_null()
        )
        .then(
            pl.col("a").shift(-1 * (pl.len() - 3))
        )
        .otherwise(
            pl.col("a").shift(3)
        )
    )
)

您的示例具有以下输出：

shape: (9, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 7   │
│ 2   ┆ 8   │
│ 3   ┆ 9   │
│ 4   ┆ 1   │
│ 5   ┆ 2   │
│ 6   ┆ 3   │
│ 7   ┆ 4   │
│ 8   ┆ 5   │
│ 9   ┆ 6   │
└─────┴─────┘

（请注意，由于

这个错误

，目前

then

分支中的移位大小必须写为

-1 * (pl.len() - 3)

而不是 3 - pl.len()。）