如何高效按分组将大型CSV文件均分至指定数量的小CSV文件？

我确定存在更优的实现方式，但目前毫无头绪。我有如下格式的CSV文件，其中ID列已排序，同ID的记录已分组：

Text                 ID
this is sample text, AAAA
this is sample text, AAAA
this is sample text, AAAA
this is sample text, AAAA
this is sample text, AAAA
this is sample text2, BBBB
this is sample text2, BBBB
this is sample text2, BBBB
this is sample text3, CCCC
this is sample text4, DDDD
this is sample text4, DDDD
this is sample text5, EEEE
this is sample text5, EEEE
this is sample text6, FFFF
this is sample text6, FFFF

我需要快速将该CSV文件拆分至X个小CSV文件中。例如当X=3时，AAAA组存入1.csv，BBBB组存入2.csv，CCCC组存入3.csv，后续组循环存入1.csv、2.csv等。由于各组大小不一，按固定行数拆分的方式不可行。

目前我使用Python的Pandas groupby实现该功能，代码如下：

file_ = 0
num_files = 3

for name, group in df.groupby(by=['ID'], sort=False):
    file_ += 1
    group['File Num'] = file_
    group.to_csv(f"{file_}.csv", index=False, header=False, mode='a')
    if file_ == num_files:
        file_ = 0

这是基于Python的方案，但我也接受使用awk或bash的实现方式，请问是否有更高效可靠的拆分方法？

补充说明：我需要将分组分配至固定数量的文件中。以X=3为例，第一组（AAAA）存入1.csv，第二组存入2.csv，第三组存入3.csv，第四组循环存入1.csv，以此类推。

示例输出1.csv：

Text                 ID
this is sample text, AAAA
this is sample text, AAAA
this is sample text, AAAA
this is sample text, AAAA
this is sample text, AAAA
this is sample text4, DDDD
this is sample text4, DDDD

示例输出2.csv：

Text                 ID
this is sample text2, BBBB
this is sample text2, BBBB
this is sample text2, BBBB
this is sample text5, EEEE
this is sample text5, EEEE

示例输出3.csv：

Text                 ID
this is sample text3, CCCC
this is sample text6, FFFF
this is sample text6, FFFF

一、Python Pandas 优化版 #

现有代码可以优化，减少不必要的列操作，同时提前处理表头避免逻辑混乱：

import pandas as pd

num_files = 3
df = pd.read_csv("input.csv")

# 提取表头并初始化目标文件
header = ','.join(df.columns)
for i in range(1, num_files + 1):
    with open(f"{i}.csv", "w", newline="") as f:
        f.write(f"{header}\n")

file_idx = 0
prev_id = None

# 逐行遍历处理分组
for _, row in df.iterrows():
    current_id = row['ID']
    if current_id != prev_id:
        file_idx = (file_idx + 1) % num_files
        file_idx = num_files if file_idx == 0 else file_idx
        prev_id = current_id
    # 追加行到对应文件
    with open(f"{file_idx}.csv", "a", newline="") as f:
        f.write(','.join(map(str, row.values)) + '\n')

优化点：避免groupby的内存开销（大文件场景更友好），提前写入表头，逻辑更清晰。

二、Awk 方案（大文件首选） #

Awk适合处理超大CSV文件，无需加载全量数据到内存，速度快、资源占用低：

BEGIN {
    FS = ","
    num_files = 3
    # 读取表头并写入所有目标文件
    getline header
    for (i=1; i<=num_files; i++) {
        print header > i".csv"
        close(i".csv") # 关闭文件避免句柄泄漏
    }
    file_idx = 0
    prev_id = ""
}
{
    # 清理ID字段的前后空格
    current_id = $2
    gsub(/^[ \t]+|[ \t]+$/, "", current_id)
    # 切换分组时更新目标文件索引
    if (current_id != prev_id) {
        file_idx = (file_idx + 1) % num_files
        file_idx = (file_idx == 0) ? num_files : file_idx
        prev_id = current_id
    }
    print $0 >> file_idx".csv"
    close(file_idx".csv") # 大文件建议每次关闭，避免缓存溢出
}

使用方式：awk -f split.awk input.csv

三、Bash 脚本方案 #

结合系统工具实现，适合bash环境下的批量处理：

#!/bin/bash
num_files=3
input_file="input.csv"

# 提取表头并初始化目标文件
header=$(head -n1 "$input_file")
for i in $(seq 1 $num_files); do
    echo "$header" > "$i.csv"
done

# 用awk处理数据行
tail -n+2 "$input_file" | awk -v num_files="$num_files" '
BEGIN {
    FS = ","
    file_idx = 0
    prev_id = ""
}
{
    current_id = $2
    gsub(/^[ \t]+|[ \t]+$/, "", current_id)
    if (current_id != prev_id) {
        file_idx = (file_idx + 1) % num_files
        file_idx = (file_idx == 0) ? num_files : file_idx
        prev_id = current_id
    }
    print $0 >> file_idx".csv"
    close(file_idx".csv")
}'

运行脚本即可完成拆分。

• Python Pandas：代码易读，适合中等规模数据，熟悉Python的用户上手快；但超大文件场景内存占用较高。
• Awk：处理超大CSV的最优选择，速度快、内存占用极低，适合服务器端批量任务。
• Bash脚本：依赖系统原生工具，适合bash环境，但逻辑复杂度高于Awk。

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