Используйте ibis-framework для расчета сдвигов (лагов) в фрейме данных

Скажем, я хочу сделать в Polars следующее:

df.with_columns(
    a_1 = pl.col('a').shift(1),
    a_2 = pl.col('a').shift(2),
    b_1 = pl.col('b').shift(1),
    b_2 = pl.col('b').shift(2),
)

Скажем, начиная с

import polars as pl

df = pl.DataFrame({'a': [1,3,2,4], 'b': [5,1,2,1]})

Итак, желаемый результат:

shape: (4, 6)
┌─────┬─────┬──────┬──────┬──────┬──────┐
│ a   ┆ b   ┆ a_1  ┆ a_2  ┆ b_1  ┆ b_2  │
│ --- ┆ --- ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64 ┆ i64  ┆ i64  ┆ i64  ┆ i64  │
╞═════╪═════╪══════╪══════╪══════╪══════╡
│ 1   ┆ 5   ┆ null ┆ null ┆ null ┆ null │
│ 3   ┆ 1   ┆ 1    ┆ null ┆ 5    ┆ null │
│ 2   ┆ 2   ┆ 3    ┆ 1    ┆ 1    ┆ 5    │
│ 4   ┆ 1   ┆ 2    ┆ 3    ┆ 2    ┆ 1    │
└─────┴─────┴──────┴──────┴──────┴──────┘

Как я могу сделать это с ibis?

Мне нужно сделать это, начиная с фрейма данных Polars df, поэтому мне нужно подключиться к нему с помощью Ibis. Глядя на документацию, я мало что могу понять с точки зрения того, что мне на самом деле следует делать.

Правильно ли сделать:

t = ibis.memtable(df)
t = t.mutate(
    ibis._['a'].lead(1),
    ibis._['a'].lead(2),
    ibis._['b'].lead(1),
    ibis._['b'].lead(2),
)
t.to_polars()