使用Python heapq.merge对大文件进行排序

Question

我想完成这样的工作,但遇到了困难:

我有一大堆文本.每一行的格式"AGTCCCGGAT filename"都是第一部分是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

您的解决方案几乎正确.但是,必须先对每个部分文件进行排序,然后再将其写入磁盘.这是一个2遍算法,用于演示它:首先,以50k行块的形式迭代文件,对块中的行进行排序,然后将这个已排序的块写入文件.在第二遍中,打开所有这些文件并合并到输出文件.

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))