Group rows by a key

GROUP BY in SQL, groupby in pandas, a shuffle in Spark. Under all of them sits a dict of lists. Build that dict.

Write group_by(rows, key):

• rows is a list of dicts.
• key is the field name to group on.
• Return a dict that maps each distinct value of row[key] to the list of rows with that value.
• Keep rows in the order they appeared in the input.
• Skip rows that don’t have the key at all.
• An empty input returns {}.

>>> group_by([{"city": "Oslo", "id": 1}, {"city": "Lima", "id": 2}, {"city": "Oslo", "id": 3}], "city")
{'Oslo': [{'city': 'Oslo', 'id': 1}, {'city': 'Oslo', 'id': 3}], 'Lima': [{'city': 'Lima', 'id': 2}]}

Approach

One pass. For each row, look up its group and append. setdefault creates the empty list on first sight.

def group_by(rows, key):
    groups = {}
    for row in rows:
        if key not in row:
            continue
        groups.setdefault(row[key], []).append(row)
    return groups

collections.defaultdict(list) does the same job. Convert it back with dict(...) if callers shouldn’t see the default behaviour.

Complexity

O(n) time: one hash lookup and one append per row. O(n) extra space for the lists, though the rows themselves are shared, not copied.

Why this matters

A hash aggregate in a database is this loop: hash the grouping key, find the bucket, update it. Distributed engines do the same thing twice. First they send every row with the same key to the same machine (the shuffle), then each machine runs this loop locally. When the dict no longer fits in memory, the engine spills buckets to disk. That spill is why a large GROUP BY can suddenly get slow.