使用Python写入大型CSV文件的最快方法

Question

我想要在一个csv文件中写入一些随机的样本数据，直到它变得1GB大。以下代码可以实现：

import numpy as np
import uuid
import csv
import os
outfile = 'data.csv'
outsize = 1024 # MB
with open(outfile, 'ab') as csvfile:
    wtr = csv.writer(csvfile)
    while (os.path.getsize(outfile)//1024**2) < outsize:
        wtr.writerow(['%s,%.6f,%.6f,%i' % (uuid.uuid4(), np.random.random()*50, np.random.random()*50, np.random.randint(1000))])    

如何更快地获取它？

- Balzer82

---

为什么你会给这个问题打上numpy的标签，但是并没有使用它（生成随机数并不需要它）？为什么要创建一个csv-writer却每行只写入一个字符串？并未说明文件大小在文件未关闭时是否更新。自己计算文件大小，不要使用getsize函数，这样更快。 - Daniel

Answer 1

问题似乎主要受IO限制。您可以通过一次写入更大的数据块而不是逐行写入来稍微改善I/O：

import numpy as np
import uuid
import os
outfile = 'data-alt.csv'
outsize = 10 # MB
chunksize = 1000
with open(outfile, 'ab') as csvfile:
    while (os.path.getsize(outfile)//1024**2) < outsize:
        data = [[uuid.uuid4() for i in range(chunksize)],
                np.random.random(chunksize)*50,
                np.random.random(chunksize)*50,
                np.random.randint(1000, size=(chunksize,))]
        csvfile.writelines(['%s,%.6f,%.6f,%i\n' % row for row in zip(*data)])   

您可以尝试更改块大小（每个块写入的行数）以查看在您的计算机上效果最佳的设置。

---

这是一个基准测试，将上述代码与您的原始代码进行比较，其中outsize设置为10 MB：

% time original.py

real    0m5.379s
user    0m4.839s
sys 0m0.538s

% time write_in_chunks.py

real    0m4.205s
user    0m3.850s
sys 0m0.351s

所以这个版本比原来的代码快了大约25%。

---

顺便说一句，我尝试用总行数的估计值替换了对os.path.getsize的调用。 不幸的是，它并没有提高速度。由于表示最终整数所需的字节数不同，因此估计也不精确--也就是说，它不能完全复制您原始代码的行为。 所以我把os.path.getsize留在原地。

- unutbu

Answer 2

清除所有不必要的东西，因此它应该更快、更容易理解：

import random
import uuid
outfile = 'data.csv'
outsize = 1024 * 1024 * 1024 # 1GB
with open(outfile, 'ab') as csvfile:
    size = 0
    while size < outsize:
        txt = '%s,%.6f,%.6f,%i\n' % (uuid.uuid4(), random.random()*50, random.random()*50, random.randrange(1000))
        size += len(txt)
        csvfile.write(txt)

- Daniel

---

len(txt) == filesize 吗？而 random.randint(1000) 需要2个参数。 - Balzer82

randint -> randrange。而 len(txt) 表示一行的长度。 - Daniel

好的。但是一行的长度或多行长度之和并不等于文件大小。顺便说一下，你的代码并不更快。试试看吧。 - Balzer82

@Balzer82，最快的写作方式可能是购买SSD :). 在IO瓶颈处优化代码相当困难。有很多低级缓冲和优化发生，我们看不到。不要太惊讶，一个应该运行得更快的代码实际上并没有显著提高速度。 - cel

Answer 3

这是在unutbu的回答基础上进行的更新：

大部分时间都花在生成随机数和检查文件大小上。

如果您提前生成行，则可以评估原始磁盘io性能：

import time
from pathlib import Path
import numpy as np
import uuid
outfile = Path('data-alt.csv')
chunksize = 1_800_000

data = [
    [uuid.uuid4() for i in range(chunksize)],
    np.random.random(chunksize) * 50,
    np.random.random(chunksize) * 50,
    np.random.randint(1000, size=(chunksize,))
]
rows = ['%s,%.6f,%.6f,%i\n' % row for row in zip(*data)]

t0 = time.time()
with open(outfile, 'a') as csvfile:
    csvfile.writelines(rows)
tdelta = time.time() - t0
print(tdelta)

在我的标准860 EVO SSD（非NVMe）上，我得到了1.43秒的时间来处理1,800,000行数据，因此每秒可以处理1,258,741行数据（在我看来还不错）。

- AustEcon