用Python的heapq.merge排序大文件

Question

我想完成一个工作，但遇到了一些困难：

我有一个很大的文本文件。每一行的格式是 "AGTCCCGGAT 文件名"，其中第一部分是跟DNA有关的东西。

教授建议我们把这个大文件分成很多临时文件，然后用 heapq.merge() 来排序。最终的目标是得到一个文件，里面包含原始文件的每一行，并且是排好序的。

我第一次尝试是把每一行都放到一个单独的临时文件里。问题是 heapq.merge() 报告说要排序的文件太多了。

我第二次尝试是每50000行分成一个临时文件。问题是，它似乎不是按行排序，而是按文件排序。例如，我们有这样的情况：

ACGTACGT filename
CGTACGTA filename
ACGTCCGT filename
CGTAAAAA filename

前两行来自一个临时文件，后两行来自第二个文件。

我用来排序的代码如下：

for line in heapq.merge(*[open('/var/tmp/L._Ipsum-strain01.fa_dir/'+str(f),'r') for f in os.listdir('/var/tmp/L._Ipsum-strain01.fa_dir')]):
     result.write(line)
result.close()

Answer 1

你的解决方案差不多是对的。不过，每个部分的文件在写入磁盘之前必须先进行排序。这里有一个两次处理的算法来演示这个过程：首先，把文件分成每次处理5万行，对这些行进行排序，然后把排序好的部分写入一个文件。第二次处理时，打开所有这些文件，把它们合并到输出文件中。

from heapq import merge
from itertools import count, islice
from contextlib import ExitStack  # not available on Python 2
                                  # need to care for closing files otherwise

chunk_names = []

# chunk and sort
with open('input.txt') as input_file:
    for chunk_number in count(1):
        # read in next 50k lines and sort them
        sorted_chunk = sorted(islice(input_file, 50000))
        if not sorted_chunk:
            # end of input
            break

        chunk_name = 'chunk_{}.chk'.format(chunk_number)
        chunk_names.append(chunk_name)
        with open(chunk_name, 'w') as chunk_file:
            chunk_file.writelines(sorted_chunk)

with ExitStack() as stack, open('output.txt', 'w') as output_file:
    files = [stack.enter_context(open(chunk)) for chunk in chunk_names]
    output_file.writelines(merge(*files))