如何在Pandas中并行化不可向量化的求根函数？

你的场景是：拥有包含a、b、c三列的多行DataFrame，需要为每行参数调用scipy.optimize.root(foo, 0.0, args=(a,b,c))求根，当前apply(axis=1)的逐行循环速度过慢，以下是几种可行的性能优化方案：

1. 用multiprocessing实现多进程并行 #

Python的multiprocessing模块可利用多核CPU并行处理任务，避开GIL（全局解释器锁）的限制，适合CPU密集型的求根计算。示例代码：

from scipy.optimize import root
import pandas as pd
from multiprocessing import Pool

# 定义单独的行处理函数，接收参数元组
def process_row(args):
    a, b, c = args
    return root(foo, 0.0, args=(a, b, c))

# 将DataFrame的a、b、c列转换为元组列表
params_list = list(zip(df['a'], df['b'], df['c']))

# 创建进程池，进程数建议设为CPU核心数
with Pool(processes=4) as pool:
    results = pool.map(process_row, params_list)

# 将结果存入DataFrame
df['root_result'] = results

2. 用concurrent.futures.ProcessPoolExecutor简化多进程逻辑 #

concurrent.futures封装了更简洁的并行接口，代码可读性更高：

from scipy.optimize import root
import pandas as pd
from concurrent.futures import ProcessPoolExecutor

def process_row(args):
    a, b, c = args
    return root(foo, 0.0, args=(a, b, c))

params_list = list(zip(df['a'], df['b'], df['c']))

with ProcessPoolExecutor(max_workers=4) as executor:
    results = list(executor.map(process_row, params_list))

df['root_result'] = results

3. 尝试向量化优化（优先级最高） #

如果能修改foo函数使其支持数组输入（即同时处理多组a、b、c参数），可以直接用向量化运算替代循环，这是效率最高的方案。部分scipy优化器原生支持向量化输入，示例如下（假设foo已适配数组）：

from scipy.optimize import root

# 一次性传入所有行的参数数组
result = root(foo, [0.0]*len(df), args=(df['a'].values, df['b'].values, df['c'].values))
# 提取每行的根结果
df['root_result'] = result.x

4. 用swifter自动选择最优执行方式 #

swifter库会自动判断函数是否适合向量化，若不适合则自动切换到并行处理，无需手动区分场景：

import pandas as pd
import swifter
from scipy.optimize import root

def process_row(row):
    return root(foo, 0.0, args=(row['a'], row['b'], row['c']))

df['root_result'] = df.swifter.apply(process_row, axis=1)

注意事项 #

• 多进程方案中，foo函数需要能被序列化（pickle），若函数包含无法序列化的内部状态或对象，需调整实现逻辑。
• 进程数不要超过CPU核心数，否则会增加进程切换开销，反而降低效率。
• 向量化是最优优化方向，优先尝试修改foo支持数组输入。

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