Shift(n) should accept a varying n

[lmocsi]

Shift(n) should accept a varying n

[mcrumiller]

import polars as pl
from polars import col

df = pl.DataFrame({
    "a": ["a", "a", "a", "b", "b", "c", "c"],
    "ii": [1, 3, 5, 7, 9, 11, 13],
})


def shift_many(value_col, shift_col):
    return col(value_col).gather(
        (pl.int_range(pl.len()).alias("len") + col(shift_col))
        .clip(0, pl.len() - 1)
    )


print(df.with_columns(shift_many("a", "ii").alias("shifted")))

shape: (7, 3)
┌─────┬─────┬─────────┐
│ a   ┆ ii  ┆ shifted │
│ --- ┆ --- ┆ ---     │
│ str ┆ i64 ┆ str     │
╞═════╪═════╪═════════╡
│ a   ┆ 1   ┆ a       │
│ a   ┆ 3   ┆ b       │
│ a   ┆ 5   ┆ c       │
│ b   ┆ 7   ┆ c       │
│ b   ┆ 9   ┆ c       │
│ c   ┆ 11  ┆ c       │
│ c   ┆ 13  ┆ c       │
└─────┴─────┴─────────┘

[ritchie46]

What do you mean a varying n. A column can only be shifted once?

[lmocsi]

I mean that n can be different in every row. See provided example.

[ritchie46]

I think you just want to do a gather in such a case.