如何流式将多输出计算密集型函数应用到Polars数组列？

问题背景 #

需要将基于Numpy的计算密集型函数应用到Polars DataFrame的pl.Array(...)类型列上，因数据量常超出内存限制，必须采用流式处理方式。当前使用iter_slices分片处理后拼接结果的方案可行，但希望利用map_batches的批量处理能力，却受限于其看似不支持多返回值的问题，询问是否有更优方案，或是否当前方案可满足需求直至实现Rust UDF。

现有方案（iter_slices） #

通过iter_slices将大DataFrame拆分为小分片，逐个处理后拼接结果，示例代码如下：

import polars as pl
import numpy as np

filter1 = np.array([-1, 1, -1, 1, 1])
filter2 = np.array([0.3, 0.3, -0.3, 0.3, 0.3])

def compute(x):
    a = np.dot(x, filter1)
    b = np.dot(x, filter2)
    return np.sin(a - b), np.cos(a + b)

# 构造测试数据
df = pl.DataFrame(
    {"a": [np.arange(5.0) + i for i in range(5)]},
    schema={"a": pl.Array(pl.Float64, 5)}
)

# 流式分片处理
dfs = []
for df_iter in df.iter_slices(10000):
    peak_x, peak_y = compute(df_iter["a"].to_numpy())
    dfs.append(pl.DataFrame({"c": peak_x, "d": peak_y}))
df_out = pl.concat(dfs)

优缺点：

• 优点：逻辑直观，完全流式处理，内存占用可控；
• 缺点：需手动管理分片与结果拼接，代码冗余，未充分利用Polars的批量API优化。

更优方案（map_batches实现多返回值） #

map_batches并非不支持多返回值，只需让处理函数返回一个包含多列的Polars DataFrame即可，Polars会自动将其与原数据（或单独作为结果）合并，且天然支持流式处理。示例代码如下：

import polars as pl
import numpy as np

filter1 = np.array([-1, 1, -1, 1, 1])
filter2 = np.array([0.3, 0.3, -0.3, 0.3, 0.3])

def compute_batch(s: pl.Series) -> pl.DataFrame:
    # 将Series中的Array列转为二维numpy数组
    x = s.to_numpy()
    a = np.dot(x, filter1)
    b = np.dot(x, filter2)
    peak_x = np.sin(a - b)
    peak_y = np.cos(a + b)
    # 返回包含多列的DataFrame
    return pl.DataFrame({"c": peak_x, "d": peak_y})

# 构造测试数据
df = pl.DataFrame(
    {"a": [np.arange(5.0) + i for i in range(5)]},
    schema={"a": pl.Array(pl.Float64, 5)}
)

# 流式批量处理
df_out = df.map_batches(compute_batch)

优缺点：

• 优点：无需手动管理分片与拼接，代码更简洁；Polars自动处理流式逻辑，内存效率更高；
• 缺点：需适配函数接收Polars Series并返回DataFrame的格式。

结论 #

1. 若暂不打算实现Rust UDF，map_batches返回DataFrame的方案是更优选择，既保留流式处理能力，又符合Polars的API设计，代码更简洁高效；
2. 原iter_slices方案完全可行，若已适配现有业务逻辑，可继续使用直至需要更高性能的Rust UDF；
3. Rust UDF会是性能最优的方案，但开发成本较高，适合对性能有极致要求的场景。

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