External "bucket" sort

Question

I'm working on a CTags plugin for Sublime Text. One of the open issues revolves around the fact that an in-memory sort of a tag file can cause some form of memory or buffer overflow errors (the stack in Sublime Text is limited to ~25 MB apparently). I've devised a simple external bucket sort to resolve this issue:

#/usr/bin/env python
#
# CSV external bucket sort

import codecs
import tempfile
import os

# column indexes
SYMBOL = 0
FILENAME = 1

TAG_FILE = 'tags'
OUT_FILE = ''.join([TAG_FILE, '_sorted_by_file'])

def sort():
    """External bucket sort of tab delimited CTag files"""
    temp_files = {}

    def get_file(filename):
        """Get a file from the store of files"""
        if not filename in temp_files:
            temp_files[filename] = tempfile.NamedTemporaryFile(delete=False)
            # close and reopen using codecs to avoid problems described here:
            #   http://stackoverflow.com/a/10490859/613428
            temp_files[filename].close()
            temp_files[filename] = codecs.open(
                temp_files[filename].name, 'w+', 'utf-8', 'ignore')
        return temp_files[filename]

    try:
        with codecs.open(TAG_FILE, 'r+', 'utf-8', 'ignore') as file_o:
            for _ in range(6):  # skip the header
                next(file_o)
            for line in file_o:
                temp_file_o = get_file(line.split('\t')[FILENAME])
                split = line.split('\t')
                split[FILENAME] = split[FILENAME].lstrip('.\\')
                temp_file_o.write('\t'.join(split))

        with codecs.open(OUT_FILE, 'w+', 'utf-8', 'ignore') as file_o:
            # we only need to sort the file names - the symbols were already
            # sorted!
            for key in sorted(temp_files):
                temp_files[key].seek(0)
                file_o.write(temp_files[key].read())
    finally:
        for key in temp_files:
            temp_files[key].close()
            os.remove(temp_files[key].name)

    os.remove(OUT_FILE)  # just for testing - remove when done

Here's what I intended to do:

For each line (a.k.a. tag) in a tag file
    Read line into memory
    Get the filename from the line.
    Use filename as key to a temp file and write whole this line to said file
Sort final list of keys (i.e. filenames)
For each file in sorted key list
    Append contents of file to sorted tag file
Close all files and delete temp files

And the performance:

$ python -m timeit -n 100 'import external_sort; external_sort.sort()'
100 loops, best of 3: 12 msec per loop

My question is this: is this as performant as it could be? Could I improve it in any way, e.g. missing corner cases?

NOTE: I'm aware that it's not really a true bucket sort as the buckets aren't arbitrary or equally sized (a project with many small files and few large files would result in unbalanced buckets). However, I'm counting on there being many average sized files and none big-enough to fill memory by themselves.

Answer 1

• It is not necessary to decode and encode UTF-8, as split('\t') and lstrip('.\\') work fine on UTF-8-encoded strings.
• The programs keeps an unlimited number of temporary files open simultaneously and will fail after exceeding the platform's maximum for open files. An approach where only a certain number of most recently used files are kept open could be appropriate here.
• file_o.write(temp_files[key].read()) reads the entire file in memory. To conserve memory, use a loop and give a block size argument to read.