如何并发统计超大压缩文件词频及列直方图？求Python兼容示例

问题分析 #

你现在面临的核心问题是超大gzip压缩TSV文件的列直方图统计，单线程下csv.DictReader的耗时占比最高，想要用生产者-消费者模式拆分任务，让多个worker进程并行处理文件行，同时需要兼容Python 3.5.3和PyPy，还要避免进程创建的开销，让worker长期存活。

可行方案 #

Python的multiprocessing模块本身就提供了原生的队列（Queue）支持，完全可以实现生产者-消费者模式；而ProcessPoolExecutor本身是基于进程池的，它的map方法确实不太适合流式处理超大文件，但我们可以结合队列手动实现任务分发，或者直接用multiprocessing.Process+Queue来构建更灵活的生产者-消费者架构——后者在Python 3.5.3下兼容性更好，也更容易控制worker的生命周期。

实现示例 #

下面是适配你需求的多进程版本代码，核心思路是：

1. 生产者进程负责逐行读取gzip文件，把原始行字符串放到任务队列（避免提前解析占用资源）
2. 多个消费者进程从队列取行，解析成csv字段，统计各自的局部Counter和统计值
3. 所有行处理完成后，生产者发送结束信号，消费者退出，最后合并所有消费者的统计结果

#!/usr/bin/env pypy3
import sys
import math
import gzip
import csv
from collections import Counter
from multiprocessing import Process, Queue, Lock

SamplesList = ('D_Crick', 'D_Watson', 'Normal_Crick', 'Normal_Watson', 'D_WGS', 'Normal_WGS')
# 结束信号，用于通知消费者停止工作
STOP_SIGNAL = None

def worker(queue, result_queue, lock):
    """消费者进程：处理行数据，统计局部结果"""
    local_cnt = {key: Counter() for key in SamplesList}
    local_stat = {key: [0, 0, 0, 0, 0] for key in SamplesList}  # sum_x, sum_x2, _, _, _
    local_max_depth = 0

    while True:
        line = queue.get()
        if line is STOP_SIGNAL:
            # 把局部统计结果发送到结果队列
            result_queue.put((local_cnt, local_stat, local_max_depth))
            break
        
        # 用csv.reader解析单行，比DictReader更高效
        reader = csv.reader([line], delimiter='\t')
        row = next(reader)
        # 对应字段：ChrID(0), Pos(1) + SamplesList(2-7)
        for idx, k in enumerate(SamplesList):
            val = int(row[2 + idx])
            if val > local_max_depth:
                local_max_depth = val
            local_cnt[k][val] += 1
            local_stat[k][0] += val
            local_stat[k][1] += val * val

def producer(queue, infile_path, lock):
    """生产者进程：读取gzip文件，逐行放入任务队列"""
    with gzip.open(infile_path, 'rt') as f:
        for line in f:
            queue.put(line.strip())
    # 给每个worker发送停止信号，确保所有进程能正常退出
    for _ in range(NUM_WORKERS):
        queue.put(STOP_SIGNAL)

def merge_results(results):
    """合并所有消费者的局部统计结果"""
    merged_cnt = {key: Counter() for key in SamplesList}
    merged_stat = {key: [0, 0, 0, 0, 0] for key in SamplesList}
    global_max_depth = 0
    total_records = 0

    for cnt, stat, max_depth in results:
        # 统计总处理行数：取任意一个Counter的总计数即可
        total_records += sum(next(iter(cnt.values())).values())
        if max_depth > global_max_depth:
            global_max_depth = max_depth
        # 合并Counter和统计值
        for k in SamplesList:
            merged_cnt[k].update(cnt[k])
            merged_stat[k][0] += stat[k][0]
            merged_stat[k][1] += stat[k][1]
    
    return total_records, global_max_depth, merged_cnt, merged_stat

def main():
    global NUM_WORKERS
    # 根据CPU核心数设置worker数量，PyPy下建议和核心数一致或略多
    NUM_WORKERS = 4 if len(sys.argv) < 4 else int(sys.argv[4])

    if len(sys.argv) < 3 :
        print('Usage:',sys.argv[0],'<samtools.depth.gz> <out.tsv> [verbose=0] [workers=4]',file=sys.stderr,flush=True)
        exit(0)
    try:
        verbose = int(sys.argv[3])
    except:
        verbose = 0
    inDepthFile = sys.argv[1]
    outFile = sys.argv[2]
    print('From:[{}], To:[{}].\nVerbose: [{}], Workers: [{}].'.format(inDepthFile,outFile,verbose,NUM_WORKERS),file=sys.stderr,flush=True)

    # 初始化队列：任务队列限制大小，避免内存溢出；结果队列用于收集统计结果
    task_queue = Queue(maxsize=NUM_WORKERS*2)
    result_queue = Queue()
    lock = Lock()  # 队列本身是进程安全的，这里预留锁用于后续可能的共享资源操作

    # 启动消费者进程
    workers = []
    for _ in range(NUM_WORKERS):
        p = Process(target=worker, args=(task_queue, result_queue, lock))
        p.start()
        workers.append(p)
    
    # 启动生产者进程，等待它完成所有行的读取
    prod_process = Process(target=producer, args=(task_queue, inDepthFile, lock))
    prod_process.start()
    prod_process.join()
    
    # 收集所有消费者的结果
    results = []
    for _ in range(NUM_WORKERS):
        results.append(result_queue.get())
    
    # 等待所有worker进程退出
    for p in workers:
        p.join()
    
    # 合并结果并计算最终统计值
    RecordCnt, MaxDepth, cDepthCnt, cDepthStat = merge_results(results)

    for k in SamplesList:
        cDepthStat[k][2] = cDepthStat[k][0] / RecordCnt  # 均值E(X)
        cDepthStat[k][3] = cDepthStat[k][1] / RecordCnt  # E(X²)
        cDepthStat[k][4] = math.sqrt(cDepthStat[k][3] - cDepthStat[k][2]*cDepthStat[k][2])  # 标准差

    # 输出结果到TSV文件
    tsvout = open(outFile, 'wt')
    print('#{}\t{}'.format('Depth','\t'.join(SamplesList)),file=tsvout)
    print( '#N={},SD:\t{}'.format(RecordCnt,'\t'.join(str(round(cDepthStat[col][4],1)) for col in SamplesList)),file=tsvout)
    for depth in range(0,MaxDepth+1):
        print( '{}\t{}'.format(depth,'\t'.join(str(cDepthCnt[col][depth]) for col in SamplesList)),file=tsvout)
    tsvout.close()

    print('[!]Lines Read:[{}], MaxDepth is [{}].'.format(RecordCnt,MaxDepth),file=sys.stderr,flush=True)

if __name__ == "__main__":
    main()

关键优化点说明 #

• 替换csv.DictReader为csv.reader：直接解析单行字符串，避免DictReader的字典映射开销，在多进程场景下效率提升更明显
• 队列大小限制：task_queue设置了maxsize，防止生产者过快写入导致内存占用过高
• PyPy友好：代码使用原生multiprocessing模块，PyPy对其支持完善，能充分发挥多进程的性能优势
• 兼容Python 3.5.3：所有API都是Python 3.5+支持的，没有使用后续版本新增特性

关于ProcessPoolExecutor的补充 #

如果一定要用ProcessPoolExecutor，也可以结合队列实现任务分发，但在Python 3.5.3中，这种模式不如Process+Queue直观，且难以控制worker的长期存活状态。上面的实现已经能很好满足你的需求，是更稳妥的选择。

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