用 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))])    

如何更快地得到它？

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)])   

您可以嘗試使用 chunksize（每個塊寫入的行數）來查看什么在您的機器上效果最好。

---

這是一個基准，將上述代碼與您的原始代碼進行比較，並將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的調用。 不幸的是，它並沒有提高速度。 由於表示最終 int 所需的字節數各不相同，因此估計也不准確——也就是說，它不能完美地復制原始代碼的行為。 所以我把os.path.getsize留在了原處。

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)

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_800_000 行我得到 1.43 秒，所以這是 1,258,741 行/秒（不太破舊的 imo）