Python拆分大CSV写入多文件的性能优化方案咨询

核心思路 #

ExitStack 支持动态添加上下文管理器，不需要提前一次性批量打开所有文件。你可以在遍历源文件的过程中，遇到新的类型时再动态打开对应文件、创建CSV写入器，用字典存储每个类型和写入器的映射关系即可，所有打开的文件句柄会由ExitStack统一管理，退出时自动批量关闭。

完整可运行代码 #

import csv
from contextlib import ExitStack

counter = 0
# 存储映射关系：key为行首的类型值，value为对应文件的CSV写入器
writers = {}

with ExitStack() as stack:
    # 打开源大文件
    large_file = stack.enter_context(open("PN000047.csv", newline=''))
    csv_reader = csv.reader(large_file)

    for row in csv_reader:
        row_type = row[0]
        # 遇到未记录的新类型时，动态创建文件和写入器
        if row_type not in writers:
            new_file = stack.enter_context(open(f"{row_type}.csv", "a", newline=''))
            writers[row_type] = csv.writer(new_file)
        # 写入当前行
        writers[row_type].writerow(row)
        counter += 1

print(counter)

注意事项 #

• 该实现的IO操作次数从原方案的百万级降到和源文件中唯一类型数量一致，性能会有数十到上百倍的提升。
• 普通操作系统默认限制单进程最大打开文件句柄数为512~1024，如果你的源文件中唯一类型数量超过该值，可以先临时调大系统句柄上限（Linux下执行ulimit -n 4096），或者补充LRU淘汰逻辑，长时间未使用的文件先关闭，后续用到时再重新打开。

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