← back to the archiveCover illustration for “Sort and materialise Python groupby batches”
TILday 136·today·Published ·by Andy Padia

Sort and materialise Python groupby batches

In short: Sort finite Python batches by the grouping key and consume each group immediately, so repeated keys and shared iterators cannot hide rows.

python ▸ from itertools import groupby
from operator import itemgetter

def group_batch(rows):
    # Finite batch, string keys; preserve within-key input order.
    key = itemgetter("document")
    ordered = sorted(rows, key=key)
    return {name: list(group) for name, group in groupby(ordered, key)}

if __name__ == "__main__":
    rows = [
        {"document": "a", "chunk": 1},
        {"document": "b", "chunk": 2},
        {"document": "a", "chunk": 3},
    ]
    grouped = group_batch(rows)
    assert [r["chunk"] for r in grouped["a"]] == [1, 3]
    assert sum(map(len, grouped.values())) == len(rows)
    print({k: [r["chunk"] for r in v] for k, v in grouped.items()})

The dictionary had the right two document names and the wrong contents. With input keys a, b, a, a local Python 3.14.6 check on 14 October 2026 retained only two of three rows after a naive groupby dictionary comprehension. The second a replaced the first group.

Sort and materialise Python groupby batches when the requirement is one retained group per key across a finite input. Sorting brings matching keys together; immediate materialisation keeps each group's rows available afterwards. My review rule is to check row conservation and delayed access, not merely the number of dictionary keys.

Python groupby recognises runs, not complete documents

The itertools reference specifies consecutive grouping: a changed key starts another group. For a, b, a, three runs are correct behaviour. Turning those runs into a dictionary is where the earlier a disappears in this example.

There is a separate lifetime constraint. The group iterators share their underlying source. Advancing the outer iterator makes a previous unconsumed group unavailable. Saving list(groupby(...)) therefore saves pairs containing iterators, not independent copies of their contents.

That gives us two different tests. Interleave matching keys to expose the ordering assumption. Then try reading saved groups after the outer loop finishes to expose the lifetime assumption. Passing one does not establish the other.

Sort and materialise Python groupby batches explicitly

Save this as group_batch.py and run python3 group_batch.py. It uses synthetic rows and the standard library. The helper assumes a finite batch with string document keys that comfortably fits in memory.

from itertools import groupby
from operator import itemgetter

def group_batch(rows):
    # Finite batch, string keys; preserve within-key input order.
    key = itemgetter("document")
    ordered = sorted(rows, key=key)
    return {name: list(group) for name, group in groupby(ordered, key)}

if __name__ == "__main__":
    rows = [
        {"document": "a", "chunk": 1},
        {"document": "b", "chunk": 2},
        {"document": "a", "chunk": 3},
    ]
    grouped = group_batch(rows)
    assert [r["chunk"] for r in grouped["a"]] == [1, 3]
    assert sum(map(len, grouped.values())) == len(rows)
    print({k: [r["chunk"] for r in v] for k, v in grouped.items()})

The same key function governs sorting and grouping. Python's sorting guide guarantees stable sorting, so rows with an equal key retain their input order. The resulting a group contains chunks 1, 3, while b contains chunk 2.

This preserves arrival order within a document, not necessarily document order. If chunks arrive out of sequence, sort or validate their sequence separately. A grouping helper cannot infer the intended order from an identifier.

What the local groupby tests established

The exact snippet ran successfully. Additional checks used every permutation of the three synthetic rows, plus empty input, a generator and duplicate rows.

StrategyInput rowsRetained rows
Unsorted groups into a dictionary32
Sorted groups, materialised immediately33
Sorted group iterators read after outer exhaustion30

Local Python 3.14.6 observations; the last row deliberately stores unconsumed iterators.

All six permutations retained every row with the helper, preserving within-key input order. Duplicate input rows remained duplicates. A missing document field raised KeyError; there is no invented fallback identity. These checks do not measure throughput or validate a production retrieval system.

My judgement for a retrieval batch

For a hypothetical retrieval pipeline assembling chunks by document, I would make the conservation assertion part of the test fixture. A dictionary with the expected document names can still be missing evidence. Counting groups alone gives false comfort.

I would also keep the batch boundary explicit. This helper sorts the entire input and stores its groups, so it is not an unbounded streaming solution. Already-clustered input may not need sorting; immediate aggregation may not need retained lists. Choose those alternatives when their assumptions are actually enforced.

The lists are not deep copies of the row dictionaries. Mutating a row later remains visible through its group. Input validation, snapshot isolation and deduplication are separate requirements, not benefits supplied by list(group).

What's in it for you

  • Test interleaved keys and retained-group access, not just tidy pre-sorted fixtures.
  • Verify that grouping preserved every expected row before handing the batch downstream.

A group is useful only if its rows survive both the ordering and the iteration.

Sources

#python#testing#data-pipelines
← older drop
Agent memory is a control plane, not a vector replacement

related drops

explore all 391 drops →
← back to the archiveday 136