ComputeError with polars dataframe while trying to pass a column expressions in a simple shift() operation

Question

Sorry, I'm starting out in Polars. Is there a way to achieve the same functionality of the shift(n) function where n is a dataframe variable.

When I try:

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

df.with_columns(shift=pl.col("a").shift(pl.col("offset")))

It gives me the error message:

Traceback (most recent call last):

  Cell In[107], line 1
    df.with_columns(shift=pl.col("a").shift(pl.col("offset")))

  File d:\anaconda3\Lib\site-packages\polars\dataframe\frame.py:8242 in with_columns
    return self.lazy().with_columns(*exprs, **named_exprs).collect(_eager=True)

  File d:\anaconda3\Lib\site-packages\polars\lazyframe\frame.py:1816 in collect
    return wrap_df(ldf.collect(callback))

ComputeError: n must be a single value.

Thank you very much in advance.

Answer 1

shift operator only allows single integer as an input. However you can try to build simple when statement, because in your example you have only two unique shift values.

import polars as pl

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

df.with_columns(shift=(
                pl.when(pl.col("offset") == 1)
                  .then(pl.col("a").shift(1))
                  .otherwise(pl.col("a").shift(-1))    
               ))

And a bit broader option where you build expression using unique values from 'shift' column.

import polars as pl

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


# create initial condition that is always non valid
expr = pl.when(1!=1).then(1)

for k in df['offset'].unique():
    # for each unique shift value add `when` clause
    expr = expr.when(pl.col("offset") == k).then(pl.col("a").shift(k))

df.with_columns(shift=expr)