Title: Sorting big files the Python 2.4 way
Submitter: Nicolas Lehuen (other recipes)
Last Updated: 2006/04/13
Version no: 1.5
Category: Files

Description:

This recipe can be used to sort big files (much bigger than the available RAM) according to a key. The sort is guaranteed to be stable on Python 2.3.

Source: Text Source

from heapq import heapify, heappop, heappush
from itertools import islice, cycle
from tempfile import gettempdir
import os

def merge(chunks,key=None):
    if key is None:
        key = lambda x : x

    values = []

    for index, chunk in enumerate(chunks):
        try:
            iterator = iter(chunk)
            value = iterator.next()
        except StopIteration:
            try:
                chunk.close()
                os.remove(chunk.name)
                chunks.remove(chunk)
            except:
                pass
        else:
            heappush(values,((key(value),index,value,iterator,chunk)))

    while values:
        k, index, value, iterator, chunk = heappop(values)
        yield value
        try:
            value = iterator.next()
        except StopIteration:
            try:
                chunk.close()
                os.remove(chunk.name)
                chunks.remove(chunk)
            except:
                pass
        else:
            heappush(values,(key(value),index,value,iterator,chunk))

def batch_sort(input,output,key=None,buffer_size=32000,tempdirs=[]):
    if not tempdirs:
        tempdirs.append(gettempdir())
    
    input_file = file(input,'rb',64*1024)
    try:
        input_iterator = iter(input_file)
        
        chunks = []
        try:
            for tempdir in cycle(tempdirs):
                current_chunk = list(islice(input_iterator,buffer_size))
                if current_chunk:
                    current_chunk.sort(key=key)
                    output_chunk = file(os.path.join(tempdir,'%06i'%len(chunks)),'w+b',64*1024)
                    output_chunk.writelines(current_chunk)
                    output_chunk.flush()
                    output_chunk.seek(0)
                    chunks.append(output_chunk)
                else:
                    break
        except:
            for chunk in chunks:
                try:
                    chunk.close()
                    os.remove(chunk.name)
                except:
                    pass
            if output_chunk not in chunks:
                try:
                    output_chunk.close()
                    os.remove(output_chunk.name)
                except:
                    pass
            return
    finally:
        input_file.close()
    
    output_file = file(output,'wb',64*1024)
    try:
        output_file.writelines(merge(chunks,key))
    finally:
        for chunk in chunks:
            try:
                chunk.close()
                os.remove(chunk.name)
            except:
                pass
        output_file.close()

if __name__ == '__main__':
    import optparse
    parser = optparse.OptionParser()
    parser.add_option(
        '-b','--buffer',
        dest='buffer_size',
        type='int',default=32000,
        help='''Size of the line buffer. The file to sort is
            divided into chunks of that many lines. Default : 32,000 lines.'''
    )
    parser.add_option(
        '-k','--key',
        dest='key',
        help='''Python expression used to compute the key for each
            line, "lambda line:" is prepended.\n
            Example : -k "line[5:10]". By default, the whole line is the key.'''
    )
    parser.add_option(
        '-t','--tempdir',
        dest='tempdirs',
        action='append',
        default=[],
        help='''Temporary directory to use. You might get performance
            improvements if the temporary directory is not on the same physical
            disk than the input and output directories. You can even try
            providing multiples directories on differents physical disks.
            Use multiple -t options to do that.'''
    )
    parser.add_option(
        '-p','--psyco',
        dest='psyco',
        action='store_true',
        default=False,
        help='''Use Psyco.'''
    )
    options,args = parser.parse_args()
    
    if options.key:
        options.key = eval('lambda line : (%s)'%options.key)
    
    if options.psyco:
        import psyco
        psyco.full()
    
    batch_sort(args[0],args[1],options.key,options.buffer_size,options.tempdirs)

Discussion:

The memory used is at most buffer_size * maximum length of line + a 64 kb buffer per chunk (there are number of lines in the file / buffer_size chunks). In practice one should try to provide big buffer_size numbers since the limiting factor is often the maximum numbers of file a process can open simultaneously. The merge process does a n-way merge between all the chunks, I should try to perform partial mergings to work around this limitation.

The Python Cookbook already contains a big file sort implementation (look for "big file sorting"), but I feel that this version is more terse and Pythonic.

Usage example :

sort.py -k "(line[5:10],int(line[10:16])" input.txt output.txt

This sort lexicographically on the text found at columns 5 to 9 inclusive, then numerically on the number found at column 10 to 15 inclusive.

Performance :

See by yourself :) I find this script quite competitive thanks to :

1) the awesome speed of Python's sort
2) the not less awesome speed of Python heapq
3) usage of file.writelines
4) a trick I've found to implement the merge sort, which is a variant of the Schwartzian transform wherein the iterator is included in the transformed tuple.

Update (1.1) : integrated changes by Ian Bicking.

Update (1.2) : added the possibility of specifying temporary directories, which can give a performance boost if they reside on different physical disks than the input or output files.

Update (1.3,1.4) : fixed typos.

Update (1.5) : This version now does its best to cleanup the temporary files even when an error is raised during the sort.