Skip to content

← Laziness, Iteration and Pipelines step 12 of 14

Hard End-to-End

The streaming sort-merge join

This is the canonical “correct algorithm, wrong memory profile” problem.

Everybody writes the hash join first: build {row[on]: row for row in right}, then scan the left. It is O(n + m), it is short, it passes every test you would think to write — and it holds the entire right side in memory. When right is a 200 GB export, the job is killed and the fix is not a tweak, it is a different algorithm.

If both inputs are already sorted by the join key — which they are, coming out of a sorted export or an ORDER BY — the sort-merge join holds exactly one key group from each side. Peak memory becomes the size of the largest key group, not the size of the input.

What to write

def _grouped(rows: Iterable[Row], on: str) -> Iterator[tuple[str, list[Row]]]
def stream_join(left: Iterable[Row], right: Iterable[Row], on: str) -> Iterator[Row]

Row is dict[str, str].

_grouped yields (key, rows_with_that_key) for each run of adjacent equal keys. It buffers one group at a time and nothing else. (This is groupby‘s job in principle, but groupby‘s group iterator is invalidated the moment you advance the outer one, which is exactly the trap a merge join walks into — writing it yourself is three lines and no surprises.)

stream_join is an inner join:

  • Advance whichever side has the smaller key. Equal keys emit the full cartesian product of the two groups: for each left row, for each right row, in order. Many-to-many is the case people forget, and it is the case that turns a duplicate key upstream into a row explosion downstream.
  • Each output row is {**left_row, **right_row}.
  • Output is ordered by key, then left-major, then right.
  • Iterate each argument exactly once, and never materialise either one.

Two hazards the type system will not catch

An iterator satisfies Iterable. Passing the same iterator as both left and right type-checks perfectly and produces nonsense, because the two groupers race each other through one shared position. Iterable[Row] cannot express “and it must be independently iterable”. The left_iters / right_iters counters in the report are the mechanical version of that review comment: exactly one __iter__ per argument.

Sorted-ness is a precondition, not a type. Nothing in Iterable[Row] says “ascending by on“. If the inputs are not sorted the join silently drops rows, which is the same failure class as the silent zip truncation. Document it, and put it in the docstring where the caller will see it.

What the report proves

  • rows — correctness, including many-to-many and non-overlapping keys.
  • pulled_before_first0. stream_join is a generator function.
  • left_pulled / right_pulled — exact pull counts. A grouper must read one row past the end of a group to know the group ended; that off-by-one is real and the expected numbers include it.
  • left_iters / right_iters — exactly 1 each.
  • With infinite=True both sides continue forever with matching keys, and a bounded consumer takes four rows. A solution that calls list(), tuple() or sorted() on either input never returns.

    Loading visualization…