如何优化Polars DataFrame行级函数执行效率并满负载利用CPU

当前代码存在两个核心问题：

• CPU利用率低：map_elements调用Python自定义函数时受GIL限制，只能单线程执行，导致CPU负载仅30%，这是耗时久的主要原因。
• 额外开销大：逐行打包struct再解包的方式存在不必要的序列化/反序列化开销，且逐行调用Python函数的成本被放大。
• 代码错误：你的my_complex_function返回的字典存在重复键{prefix_name}_result_B，会导致后一个值覆盖前一个，需先修正。

1. 优先用Polars矢量化操作重构函数 #

这是最优解——完全绕过Python GIL，利用Polars原生多线程引擎拉满CPU负载，同时消除逐行调用的开销。如果my_complex_function的逻辑能转换成Polars矢量化API实现，速度提升最明显：

# 重构为矢量化版本，直接处理整个Series
def vectorized_complex_function(
    param_A: pl.Series, param_B: pl.Series, param_C: pl.Series,
    param_D: float, param_E: float, param_F: float, prefix_name: str
) -> pl.DataFrame:
    # 用Polars矢量化API实现原逻辑（示例）
    result_A = param_A * param_B + param_D
    result_B = param_C + param_E - param_F
    result_C = (param_A + param_C) / 2
    
    # 直接返回DataFrame，避免后续unnest操作
    return pl.DataFrame({
        f"{prefix_name}_result_A": result_A,
        f"{prefix_name}_result_B": result_B,
        f"{prefix_name}_result_C": result_C
    })

# 调用方式：直接传入列，无需打包struct
df = df.with_columns(vectorized_complex_function(
    df['float_param_A'],
    df['float_param_B'],
    df['float_param_C'],
    float_param_D,
    float_param_E,
    float_param_F,
    prefix_name
))

2. 无法矢量化时，用多进程绕过GIL #

如果函数逻辑无法用Polars矢量化API重构，就用多进程充分利用CPU核心：

from concurrent.futures import ProcessPoolExecutor

# 打包所有需要的参数为可序列化的元组列表
inputs = [
    (row['float_param_A'], row['float_param_B'], row['float_param_C'],
     float_param_D, float_param_E, float_param_F, prefix_name)
    for row in df.to_dicts()
]

# 用进程池并行执行函数
with ProcessPoolExecutor() as executor:
    results = list(executor.map(lambda args: my_object.my_complex_function(*args), inputs))

# 将结果合并到原DataFrame
results_df = pl.DataFrame(results)
df = df.hstack(results_df)

3. 减少序列化开销 #

当前的struct.map_elements逐行序列化数据开销大，优化方式：

• 直接传递列而非打包struct，避免逐行序列化。
• 若必须用struct，改用map_batches处理整个批次的struct，减少序列化次数。

完全不是最优。原因：

• 单线程执行浪费CPU资源，导致负载低、耗时久。
• 逐行打包/解包struct带来不必要的序列化开销。
• 逐行调用Python函数的开销被放大，进一步拖慢速度。

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