如何用Ibis框架实现Polars中DataFrame的移位（Shift）操作？

在Polars中，我们可以通过以下代码为DataFrame生成移位列：

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),
)

初始DataFrame定义如下：

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的shift(n)是取当前行的前n行数据（向下偏移），对应Ibis中的lag(n)方法；而你使用的lead(n)是取当前行的后n行数据，方向完全相反。
• 未指定列名：原Polars代码为新列命名为a_1、a_2等，你的Ibis代码没有自定义列名，会生成默认命名（如a_lead_1），不符合期望输出的列名要求。

正确的Ibis实现代码如下：

import ibis

t = ibis.memtable(df)
t = t.mutate(
    a_1=ibis._['a'].lag(1),
    a_2=ibis._['a'].lag(2),
    b_1=ibis._['b'].lag(1),
    b_2=ibis._['b'].lag(2),
)
t.to_polars()

执行这段代码后，就能得到和Polars完全一致的输出结果。

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