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

时间:2014-05-03 21:42:59

标签: python sorting

我正在寻找完成这项工作,但遇到了困难:

我有大量的文字。每一行的格式为"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()

1 个答案:

答案 0 :(得分:6)

您的解决方案几乎是正确的。但是,在将 它们写入磁盘之前,每个部分文件必须排序。这是一个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))