如何在Polars中实现当前值与前一个非空值的差值计算？

给定如下Polars DataFrame：

import polars as pl

df = pl.DataFrame({
    "a": [1,15,None,20,None]
})

需要生成新列b，满足以下规则：

• 若列a的值为null，则列b的值也为null；
• 若列a有值，则用当前值减去前一个非null值（第一个非空值对应的b为0）。

期望输出：

expected = pl.from_repr("""
┌──────┬──────┐
│ a    ┆ b    │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 1    ┆ 0    │
│ 15   ┆ 14   │ # b = 15 - 1
│ null ┆ null │
│ 20   ┆ 5    │ # b = 20 - 15
│ null ┆ null │
└──────┴──────┘
""")

此前尝试过整列聚合取最后值的方案，不符合需求；也试过rolling方法，但空值无规律无法确定周期，因此寻求解决方法。

通过向前填充非null值结合位移计算即可实现需求，代码如下：

import polars as pl

df = pl.DataFrame({
    "a": [1,15,None,20,None]
})

result = df.with_columns(
    b=pl.when(pl.col("a").is_not_null())
        .then(pl.col("a") - pl.col("a").fill_null(strategy="forward").shift(1).fill_null(pl.col("a")))
        .otherwise(None)
)

print(result)

运行结果与期望输出一致：

┌──────┬──────┐
│ a    ┆ b    │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 1    ┆ 0    │
│ 15   ┆ 14   │
│ null ┆ null │
│ 20   ┆ 5    │
│ null ┆ null │
└──────┴──────┘

代码逻辑说明 #

• pl.col("a").fill_null(strategy="forward")：向前填充a列的空值，得到每个位置对应的前一个非null值；
• .shift(1)：将填充后的序列向下位移一位，此时第一个位置变为空值；
• .fill_null(pl.col("a"))：把位移后的第一个空值替换为原列a的第一个值，确保第一个非空值的差值为0；
• pl.when(...).then(...).otherwise(None)：仅在原列a非空时计算差值，空值位置直接设为null。

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