如何在Python中将大csv文件拆分为均匀大小的块？

Question

在一个基本的程序中，我有以下的处理过程。

import csv
reader = csv.reader(open('huge_file.csv', 'rb'))

for line in reader:
    process_line(line)

可以看看这个相关的问题。我想每100行发送一次处理结果，以实现批量分片。

实现这个相关答案的问题是，csv对象不能直接用下标访问，也不能使用len函数来获取长度。

>>> import csv
>>> reader = csv.reader(open('dataimport/tests/financial_sample.csv', 'rb'))
>>> len(reader)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: object of type '_csv.reader' has no len()
>>> reader[10:]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: '_csv.reader' object is unsubscriptable
>>> reader[10]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: '_csv.reader' object is unsubscriptable

我该怎么解决这个问题呢？

Answer 1

对于所有的 .csv 文件来说，没有一个特别好的方法来处理这个问题。你可以使用 file.seek 来跳过文件的一部分，从而把文件分成几个小块。然后，你需要逐个字节地扫描，找到每一行的结束位置。这样你就可以独立处理这两个小块了。下面的代码示例（未经测试）可以帮助你入门。

file_one = open('foo.csv')
file_two = open('foo.csv') 
file_two.seek(0, 2)     # seek to the end of the file
sz = file_two.tell()    # fetch the offset
file_two.seek(sz / 2)   # seek back to the middle
chr = ''
while chr != '\n':
    chr = file_two.read(1)
# file_two is now positioned at the start of a record
segment_one = csv.reader(file_one)
segment_two = csv.reader(file_two)

我不太确定你怎么知道什么时候完成了对 segment_one 的遍历。如果你的 CSV 文件中有一列是行 ID，那么当你遇到 segment_two 中第一行的行 ID 时，你就可以停止处理 segment_one 了。

Answer 2

我们可以使用pandas这个模块来处理这些大的csv文件。

df = pd.DataFrame()
temp = pd.read_csv('BIG_File.csv', iterator=True, chunksize=1000)
df = pd.concat(temp, ignore_index=True)

Answer 3

你只需要把你的 reader 变得可以用下标访问，就可以把它包裹在一个 list 里。显然，这在处理非常大的文件时会出现问题（可以参考下面的 更新 中的其他方法）：

>>> reader = csv.reader(open('big.csv', 'rb'))
>>> lines = list(reader)
>>> print lines[:100]
...

进一步阅读：如何在Python中将列表拆分成均匀大小的块？

---

更新 1（列表版本）：另一种可能的方法是，在遍历行时，处理每个到达的块：

#!/usr/bin/env python

import csv
reader = csv.reader(open('4956984.csv', 'rb'))

chunk, chunksize = [], 100

def process_chunk(chuck):
    print len(chuck)
    # do something useful ...

for i, line in enumerate(reader):
    if (i % chunksize == 0 and i > 0):
        process_chunk(chunk)
        del chunk[:]  # or: chunk = []
    chunk.append(line)

# process the remainder
process_chunk(chunk)

---

更新 2（生成器版本）：我还没有进行性能测试，但也许你可以通过使用块 生成器 来提高性能：

#!/usr/bin/env python

import csv
reader = csv.reader(open('4956984.csv', 'rb'))

def gen_chunks(reader, chunksize=100):
    """ 
    Chunk generator. Take a CSV `reader` and yield
    `chunksize` sized slices. 
    """
    chunk = []
    for i, line in enumerate(reader):
        if (i % chunksize == 0 and i > 0):
            yield chunk
            del chunk[:]  # or: chunk = []
        chunk.append(line)
    yield chunk

for chunk in gen_chunks(reader):
    print chunk # process chunk

# test gen_chunk on some dummy sequence:
for chunk in gen_chunks(range(10), chunksize=3):
    print chunk # process chunk

# => yields
# [0, 1, 2]
# [3, 4, 5]
# [6, 7, 8]
# [9]

有一个小问题，正如 @totalhack 指出的：

要注意，这会不断返回同一个对象，只是内容不同。如果你打算在每次迭代之间对这个块进行所有需要的操作，这样是没问题的。