如何为大型DataFrame并行执行按玩家分组的替换函数？

我需要实现一个函数，在DataFrame中搜索并将旧条目替换为新条目。该DataFrame基于玩家及其门票数据，因此必须按PLAYER_ID分组拆分，而非均分数据块。该函数在小型DataFrame上运行正常，但处理大型DataFrame时，用multiprocessing并行化后计算时间无明显提升，现有代码如下：

# Callback function for updating progress
def update_progress(result):
    with progress_lock:
        progress.value += 1
        print(f"Progress: {progress.value}/{num_chunks} chunks processed.")

def split_dataframe_by_player(df, num_chunks):
    grouped = df.groupby('PLAYER_ID')
    sorted_groups = sorted(grouped, key=lambda x: len(x[1]), reverse=True)
    chunks = [[] for _ in range(num_chunks)]

    for i, group in enumerate(sorted_groups):
        chunks[i % num_chunks].append(group[1])

    return [pd.concat(chunk) if chunk else pd.DataFrame() for chunk in chunks]

# Parallel processing function
def parallel_process(df, func):
    global df_chunks, progress, progress_lock
    num_cores = mp.cpu_count()
    df_chunks = split_dataframe_by_player(df, num_cores)

    progress = Value('i', 0)
    progress_lock = Lock()

    pool = mp.Pool(num_cores)
    results = [pool.apply_async(func, args=(chunk,), callback=update_progress) for chunk in df_chunks]

    pool.close()
    pool.join()

    return pd.concat([r.get() for r in results])

现有代码的问题分析 #

1. 数据拆分开销：split_dataframe_by_player中先分组、排序，再逐个将子DataFrame添加到chunk后执行pd.concat，这个过程会产生额外的内存开销和计算成本，尤其在处理大型DataFrame时。
2. 进程间数据传递成本：每个chunk作为DataFrame在进程间传递时，需要通过pickle序列化/反序列化，大型DataFrame的序列化开销极高，会抵消并行带来的收益。
3. 进度更新的IO开销：update_progress中频繁的print操作属于同步IO，会拖慢整体处理速度，且全局变量的使用增加了代码复杂度。
4. 负载均衡隐患：尽管按组大小排序后分配，但仍可能出现部分chunk数据量远大于其他chunk的情况，导致部分进程长期闲置。

优化方案 #

1. 优先优化替换逻辑的矢量化 #

绝大多数替换操作都可以通过pandas内置的矢量化API实现，效率远高于逐行/逐组的循环处理，甚至无需并行。例如：

• 单值替换：df['target_col'] = df['target_col'].replace(old_val, new_val)
• 多值映射：df['target_col'] = df['target_col'].map(replace_dict).fillna(df['target_col'])
• 条件替换：df.loc[df['target_col'] == old_val, 'target_col'] = new_val

如果你的替换逻辑可以转化为矢量化操作，这是最有效的提速方式，比并行更高效。

2. 优化多进程实现 #

（1）减少数据拆分与传递开销

修改数据拆分逻辑，避免提前concat chunk，直接传递分组键让子进程自行提取对应数据：

def split_player_keys(df, num_chunks):
    grouped = df.groupby('PLAYER_ID')
    # 按组大小排序分组键
    sorted_keys = sorted(grouped.groups.keys(), key=lambda k: len(grouped.get_group(k)), reverse=True)
    chunks = [[] for _ in range(num_chunks)]
    for i, key in enumerate(sorted_keys):
        chunks[i % num_chunks].append(key)
    return chunks, grouped

def process_chunk(keys, grouped):
    # 子进程中提取对应分组并处理
    chunk_df = pd.concat([grouped.get_group(k) for k in keys])
    # 执行你的替换逻辑
    # chunk_df = your_replace_func(chunk_df)
    return chunk_df

def parallel_process_optimized(df, func):
    num_cores = mp.cpu_count()
    key_chunks, grouped = split_player_keys(df, num_cores)
    
    # 使用Manager共享分组数据
    from multiprocessing import Manager
    manager = Manager()
    shared_groups = manager.dict({k: grouped.get_group(k) for k in grouped.groups.keys()})
    
    pool = mp.Pool(num_cores)
    # 用imap_unordered提高任务调度效率
    results = pool.imap_unordered(lambda keys: func(keys, shared_groups), key_chunks)
    
    pool.close()
    pool.join()
    
    return pd.concat(results)

（2）替换低效的进度更新

去掉频繁的print，改用tqdm实现高效进度条，同时避免全局变量：

from tqdm import tqdm

def parallel_process_with_progress(df, func):
    num_cores = mp.cpu_count()
    key_chunks, grouped = split_player_keys(df, num_cores)
    
    manager = Manager()
    shared_groups = manager.dict({k: grouped.get_group(k) for k in grouped.groups.keys()})
    
    pool = mp.Pool(num_cores)
    results = []
    with tqdm(total=len(key_chunks)) as pbar:
        for res in pool.imap_unordered(lambda keys: func(keys, shared_groups), key_chunks):
            results.append(res)
            pbar.update(1)
    
    pool.close()
    pool.join()
    
    return pd.concat(results)

3. 用Dask替代手动多进程 #

Dask是专门处理大型数据集的并行计算框架，可自动管理数据分区和进程，且能保证同一PLAYER_ID在同一个分区中，无需手动拆分：

import dask.dataframe as dd

def process_chunk(df):
    # 你的替换逻辑
    # df['target_col'] = df['target_col'].replace(old_vals, new_vals)
    return df

# 将pandas DataFrame转为Dask DataFrame，按CPU核心数设置分区
ddf = dd.from_pandas(df, npartitions=mp.cpu_count())
# 设置PLAYER_ID为索引，确保同一玩家在同一分区
ddf = ddf.set_index('PLAYER_ID', npartitions=mp.cpu_count())
# 并行处理每个分区
result_df = ddf.map_partitions(process_chunk).compute()
# 恢复原索引
result_df = result_df.reset_index()

Dask会自动处理数据拆分、进程管理和结果合并，代码更简洁，且能有效避免手动多进程的常见陷阱。

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