创建批处理文件按指定数量唯一ID拆分大型CSV文件

这里有个实用的Python实现方案，完全满足你的需求——按指定数量的唯一ID拆分大型CSV，保证每个ID的所有记录都归属到同一个文件，同时自动保留表头：

实现思路 #

• 先读取原CSV的表头信息，确保每个拆分后的小文件都能继承表头
• 遍历原文件一次，收集所有唯一ID并按指定数量分组（比如每2个ID一组）
• 采用迭代方式处理原文件（避免内存过载），将每条记录写入对应ID组的文件，每个文件仅在首次写入时添加表头

完整代码实现 #

import csv
from collections import defaultdict

def split_csv_by_unique_ids(input_file, chunk_size, output_prefix):
    # 第一步：收集所有唯一ID并分组
    unique_ids = set()
    with open(input_file, 'r', newline='', encoding='utf-8') as infile:
        reader = csv.reader(infile)
        next(reader)  # 跳过表头
        for row in reader:
            if row:  # 跳过空行
                unique_ids.add(row[0])  # 假设ID在第一列，可根据实际修改索引

    # 将唯一ID按chunk_size分组
    id_groups = []
    current_group = []
    for idx, id_val in enumerate(sorted(unique_ids)):  # 排序保证分组顺序稳定
        current_group.append(id_val)
        if (idx + 1) % chunk_size == 0:
            id_groups.append(current_group)
            current_group = []
    if current_group:  # 处理剩余不足chunk_size的ID
        id_groups.append(current_group)

    # 第二步：建立ID到文件组的映射
    id_to_group = {}
    for group_idx, group in enumerate(id_groups):
        for id_val in group:
            id_to_group[id_val] = group_idx + 1  # 文件编号从1开始

    # 第三步：遍历原文件，写入对应分组的文件
    file_handles = {}
    headers = None
    with open(input_file, 'r', newline='', encoding='utf-8') as infile:
        reader = csv.reader(infile)
        headers = next(reader)  # 获取表头

        for row in reader:
            if not row:
                continue
            id_val = row[0]
            group_num = id_to_group[id_val]
            output_file = f"{output_prefix}_{group_num}.csv"

            # 如果文件还没打开，先打开并写入表头
            if output_file not in file_handles:
                fh = open(output_file, 'w', newline='', encoding='utf-8')
                writer = csv.writer(fh)
                writer.writerow(headers)
                file_handles[output_file] = (fh, writer)
            # 写入当前行
            file_handles[output_file][1].writerow(row)

    # 关闭所有打开的文件句柄
    for fh, _ in file_handles.values():
        fh.close()

# 示例调用（匹配你的需求）
split_csv_by_unique_ids(
    input_file='data.csv',
    chunk_size=2,  # 每个文件包含2个唯一ID
    output_prefix='data'
)

代码说明 #

• ID列索引：如果你的ID不在第一列，修改代码中row[0]的索引值即可（比如ID在第二列就改成row[1]）
• 内存友好：采用迭代读取方式，不会把整个大文件加载到内存，适合处理GB级的大型CSV
• 排序稳定：对唯一ID进行排序后分组，保证拆分后的文件内容顺序可预期
• 编码处理：指定encoding='utf-8'避免中文或特殊字符乱码，可根据实际文件编码调整

效果验证（对应你的示例） #

运行代码后：

• data_1.csv会包含ID=1和ID=2的所有5条记录，带表头
• data_2.csv会包含ID=3和ID=4的所有4条记录，带表头

注意事项 #

• 如果原CSV文件有重复的ID行，代码会自动将所有同ID行归入同一文件
• 确保输出目录有写入权限，避免文件创建失败
• 若需要按ID的出现顺序分组（而非排序后的顺序），可以把sorted(unique_ids)改成list(unique_ids)（注意：这种情况下唯一ID的收集顺序取决于文件读取顺序）

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