如何在Polars数据框中实现列的循环移位？

假设我们有如下Polars数据框：

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

需要添加一列b，使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   │
└─────┴─────┘

以下是两种实现方法：

方法1：拼接头尾数据 #

循环移位n位（此处n=3），本质是把列的后n个元素移到前面，剩余元素接在后面。可以直接用tail()和head()截取对应部分，再拼接成新列：

n = 3
df = df.with_columns(
    b=pl.concat([pl.col("a").tail(n), pl.col("a").head(len(df) - n)])
)

方法2：通过索引选取（通用场景适用） #

构造循环移位后的索引序列，用take()方法按索引提取元素，这种方式更适合需要动态计算移位位数的场景：

n = 3
df = df.with_columns(
    b=pl.col("a").take(pl.arange(n, len(df)).append(pl.arange(0, n)))
)

这里pl.arange(n, len(df))获取从第n位到末尾的索引，pl.arange(0, n)获取前n位索引，拼接后的索引序列就是循环移位后的顺序，take()会按该顺序提取对应元素。

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