Pandas多进程Apply内存不足问题及异构DataFrame处理

问题背景 #

现有一个Pandas DataFrame，price_list列数据类型不统一，包含字典格式字符串、单个数值或NaN，需要根据shop_id匹配字典中的对应价格，转换为目标格式。单进程apply逻辑可行但处理大数据量耗时过长，改用multiprocessing.Pool4核运行时出现内存不足、内核崩溃（环境：16GB内存，启动时已占用8GB，Linux Python3.6）。

原始数据示例：

import pandas as pd
import numpy as np
import multiprocessing as mp
from functools import partial

df = pd.DataFrame({'item_id':[1,1,1,2,3,4,4],
'shop_id':['S1','S2','S3','S2','S3','S1','S2'], 
'price_list':["{'10':['S1','S2'], '20':['S3'], '30':['S4']}","{'10':['S1','S2'], '20':['S3'], '30':['S4']}","{'10':['S1','S2'], '20':['S3'], '30':['S4']}",'50','NaN',"{'10':['S1','S2','S3'],'25':['S4']}","{'10':['S1','S2','S3'],'25':['S4']}"]})

目标转换格式：

原始单进程逻辑：

def get_price(row):
    if isinstance(row['price_list'], str) and row['price_list'].startswith('{'):
        prices = eval(row['price_list'])
        price = np.nan
        for key, value in prices.items():
            if row['shop_id'] in value:
                price = key
                break
    else:
        price = row["price_list"]
        # 处理字符串NaN转为实际NaN
        if price == 'NaN':
            price = np.nan
    return price

df['price'] = df.apply(get_price, axis=1)

优化方案（不减少核数） #

1. 优化多进程数据传递，减少内存拷贝 #

原生multiprocessing.Pool在传递大DataFrame分块时会做完整拷贝，内存开销翻倍。改用concurrent.futures.ProcessPoolExecutor配合迭代分块，内存管理更高效：

from concurrent.futures import ProcessPoolExecutor

def process_chunk(chunk):
    chunk['price'] = chunk.apply(get_price, axis=1)
    return chunk

def parallel_process(df, num_processes=4):
    chunks = np.array_split(df, num_processes)
    with ProcessPoolExecutor(max_workers=num_processes) as executor:
        processed_chunks = list(executor.map(process_chunk, chunks))
    return pd.concat(processed_chunks, ignore_index=True)

# 调用
df = parallel_process(df, num_processes=4)

2. 预处理price_list，避免重复解析字典字符串 #

大量行的price_list内容重复，先对其去重解析，再映射回原DataFrame，减少重复计算和内存占用：

# 第一步：创建price_list到解析后数据的映射
unique_prices = df['price_list'].unique()
price_map = {}
for p_str in unique_prices:
    if isinstance(p_str, str) and p_str.startswith('{'):
        price_map[p_str] = eval(p_str)
    else:
        price_map[p_str] = p_str

# 第二步：改写处理函数，直接复用映射结果
def get_price_mapped(row):
    p_data = price_map[row['price_list']]
    if isinstance(p_data, dict):
        price = np.nan
        for key, value in p_data.items():
            if row['shop_id'] in value:
                price = key
                break
    else:
        price = np.nan if p_data == 'NaN' else p_data
    return price

# 结合多进程处理
df = parallel_process(df, num_processes=4)

3. 改用Dask处理大DataFrame #

Dask专为大数据集设计，自动分块并管理内存，避免内存溢出：

import dask.dataframe as dd

# 将Pandas DataFrame转为Dask DataFrame，分区数等于核数
ddf = dd.from_pandas(df, npartitions=4)

# 应用处理函数
ddf['price'] = ddf.apply(get_price_mapped, axis=1, meta=('price', object))

# 转换回Pandas DataFrame
df = ddf.compute()

4. 优化处理函数，提升内存效率 #

用ast.literal_eval替代eval（更安全且内存占用更低），同时优化逻辑减少冗余操作：

import ast

def get_price_optimized(row):
    p_list = row['price_list']
    if isinstance(p_list, str):
        if p_list.startswith('{'):
            try:
                prices = ast.literal_eval(p_list)
                price = np.nan
                for key, shops in prices.items():
                    if row['shop_id'] in shops:
                        price = key
                        break
                return price
            except:
                return np.nan
        elif p_list == 'NaN':
            return np.nan
    return p_list

# 替换原函数后复用多进程逻辑
df = parallel_process(df, num_processes=4)

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