Python Multiprocessing处理超大Pandas DataFrame时程序卡顿求助

核心问题诊断 #

你的代码在大数据场景下卡顿的根源有三个：

• 内存过载：每个子进程都会复制一份1亿行的DataFrame，63个进程会把内存占用放大63倍，直接导致系统内存耗尽，进程陷入无响应状态。
• 重复计算浪费：每个task函数都重新执行一次全量DataFrame的groupby+agg操作，这是完全冗余的——全量分组只需要执行一次，不需要每个任务都重复计算。
• 无效任务量过大：stores_list生成了1亿个元素（和DataFrame行数一致），但实际唯一的store数量只有1000万（0.1*nrows），大量重复任务占用了进程资源。

优化方案 #

1. 预先完成全局聚合（最优方案） #

既然你的目标是计算每个store的均价，最直接高效的方式是在主进程中一次性完成groupby计算，完全不需要多进程——Pandas的groupby本身已经做了优化，单进程处理1亿行的聚合比多进程重复计算要快得多。

2. 必要时的多进程优化（如果需拆分计算） #

如果业务逻辑必须用多进程（比如更复杂的每个store处理），需要做以下调整：

• 只传递必要数据：预先拆分数据，每个进程只处理部分store对应的子DataFrame，而非全量数据。
• 减少任务数量：基于唯一store列表生成任务，避免重复处理同一store。
• 避免全局变量传递大对象：改用multiprocessing.Manager或者共享内存（如pandas.DataFrame的share_memory方法）来共享数据，避免每个进程复制全量DataFrame。

完整修改代码（最优方案） #

直接在主进程完成聚合，效率远高于原多进程方案：

import time
import pandas as pd
import random as rd
import psutil

def create_datafile(nrows):
    price = [rd.randint(0,300) for i in range(nrows)]
    stores = [i % (0.1 * nrows) for i in range(nrows)]
    data = zip(stores, price)
    return pd.DataFrame(data=data, columns=["stores", "price"])

def get_ram_usage_pct():
    return psutil.virtual_memory().percent

if __name__ == "__main__":
    nrows = 100000000

    print('Creating data...')
    df = create_datafile(nrows)
    print('Data created.')
    print('RAM usage after data creation is {} %'.format(get_ram_usage_pct()))

    tic = time.time()
    # 主进程一次性完成所有store的均值计算
    product_df = df.groupby('stores', as_index=False).agg(mean_price=("price", 'mean'))
    # 转换为字典格式
    dic_mean = product_df.set_index('stores')['mean_price'].to_dict()
    toc = time.time()

    print(f'Processed data in {round((toc-tic)/60,1)} minutes (rounded to 0.1).')
    # print(dic_mean)

多进程版本优化（针对复杂场景） #

如果必须用多进程处理每个store的逻辑，修改如下：

import time
import pandas as pd
import random as rd
import multiprocessing as mp
import psutil

def create_datafile(nrows):
    price = [rd.randint(0,300) for i in range(nrows)]
    stores = [i % (0.1 * nrows) for i in range(nrows)]
    data = zip(stores, price)
    return pd.DataFrame(data=data, columns=["stores", "price"])

def task(store_data):
    store, group_df = store_data
    mean_price = group_df['price'].mean()
    # 模拟存活打印，每1000个任务打印一次
    if store % 1000 == 0:
        print(f'Processing store {store}, mean price: {mean_price}')
    return (store, mean_price)

def get_ram_usage_pct():
    return psutil.virtual_memory().percent

if __name__ == "__main__":
    nrows = 100000000
    nb_jobs = 63

    print('Creating data...')
    df = create_datafile(nrows)
    print('Data created.')
    print('RAM usage after data creation is {} %'.format(get_ram_usage_pct()))

    # 按store分组，生成(store, group_df)的迭代器
    grouped = df.groupby('stores')
    store_groups = [(store, group) for store, group in grouped]

    tic = time.time()
    print(f'Max number of jobs: {mp.cpu_count() - 1}')
    print(f'Running: {min(nb_jobs, mp.cpu_count() - 1)} jobs...')
    
    with mp.Pool(processes=min(nb_jobs, mp.cpu_count() - 1)) as pool:
        results = pool.map(task, store_groups)
        dic_mean = dict(results)
    
    toc = time.time()
    print(f'Processed data in {round((toc-tic)/60,1)} minutes (rounded to 0.1).')
    # print(dic_mean)

关键优化点说明 #

• 原代码中每个task重复执行全量groupby，修改后仅在主进程做一次分组（或直接聚合），避免了99.9%的重复计算。
• 多进程版本中，每个任务仅处理单个store的子DataFrame，内存占用大幅降低，不会出现多进程复制全量DataFrame的情况。
• 任务数量从1亿减少到1000万（唯一store数量），避免了大量重复任务的资源消耗。

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