如何快速處理大 csv 文件？

Question

我有一個超過 100 萬行的大型 csv 文件。 每行有兩個特征，調用站點（API 調用的位置）和調用站點的一系列標記。 它們寫成：

callsite 1, token 1, token 2, token 3, ...
callsite 1, token 3, token 4, token 4, token 6, ...
callsite 2, token 3, token 1, token 6, token 7, ... 

我想打亂行並將它們分成兩個文件（用於訓練和測試）。 問題是我想根據調用點而不是行進行拆分。 可能有不止一行屬於一個調用點。 所以我首先閱讀了所有的調用站點，將它們隨機排列並拆分如下：

import csv
import random
with open(file,'r') as csv_file:
    reader = csv.reader(csv_file)
    callsites = [row[0] for row in reader]
random.shuffle(callsites)
test_callsites = callsites[0:n_test] //n_test is the number of test cases

然后，我從 csv 文件中讀取每一行，並比較調用站點以將其放入 train.csv 或 test.csv 中，如下所示：

with open(file,'r') as csv_file, open('train.csv','w') as train_file, open('test.csv','w') as test_file:
    reader = csv.reader(csv_file)
    train_writer = csv.writer(train_file)
    test_writer = csv.writer(test_file)
    for row in reader:
        if row[0] in test_callsites:
            test_writer.writerow(row)
        else:
            train_writer.writerow(row)

問題是代碼運行速度太慢，一天多才能完成。 每行的比較導致復雜度 O(n^2)。 而且逐行讀寫也可能效率不高。 但是我擔心加載 memory 中的所有數據會導致 memory 錯誤。 有沒有更好的方法來處理這樣的大文件？

如果我用dataframe讀寫會不會更快？ 但是每行的序列長度是不同的。 我嘗試將數據寫為（將所有標記作為列表放在一列中）：

callsite,     sequence
callsite 1, [token1||token2||token 3]

但是，將 [token 1||token 2||token 3] 恢復為序列似乎並不方便。 有沒有更好的做法來存儲和恢復這種可變長度的數據？

Answer 1

這樣的事情呢？

import csv
import random
random.seed(42) # need this to get reproducible splits

with open("input.csv", "r") as input_file, open("train.csv", "w") as train_file, open(
    "test.csv", "w"
) as test_file:
    reader = csv.reader(input_file)
    train_writer = csv.writer(train_file)
    test_writer = csv.writer(test_file)

    test_callsites = set()
    train_callsites = set()

    for row in reader:
        callsite = row[0]

        if callsite in test_callsites:
            test_writer.writerow(row)
        elif callsite in train_callsites:
            train_writer.writerow(row)
        elif random.random() <= 0.2: # put here the train/test split you need
            test_writer.writerow(row)
            test_callsites.add(callsite)
        else:
            train_writer.writerow(row)
            train_callsites.add(callsite)

這樣，您將需要對文件進行一次傳遞。 缺點是您將獲得大約20% 的拆分。

在 1Mx100 行 (~850mb) 上進行了測試，看起來相當可用。

Answer 2

最簡單的解決方法是更改：

test_callsites = callsites[0:n_test]

至

test_callsites = frozenset(callsites[:n_test])  # set also works; frozenset just reduces chance of mistakenly modifying it

這將減少 test_callsites: 中if row[0] in test_callsites:從O(n_test)到O(1) ，如果n_test大約為四位數或更多（可能，當我們正在談論數百萬行）。

您還可以通過更改首先創建它來稍微減少工作（主要是通過選擇較小的事物箱來改善 memory 位置）：

random.shuffle(callsites)
test_callsites = callsites[0:n_test]

至：

test_callsites = frozenset(random.sample(callsites, n_test))

這避免了重新洗牌整個callsites ，有利於從中選擇n_test值（然后您將其轉換為frozenset ，或只是set ，以進行廉價查找）。 獎金，這是一個單行。 :-)

---

旁注：您的代碼可能是錯誤的。 您必須將newline=''傳遞給您的各種調用才能open ，以確保遵守所選 CSV 方言的換行首選項。