如何更快地将一个大的CSV文件均匀地按照组拆分成多个较小的CSV文件？

Question

我相信肯定有更好的方法，但是我现在脑子一片空白。我有一个以这种格式排列的CSV文件。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"。

这些组的大小不同，所以硬编码的数字分割方法在这里行不通。

有没有比我目前使用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(file_+'.csv',index=False, header=False, mode='a')

        if file_ == num_files:

            file_ = 0

这是一个基于Python的解决方案，但如果能完成任务，我也可以接受使用awk或bash的方法。

编辑：

为了澄清，我希望将组分割成固定数量的文件，我可以设置。

在这种情况下，是3个文件（所以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

- GreenGodot

---

你的ID列是否总是按照排序值排列？如示例所示。 - RavinderSingh13

是的，ID列已经排序了。忘记澄清这一点。 - GreenGodot

你的真实数据集有多少行？对所有提供的答案进行快速测试后，我发现Python要快得多。你能在你的数据集上进行测试吗？ - mozway

@mozway我更新了我的答案中的第一段脚本，以避免持续打开/关闭文件，以防止出现“太多打开的文件”错误（我怀疑Python脚本也更快，因为它没有这样做），因为我不认为对于这个应用程序是必要的，请再次使用此速度测试，应该会比在每次更改最后一个字段时都打开/关闭每个文件时快得多。 - Ed Morton

Answer 1

在任何Unix系统上的任何shell中使用awk：

$ cat tst.awk
NR==1 {
    hdr = $0
    next
}
$NF != prev {
    out = (((blockCnt++) % X) + 1) ".csv"
    if ( blockCnt <= X ) {
        print hdr > out
    }
    prev = $NF
}
{ print > out }

$ awk -v X=3 -f tst.awk input.csv

$ head [0-9]*.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

If X was some large enough number that you exceed your system limit for concurrently open files and you start getting a "too many open files" error then you'd need to use GNU awk as it handles that internally or change the code to only have 1 file open at a time:

NR==1 {
    hdr = $0
    next
}
$NF != prev {
    close(out)
    out = (((blockCnt++) % X) + 1) ".csv"
    if ( blockCnt <= X ) {
        print hdr > out
    }
    prev = $NF
}
{ print >> out }

or implement your own way of managing how many files are open concurrently.

---

EDIT: here's what the suggestion by @PaulHodges in the comments would result in a script like:

NR == 1 {
    for ( i=1; i <= X; i++ ) {
        print > (i ".csv")
    }
    next
}
$NF != prev {
    out = (((NR-1) % X) + 1) ".csv"
    prev = $NF
}
{ print > out }

- Ed Morton

---

将blockCnt测试移动到一个BEGIN块来写入已知输出文件集的标题，对于大型输入文件而言，是否会对速度产生显著影响？ - Paul Hodges

@PaulHodges 是的，这样会稍微快一点，但是你需要添加一个调用getline函数来读取头部，并将其输入重定向到文件，因为awk还没有打开输入文件，然后文件将在主循环中重新打开，而头部行仍然存在，所以你需要添加NR==1 { next}或类似的内容。考虑到所有这些，最好在一个NR==1部分完成。我刚刚在我的答案末尾添加了一个类似的脚本-唯一的潜在缺点是即使输入中的ID少于X个，它也会始终创建X个输出文件。 - Ed Morton

Answer 2

你可以使用这个awk解决方案：

awk -v X=3 '
FNR == 1 {   # save 1st record as header 
   hdr = $0
   next
}
p != $NF {   # ID field changes, move to new output csv file 
   close(fn)
   fn = ((n++ % X) + 1)".csv" # construct new file name
}
!seen[fn]++ {                 # do we need to print header
   print hdr > fn 
}
{
   print >> fn                # append each record to output
   p = $NF                    # save last field in variable p
}' file

- anubhava

Answer 3

请使用您提供的示例尝试以下代码。如示例所示，最后一列已按排序。

awk -v x="3" '
BEGIN{
  count=1
  outFile=count".csv"
}
FNR==1{
  print
  next
}
prev!=$NF && prev{
  close(outFile)
  count++
  outFile=count".csv"
}
{
  print >> (outFile)
  prev=$NF
}
x==count{ count=1 }
' Input_file

- RavinderSingh13

---

@GreenGodot，请问你能告诉我这里的x==3是什么意思吗？谢谢 - RavinderSingh13

我们将把这些团队分成我可以设定的固定数量的文件中。在这种情况下，是三个文件（所以x = 3）。第一组将放入1.csv，第二组将放入2.csv，第三组将放入3.csv，然后第四组将循环回来并插入到1.csv中。如果你能理解的话？ - GreenGodot

@GreenGodot，请你帮我检查一下我更新的代码，然后告诉我进展如何。 - RavinderSingh13

@RavinderSingh13。感谢您的耐心等待，但实际操作中，当我尝试运行您的代码时，出现了“第19行：if附近的语法错误”。 - GreenGodot

@GreenGodot，我已经更新了代码，请尝试修改后的代码，谢谢。 - RavinderSingh13

Answer 4

使用 groupby 和 factorize 对所需分组数（N）取模：

N = 3

for i, g in df.groupby(pd.factorize(df['ID'])[0]%N):
    g.to_csv(f'chunk{i+1}.csv', index=False)

输出文件：

# chunk1.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

# chunk2.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

# chunk3.csv
Text,ID
this is sample text3,CCCC
this is sample text6,FFFF
this is sample text6,FFFF

时间

在1400万行上进行测试：

15.8 s ± 687 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

大约14秒是由于输入/输出引起的。

与其他答案进行比较（使用shell中的time）：

# @mozway as a python script including imports and reading the file
real    0m20,834s

# @RavinderSingh13
real    1m22,952s

# @anubhava
real    1m23,790s

# @Ed Morton (updated code, original solution was 2m58,171s)
real    0m8,599s

作为一个函数：

import pandas as pd

def split_csv(filename, N=3, id_col='ID', out_basename='chunk'):
    df = pd.read_csv(filename)
    for i, g in df.groupby(pd.factorize(df[id_col])[0]%N):
        g.to_csv(f'{out_basename}{i+1}.csv', index=False)

split_csv('my_file.csv', N=3)

- mozway

---

我知道Python 2已经被弃用，但有时候出于组织或合同上的原因，我们被要求使用它，这就是我现在所处的情况，对此我无能为力。唉，算了。 - Ed Morton

@EdMorton 当然，我明白，但是我已经不再使用Python2编码了 ;) 我测试了你更新的脚本，现在运行时间大约为10秒 - mozway

是的，计划已经确定，我只是重新检查旧脚本，虽然我相当确定在时间安排上没有犯错误。但这是有道理的，如果你一遍又一遍地重新打开文件，行为是二次方的，因为每次迭代都需要重新到达结尾。编辑：再次测试，确实超过2分钟。 - mozway

@EdMorton 你可以在另一台电脑上测试一下。我使用了pandas和df2 = pd.concat([df]*1_000_000, ignore_index=True) ; df2['ID'] = df2['ID'].ne(df2['ID'].shift()).cumsum().astype(str).radd('XXX') ; df2.to_csv('split_test.csv', index=False)来生成具有不同ID的CSV文件（从原始df中）。然后使用time awk -v X=3 -f tst.awk split_test.csv来测试你的解决方案。 - mozway

@EdMorton，我的意思是你的文件是AAA...BBB...CCC...DDD...EEE...FFF... 你应该在AA...B...C...D...E...F...G...H...[...]Z123...上进行测试。只有6个不同的ID，你很少循环遍历文件。尝试使用许多排序后的ID，而不仅仅是6个。 - mozway

Answer 5

这里

group.to_csv(file_+'.csv',index=False, header=False, mode='a')

你提供的第一个参数是字符串，但是to_csv方法允许你提供类似文件的对象作为第一个参数，在这种情况下，你可以避免多次进行文件打开相关的操作。考虑以下简单的比较。

import os
import time
import pandas as pd
REPEAT = 1000
df = pd.DataFrame({'col1':range(100)})
t1 = time.time()
for _ in range(REPEAT):
    df.to_csv('file.csv',index=False,header=False,mode='a')
t2 = time.time()
os.remove('file.csv')
t3 = time.time()
with open('file.csv','a') as f:
    for _ in range(REPEAT):
        df.to_csv(f,index=False,header=False)
t4 = time.time()
print('Using filename',t2-t1)
print('Using filehandle',t4-t3)

输出结果

Using filename 0.35850977897644043
Using filehandle 0.2669696807861328

注意到第二种方式花费了大约第一种方式的75%的时间，所以虽然它更快，但仍然是同一数量级。

- Daweo