How can I get a circular shift of a column in a Polars dataframe?
For example, say I start with this dataframe,
df = pl.DataFrame(dict(a=range(1, 10)))
and I want to add a column b that is a circular shift of column a, where the number of places shifted is 3. Then the expected output is this:
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 │
└─────┴─────┘
Here is an approach that uses a when/then/otherwise expression to combine a forward shift and a backward shift to get one circular shift:
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)
)
)
This has the following output for your example:
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 │
└─────┴─────┘
(Note that at the moment the shift size in the then branch has to be written as -1 * (pl.len() - 3) instead of 3 - pl.len(), due to this bug.)