Polars中借助Numba JIT实现多列迭代计算的方案问询

针对Polars v0.20.26中需要带状态迭代更新A、B、D三列的场景，推荐以下高效实现方案：

核心思路 #

利用Polars的map_batches接口结合Numba JIT批处理，将带状态的迭代逻辑封装在Numba函数中，一次性返回所有列的计算结果，既保留Polars的高效数据处理特性，又解决多返回值的问题。

具体实现 #

1. 定义Numba批处理函数 #

编写带状态的迭代计算函数，使用Numba的nopython=True模式编译，接近C级执行效率：

import polars as pl
from numba import jit
import numpy as np

@jit(nopython=True)
def iterative_calc(A_arr, B_arr, initial_val):
    n = len(A_arr)
    A_out = np.empty(n, dtype=np.float64)
    B_out = np.empty(n, dtype=np.float64)
    D_out = np.empty(n, dtype=np.float64)
    
    current_val = initial_val
    for i in range(n):
        A_out[i] = A_arr[i] * current_val
        B_out[i] = B_arr[i] * current_val
        D_out[i] = A_out[i] + B_out[i]
        current_val = D_out[i]
    
    return A_out, B_out, D_out

2. 结合Polars的map_batches执行计算 #

通过map_batches将Numba函数应用到DataFrame的批处理数据上，一次性生成更新后的A、B、D列：

# 构造示例数据
df = pl.DataFrame({
    "A": [1.0, 2.0, 3.0],
    "B": [4.0, 5.0, 6.0],
    "D": [0.0, 0.0, 0.0]
})

initial_value = 2.0

# 执行批量计算
result_df = df.map_batches(
    lambda batch: pl.DataFrame(
        {
            "A": a,
            "B": b,
            "D": d
        }
        for a, b, d in [iterative_calc(batch["A"].to_numpy(), batch["B"].to_numpy(), initial_value)]
    )
)

print(result_df)

方案优势 #

• 低开销：map_batches是Polars原生批处理接口，仅需一次Python到Numba的调用，避免了逐元素计算的开销。
• 高效计算：Numba的nopython模式编译后的函数执行效率接近纯C代码，远高于自定义C函数的Python调用开销。
• 完整保留结果：一次性返回A、B、D三列的计算值，无需重复计算或额外存储。

扩展优化 #

如果处理超大规模数据集，可结合Polars的scan_parquet等懒加载接口，配合map_batches实现流式计算，避免全量数据加载到内存。

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