如何在Polars中用自定义函数填充指定列空值（替代Pandas实现）

我有一段Pandas代码：

df[col] = (
            df[col]
            .fillna(method="ffill", limit=1)
            .apply(lambda x: my_function(x))
        )

想要改写成Polars代码，尝试了以下实现：

df = df.with_columns(
            pl.col(col)
            .fill_null(strategy="forward", limit=1)
            .map_elements(lambda x: my_function(x))
        )

但代码未正常工作：向前填充策略生效，但自定义函数未处理剩余缺失值。

测试代码 #

import polars as pl
import pandas as pd
import numpy as np

df_polars = pl.DataFrame(
    {"A": [1, 2, None, None, None, None, 4, None], "B": [5, None, None, None, None, 7, None, 9]}
)

df_pandas = pd.DataFrame(
    {"A": [1, 2, None, None, None, None, 4, None], "B": [5, None, None, None, None, 7, None, 9]}
)

last_valid_data: int


def my_function(x):
    global last_valid_data
    if x == None or np.isnan(x):
        result = last_valid_data * 10
    else:
        last_valid_data = x
        result = x
    return result


col = "A"

last_valid_data = df_pandas[col][0]
df_pandas[col] = df_pandas[col].fillna(method="ffill", limit=1).apply(lambda x: my_function(x))

last_valid_data = df_polars[col][0]
df_polars = df_polars.with_columns(
    pl.col(col).fill_null(strategy="forward", limit=1).map_elements(lambda x: my_function(x))
)

Pandas期望输出 #

A    B
0   1.0  5.0
1   2.0  NaN
2   2.0  NaN
3  20.0  NaN
4  20.0  NaN
5  20.0  7.0
6   4.0  NaN
7   4.0  9.0

Polars当前输出 #

┌──────┬──────┐
│ A    ┆ B    │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 1    ┆ 5    │
│ 2    ┆ null │
│ 2    ┆ null │
│ null ┆ null │
│ null ┆ null │
│ null ┆ 7    │
│ 4    ┆ null │
│ 4    ┆ 9    │
└──────┴──────┘

问题根源 #

Polars的map_elements默认参数skip_nulls=True，会自动跳过null值，不会把这些缺失值传入my_function处理，所以剩余的null值没被替换成预期的last_valid_data *10。

修改后的代码 #

给map_elements加上skip_nulls=False参数，强制所有值（包括null）都进入自定义函数处理：

last_valid_data = df_polars[col][0]
df_polars = df_polars.with_columns(
    pl.col(col)
    .fill_null(strategy="forward", limit=1)
    .map_elements(lambda x: my_function(x), skip_nulls=False)
)

验证结果 #

修改后运行代码，Polars输出会和Pandas的期望输出一致：

┌──────┬──────┐
│ A    ┆ B    │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 1    ┆ 5    │
│ 2    ┆ null │
│ 2    ┆ null │
│ 20   ┆ null │
│ 20   ┆ null │
│ 20   ┆ 7    │
│ 4    ┆ null │
│ 4    ┆ 9    │
└──────┴──────┘

更优实现（避免全局变量） #

使用全局变量存在并行处理时的安全隐患，推荐用Polars原生的窗口状态函数替代，更符合向量化处理逻辑：

def polars_style_process(col_name):
    return (
        pl.col(col_name)
        .fill_null(strategy="forward", limit=1)
        .scan(lambda s: s.running_map(lambda acc, x: x if x is not None else acc * 10, initial=s.first()))
    )

df_polars = df_polars.with_columns(polars_style_process(col))

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