如何并行排序并提取大型嵌套字典中每个查询的Top-K文档

问题背景 #

我们有一个存储排序结果的大型嵌套字典，外层键为查询ID，内层字典是对应查询的文档-评分映射。需要对每个内层字典排序并提取Top-K项，原串行处理逻辑已基于heapq.nlargest实现，由于各查询的排序操作完全独立，可通过并行化提升处理效率。

并行实现方案 #

方案1：使用concurrent.futures.ProcessPoolExecutor #

这是Python标准库中简洁易用的并行实现方式，适配CPU密集型任务场景。

import json
import heapq
from concurrent.futures import ProcessPoolExecutor

top_k = 5

def process_query(query_item):
    """处理单个查询的Top-K提取逻辑"""
    query_idx, docs_scores = query_item
    top_docs = heapq.nlargest(top_k, docs_scores.items(), key=lambda x: x[1])
    return (query_idx, dict(top_docs))

if __name__ == "__main__":
    # 假设ranking是你的大型嵌套字典
    # ranking = ... 加载目标数据
    
    # 启动进程池并行处理所有查询
    with ProcessPoolExecutor() as executor:
        results = list(executor.map(process_query, ranking.items()))
    
    # 重新构建处理后的结果字典
    parallel_ranking = dict(results)
    
    # 输出格式化结果
    print(json.dumps(parallel_ranking, indent=4))

方案2：使用multiprocessing.Pool #

这是传统的多进程池实现方式，对进程调度有更灵活的控制空间。

import json
import heapq
from multiprocessing import Pool

top_k = 5

def process_query(query_item):
    query_idx, docs_scores = query_item
    top_docs = heapq.nlargest(top_k, docs_scores.items(), key=lambda x: x[1])
    return (query_idx, dict(top_docs))

if __name__ == "__main__":
    # 假设ranking是你的大型嵌套字典
    # ranking = ... 加载目标数据
    
    # 创建进程池，默认使用CPU核心数，可通过参数指定进程数量
    with Pool() as pool:
        results = pool.map(process_query, ranking.items())
    
    # 构建最终结果字典
    parallel_ranking = dict(results)
    
    print(json.dumps(parallel_ranking, indent=4))

关键注意事项 #

• 进程安全：多进程模式下无法直接修改原字典，需通过收集各进程返回结果，重新构建最终字典。
• 数据传递开销：若嵌套字典规模极大，可考虑分块处理或使用共享内存（如multiprocessing.Manager）减少数据复制开销；多数场景下直接传递查询键值对的开销可接受。
• 进程数量控制：ProcessPoolExecutor和Pool默认使用CPU核心数，也可通过max_workers参数手动指定进程数，避免过多进程导致调度开销激增。

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