我有一个超过100万行的大型csv文件。每行有两个特性，callsite(API调用的位置(和指向callsite的令牌序列。它们被写成：
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, ...
我想打乱行并将它们拆分为两个文件(用于训练和测试(。问题是，我想根据调用位置而不是行进行拆分。可能有多个行属于一个调用站点。因此，我首先阅读了所有的callsite，将它们打乱并拆分如下：
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(。并且逐行读取和写入也可能不是有效的。但我担心加载内存中的所有数据会导致内存错误。有没有更好的方法来处理这样的大文件？
如果我使用数据帧来读写它会更快吗？但序列长度每行都不同。我试图将数据写为(将所有令牌作为列表放在一列中(：
callsite, sequence
callsite 1, [token1||token2||token 3]
然而，将[令牌1||令牌2||令牌3]恢复为序列似乎并不方便。有没有更好的做法来存储和恢复可变长度的数据？
最简单的修复方法是更改：
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))
这避免了为了从中选择n_test值而重新排列整个callsites(然后将其转换为frozenset或仅set，以进行廉价查找(。奖金，这是一句俏皮话。：-(
旁注：您的代码可能是错误的。必须将newline=''传递给对open的各种调用，以确保所选CSV方言的换行首选项得到遵守。
这样的东西怎么样？
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(上进行了测试，似乎可以合理使用。