如何快速处理大 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 方言的换行首选项。