如何快速处理大型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, ...

我想打乱行并将它们分成两个文件（用于训练和测试）。问题是我想根据调用站点而不是行进行拆分。可能有不止一行属于一个调用点。所以我先把所有的callsites都读一遍，洗牌拆分如下：

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)。而且逐行读写也可能效率不高。但是我担心加载内存中的所有数据会导致内存错误。有没有更好的方法来处理这样的大文件？

用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

这会将 if row[0] in test_callsites: 的每次测试的工作量从 O(n_test) 减少到 O(1)，如果 n_test 的顺序是四位数或更多（可能，当我们谈论数百万行时）。

您还可以通过更改首先创建它来稍微减少工作量（主要是通过选择更小的东西来改善内存局部性）：

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 方言的换行首选项得到尊重。

如何快速处理大型csv文件？

How to deal with large csv file quickly?

python

csv

large-data

dataframe