Polars中按列应用UDF的高效方法及结果解包问题求助

1. 修正map_batches的返回值类型 #

你当前代码产生嵌套结构的核心原因是：map_batches接收列Series后，lambda直接返回了numpy数组，Polars会把这个数组当作单个元素封装进新的Series，最终导致每列变成仅含一个数组元素的1行Series。

只需要把返回的numpy数组包装成Polars Series，就能得到和原DataFrame形状一致的结果：

df_baseline_series = df.select(
    pl.all().map_batches(lambda y: pl.Series(get_baseline(y.to_numpy())))
)

处理后df_baseline_series会保持原2000×n的结构，无需额外解包。

2. 批量处理整个DataFrame（更高效） #

如果baseline_dt函数支持直接处理二维numpy数组（即同时处理所有列），可以跳过逐列map_batches，直接转换整个DataFrame为numpy数组处理，再转回Polars：

# 假设baseline_dt输入为(样本数, 特征数)的二维数组，返回同形状结果
processed_np = baseline_dt(df.to_numpy(), wavelet='qshift3', level=6, max_iter=1)
df_baseline = pl.DataFrame(processed_np, columns=df.columns)

这种方式减少了Polars和numpy之间的多次转换开销，性能会更优，尤其当n较大时。

3. 自定义向量化UDF（进阶） #

如果需要重复使用这个逻辑，可以把get_baseline封装成Polars的向量化UDF，让代码更简洁：

import polars as pl

@pl.api.register_expr_namespace("baseline")
class BaselineNamespace:
    def __init__(self, expr: pl.Expr):
        self._expr = expr
    
    def dt(self, wavelet='qshift3', level=6, max_iter=1):
        return self._expr.map_batches(
            lambda s: pl.Series(baseline_dt(s.to_numpy(), wavelet=wavelet, level=level, max_iter=max_iter))
        )

# 使用时直接调用
df_baseline = df.select(pl.all().baseline.dt())

这种方式让代码更具可读性，适合长期维护。

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