Python multiprocess.Pool.map无法处理大型数组，求DataFrame并行处理优化方案

嘿，我碰到过一模一样的问题——用multiprocessing的Pool.map处理大型DataFrame分片时，要么慢得离谱，要么直接报错对吧？这大多是因为进程间传递大对象的序列化开销太高，或者内存顶不住了。咱们来一步步解决：

1. 优化分片策略，别让单个分片太大 #

你现在是按CPU核心数均分DataFrame，但如果DataFrame本身特别大，每个分片还是会占超多内存，进程间传递的时候就容易出问题。可以把分片数量调得比核心数多一些，比如核心数的2-4倍，让每个分片更小更轻：

# 把原来的num_cores改成num_cores*2，根据实际情况调整
num_partitions = num_cores * 2
partitions = np.linspace(0, len(df), num_partitions + 1, dtype=np.int64)
df_split = [df.iloc[partitions[i]:partitions[i + 1]] for i in range(num_partitions)]

2. 用imap替代map，减少内存压力 #

Pool.map会一次性把所有分片都塞给进程池，内存瞬间就被占满了。换成imap就不一样了，它是迭代着给进程传递分片，用多少传多少，内存压力会小很多：

# 把pool.map改成pool.imap
series = pd.concat(pool.imap(partial(apply_wrapper, func=func, **kargs), df_split))

要是你不需要保持原DataFrame的行顺序，用imap_unordered会更快——它一拿到进程的结果就返回，不用等所有进程都跑完。

3. 检查你的apply_wrapper函数，别做多余的事 #

确保你的apply_wrapper只处理必要的数据，比如如果你的目标是对每行执行func，那wrapper应该直接让分片调用apply，别搞复杂了。举个正确的wrapper例子：

def apply_wrapper(df_slice, func, **kwargs):
    # axis=1表示对每行操作，根据你的需求调整
    return df_slice.apply(func, axis=1, **kwargs)

另外，尽量别用复杂的lambda当func，换成普通的函数——lambda有时候序列化会出问题，而且调试起来也麻烦。

4. 实在不行，换个更适合的库 #

如果multiprocessing还是搞不定，试试专门针对pandas优化的并行库：

• swifter：它会自动判断是用普通apply还是并行apply，用法超简单：

  import swifter
  df['result'] = df.swifter.apply(func, **kargs)
  

• dask.dataframe：专门处理超大数据集，能把数据拆成小块在磁盘和内存间切换，不会爆内存：

  import dask.dataframe as dd
  ddf = dd.from_pandas(df, npartitions=num_cores*2)
  # meta参数要指定结果的类型，比如('result', 'float64')
  result = ddf.apply(func, axis=1, meta=('result', 'float64')).compute()
  

修正后的完整代码示例 #

from multiprocessing import cpu_count, Pool
from functools import partial
import numpy as np
import pandas as pd
from pandas import DataFrame

def apply_wrapper(df_slice, func, **kwargs):
    # 按需求调整axis，1是行，0是列
    return df_slice.apply(func, axis=1, **kwargs)

def parallel_applymap_df(df: DataFrame, func, num_cores=cpu_count(), **kargs):
    # 优化分片数量，避免单个分片过大
    num_partitions = num_cores * 2
    partitions = np.linspace(0, len(df), num_partitions + 1, dtype=np.int64)
    df_split = [df.iloc[partitions[i]:partitions[i + 1]] for i in range(num_partitions)]
    
    # 用with语句管理进程池，自动关闭释放资源
    with Pool(num_cores) as pool:
        # 用imap减少内存压力
        results = pool.imap(partial(apply_wrapper, func=func, **kargs), df_split)
        series = pd.concat(results)
    
    return series

最后提几个注意点 #

• 一定要用with Pool(...)来创建进程池，这样进程会自动关闭，不会留着占资源。
• 如果你的func需要用外部变量，尽量通过partial传进去，别用全局变量——大对象全局变量会被重复序列化到每个进程，浪费内存。
• 如果DataFrame里有复杂类型（比如列表、自定义对象），先转换成基本类型再处理，不然序列化会慢到离谱。

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