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

Posted 2023-03-11

【中文标题】如何快速处理大型csv文件？

【英文标题】：How to deal with large csv file quickly?

【发布时间】：2021-04-17 18:28:47

【问题描述】：

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

如果我使用dataframe读写会更快吗？但是每行的序列长度是不同的。我尝试将数据写为（将所有标记作为列表放在一列中）：

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

但是，将 [token 1||token 2||token 3] 还原为序列似乎并不方便。 有没有更好的做法来存储和恢复这种可变长度的数据？

【问题讨论】：

特定呼叫站点的线路是否相邻，就像您的示例中一样？或者它的行可以散布在整个文件中吗？

为什么需要读取两次文件？每次看到以前从未见过的呼叫站点时，只需随机决定选择哪个集合。

【参考方案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) 上进行测试，看起来相当可用。

【参考方案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 方言的换行首选项。