Problem 89 · Data Engineering & Streaming

External Merge Sort

HardClassic DSA & systems

01Problem

Sort a stream of integers too large for memory, given that at most chunk_size items fit in RAM: produce sorted runs on disk, then k-way merge them lazily.

02Key insight

Two phases: sort fixed-size chunks into spilled runs, then a heap-based k-way merge that streams with O(k) memory. With many runs, do multi-pass merges (fan-in limited by file handles/buffer size) - this is exactly how databases spill ORDER BY.

03Solution

solution.pypython44 lines

import heapq
import itertools
import os
import tempfile
from contextlib import ExitStack
from typing import Iterable, Iterator


def external_sort(items: Iterable[int], chunk_size: int, max_fan_in: int = 64) -> Iterator[int]:
    if max_fan_in < 2:
        raise ValueError("max_fan_in must be >= 2")
    with tempfile.TemporaryDirectory() as tmp:
        ids = itertools.count()

        def write_run(values: Iterable[int]) -> str:
            path = os.path.join(tmp, f"run{next(ids)}")
            with open(path, "w") as f:
                f.writelines(f"{x}\n" for x in values)
            return path  # closed on disk: runs don't hold file handles

        def merged(stack: ExitStack, paths: list[str]) -> Iterator[int]:
            files = [stack.enter_context(open(p)) for p in paths]
            return heapq.merge(*((int(line) for line in f) for f in files))  # one head per run

        def merge_to_run(paths: list[str]) -> str:
            with ExitStack() as stack:
                out = write_run(merged(stack, paths))
            for p in paths:
                os.remove(p)
            return out

        runs: list[str] = []
        chunk: list[int] = []
        for x in items:
            chunk.append(x)
            if len(chunk) >= chunk_size:
                runs.append(write_run(sorted(chunk)))
                chunk.clear()
        if chunk:
            runs.append(write_run(sorted(chunk)))
        while len(runs) > max_fan_in:  # intermediate passes: at most max_fan_in inputs open
            runs = [merge_to_run(runs[i:i + max_fan_in]) for i in range(0, len(runs), max_fan_in)]
        with ExitStack() as stack:
            yield from merged(stack, runs)

Try it first.

TimeO(n log n) CPU, O(n) IO per pass

SpaceO(chunk_size + k) memory

tests.py · 5 passing

tests.pypython24 lines

import random
rnd = random.Random(1)
data = [rnd.randint(-1000, 1000) for _ in range(5003)]
assert list(external_sort(data, chunk_size=100)) == sorted(data)
assert list(external_sort([], chunk_size=10)) == []
assert list(external_sort(iter([3, 1, 2]), chunk_size=1)) == [1, 2, 3]
import builtins
live, peak = [0], [0]
class _Tracked:  # counts input runs held open concurrently
    def __init__(self, f, reading): self.f, self.reading = f, reading
    def __enter__(self):
        live[0] += self.reading; peak[0] = max(peak[0], live[0])
        return self.f.__enter__()
    def __exit__(self, *exc):
        live[0] -= self.reading
        return self.f.__exit__(*exc)
def open(path, mode="r", *a, **kw):
    return _Tracked(builtins.open(path, mode, *a, **kw), "r" in mode)
big = [rnd.randint(-10**6, 10**6) for _ in range(1000)]
for fan_in in (64, 8, 2):
    peak[0] = 0
    assert list(external_sort(big, chunk_size=2, max_fan_in=fan_in)) == sorted(big)  # 500 runs
    assert 0 < peak[0] <= fan_in
del open

04Staff-level follow-up

Data is 10 TB across a cluster - how do you sort it? Sample to pick range-partition boundaries, shuffle each record to its range's node, sort locally (external sort), and concatenate partitions in order (TeraSort); skew is handled by finer sampling.