Python：如何将大型文本文件读入内存

Question

我在一台配有1GB内存的Mac Mini上使用Python 2.6。我想读取一个很大的文本文件。

$ ls -l links.csv; file links.csv; tail links.csv 
-rw-r--r--  1 user  user  469904280 30 Nov 22:42 links.csv
links.csv: ASCII text, with CRLF line terminators
4757187,59883
4757187,99822
4757187,66546
4757187,638452
4757187,4627959
4757187,312826
4757187,6143
4757187,6141
4757187,3081726
4757187,58197

这个文件的每一行都是由两个用逗号分隔的整数值组成的元组。我想把整个文件读进来，并根据第二列进行排序。我知道其实可以不把整个文件都读进内存就进行排序。但我觉得对于一个500MB的文件，既然我有1GB的内存，应该还是可以在内存中处理的。

但是，当我尝试读取这个文件时，Python似乎分配了比文件在磁盘上占用的空间多得多的内存。所以即使有1GB的内存，我也无法把这个500MB的文件读进内存。

我用来读取文件并打印一些内存使用信息的Python代码是：

#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys

infile=open("links.csv", "r")

edges=[]
count=0
#count the total number of lines in the file
for line in infile:
 count=count+1

total=count
print "Total number of lines: ",total

infile.seek(0)
count=0
for line in infile:
 edge=tuple(map(int,line.strip().split(",")))
 edges.append(edge)
 count=count+1
 # for every million lines print memory consumption
 if count%1000000==0:
  print "Position: ", edge
  print "Read ",float(count)/float(total)*100,"%."
  mem=sys.getsizeof(edges)
  for edge in edges:
   mem=mem+sys.getsizeof(edge)
   for node in edge:
    mem=mem+sys.getsizeof(node) 

  print "Memory (Bytes): ", mem 

我得到的输出是：

Total number of lines:  30609720
Position:  (9745, 2994)
Read  3.26693612356 %.
Memory (Bytes):  64348736
Position:  (38857, 103574)
Read  6.53387224712 %.
Memory (Bytes):  128816320
Position:  (83609, 63498)
Read  9.80080837067 %.
Memory (Bytes):  192553000
Position:  (139692, 1078610)
Read  13.0677444942 %.
Memory (Bytes):  257873392
Position:  (205067, 153705)
Read  16.3346806178 %.
Memory (Bytes):  320107588
Position:  (283371, 253064)
Read  19.6016167413 %.
Memory (Bytes):  385448716
Position:  (354601, 377328)
Read  22.8685528649 %.
Memory (Bytes):  448629828
Position:  (441109, 3024112)
Read  26.1354889885 %.
Memory (Bytes):  512208580

在读取了500MB文件的25%后，Python就消耗了500MB的内存。所以看起来把文件内容存储为整数元组的列表并不是很节省内存。

有没有更好的方法，让我可以把这个500MB的文件读进我1GB的内存里呢？

Answer 1

把输入的每一行存储在内存中最省钱的方法是用数组，也就是用array.array('i')这种形式——假设每个数字都能放进一个有符号的32位整数里。这样存储的内存开销是8N字节，其中N是行数。

下面是如何排序并把输出文件按顺序写出的：

from array import array
import csv
a = array('i')
b = array('i')
for anum, bnum in csv.reader(open('input.csv', 'rb')):
    a.append(int(anum))
    b.append(int(bnum))
wtr = csv.writer(open('output.csv', 'wb'))
for i in sorted(xrange(len(a)), key=lambda x: b[x]):
    wtr.writerow([a[i], b[i]])

不幸的是，sorted()返回的是一个列表，而不是一个迭代器，这个列表会比较大：指针占用4N字节，整数对象占用12N字节，也就是说，sorted()的输出总共需要16N字节。注意：这是基于32位机器上的CPython 2.X版本；在3.X版本和64位机器上情况会更糟。总的来说，这样加起来是24N字节。假设你有3100万行数据，那么你需要的内存就是31 * 24 = 744 MB……看起来应该是可以的；不过要注意，这个计算没有考虑排序过程中可能分配的额外内存，但你还有一个合理的安全余地。

顺便问一下，额外的1GB或3GB内存，按照你的工资水平换算成小时，成本是多少呢？

Answer 2

所有的Python对象在存储数据的基础上，还会占用一些额外的内存。根据我在32位的Ubuntu系统上使用getsizeof的结果，一个元组（tuple）会额外占用32个字节，而一个整数（int）则占用12个字节。所以你文件中的每一行大约会占用56个字节，加上列表中的一个4字节的指针。我猜在64位系统上，这个占用会更多。这个数据和你提供的数字是一致的，这意味着你的3000万行数据大约会占用1.8GB的内存。

我建议你不要用Python，而是使用Unix的排序工具。我不是Mac用户，但我猜OS X的排序选项和Linux的版本是一样的，所以你可以试试这个：

sort -n -t, -k2 links.csv

-n表示按数字排序

-t,表示用逗号作为字段分隔符

-k2表示按第二个字段排序

这个命令会对文件进行排序，并将结果输出到标准输出（stdout）。你可以把结果重定向到另一个文件，或者通过管道传递给你的Python程序进行进一步处理。

编辑：如果你不想在运行Python脚本之前先对文件进行排序，你可以使用subprocess模块来创建一个管道，连接到shell的排序工具，然后从管道的输出中读取排序后的结果。

Answer 3

这里有一个关于如何处理比内存大的文件的方案，具体可以参考这个页面，不过你需要根据自己的情况，特别是涉及到CSV格式的数据，进行一些调整。页面上还有其他资源的链接。

补充：确实，磁盘上的文件并不是“比内存大”，但在内存中表示的数据可能会比可用内存大得多。首先，你的程序并不能使用整整1GB的内存（因为操作系统也占用内存等）。其次，即使你用最紧凑的方式存储这些数据（比如在32位机器上用两个整数列表），这30M对整数的存储也会占用934MB。

如果使用numpy，你只需要大约250MB就能完成这个任务。虽然这种加载方式不是特别快，因为你需要先计算行数并预先分配数组，但在内存中进行排序时，这可能是最快的实际排序方法：

import time
import numpy as np
import csv

start = time.time()
def elapsed():
    return time.time() - start

# count data rows, to preallocate array
f = open('links.csv', 'rb')
def count(f):
    while 1:
        block = f.read(65536)
        if not block:
             break
        yield block.count(',')

linecount = sum(count(f))
print '\n%.3fs: file has %s rows' % (elapsed(), linecount)

# pre-allocate array and load data into array
m = np.zeros(linecount, dtype=[('a', np.uint32), ('b', np.uint32)])
f.seek(0)
f = csv.reader(open('links.csv', 'rb'))
for i, row in enumerate(f):
    m[i] = int(row[0]), int(row[1])

print '%.3fs: loaded' % elapsed()
# sort in-place
m.sort(order='b')

print '%.3fs: sorted' % elapsed()

在我的机器上，使用一个与你提供的样本文件类似的文件时，输出结果如下：

6.139s: file has 33253213 lines
238.130s: read into memory
517.669s: sorted

numpy的默认排序算法是快速排序。ndarray.sort()这个函数（它是在原地排序）也可以接受关键字参数kind="mergesort"或kind="heapsort"，但这两种方法似乎都不能在记录数组上进行排序。顺便提一下，我使用记录数组是因为这是我能想到的唯一方法，可以将列一起排序，而不是默认的独立排序（这样会完全搞乱你的数据）。