Problem 99 · Data Engineering & Streaming

Distributed Group-By with Partial Aggregates

MediumClassic DSA & systems

01Problem

Compute per-key count, sum, min, max and avg over rows spread across partitions using a map-side combine, a hash shuffle to R reducers, and a final reduce.

02Key insight

Only algebraic aggregates can be combined early: AVG is carried as (sum, count) and finalized at the end. Map-side combining shrinks shuffle volume from rows to distinct keys per partition; the merge function must be associative and commutative so partial order doesn't matter. Use a stable hash for partitioning (Python's hash() is salted per process).

03Solution

solution.pypython39 lines

import zlib
from collections import defaultdict
from typing import Any, Callable, Iterable

Partial = tuple[int, float, float, float]  # count, sum, min, max


def combine(rows: Iterable[dict], key: str, value: str) -> dict[Any, Partial]:
    acc: dict[Any, Partial] = {}
    for row in rows:
        k, v = row[key], row[value]
        acc[k] = merge(acc[k], (1, v, v, v)) if k in acc else (1, v, v, v)
    return acc


def merge(a: Partial, b: Partial) -> Partial:  # associative + commutative
    return a[0] + b[0], a[1] + b[1], min(a[2], b[2]), max(a[3], b[3])


def stable_partition(k: Any, num_reducers: int) -> int:
    return zlib.crc32(repr(k).encode()) % num_reducers


def distributed_group_by(partitions: list[list[dict]], key: str, value: str,
                         num_reducers: int = 3,
                         partitioner: Callable[[Any, int], int] = stable_partition) -> dict[Any, dict]:
    shuffle: list[dict[Any, list[Partial]]] = [defaultdict(list) for _ in range(num_reducers)]
    for part in partitions:                       # map + combine, then shuffle
        for k, partial in combine(part, key, value).items():
            shuffle[partitioner(k, num_reducers)][k].append(partial)
    result: dict[Any, dict] = {}
    for bucket in shuffle:                         # reduce + finalize
        for k, partials in bucket.items():
            total = partials[0]
            for p in partials[1:]:
                total = merge(total, p)
            n, s, lo, hi = total
            result[k] = {"count": n, "sum": s, "min": lo, "max": hi, "avg": s / n}
    return result

Try it first.

TimeO(n) total; shuffle O(partitions * distinct keys)

SpaceO(distinct keys per partition)

tests.py · 2 passing

tests.pypython10 lines

import random
rng = random.Random(5)
rows = [{"region": rng.choice("ABCD"), "amt": rng.randint(1, 100)} for _ in range(2000)]
parts = [rows[i::4] for i in range(4)]
res = distributed_group_by(parts, "region", "amt")
for r in "ABCD":
    vals = [x["amt"] for x in rows if x["region"] == r]
    assert res[r] == {"count": len(vals), "sum": sum(vals), "min": min(vals),
                      "max": max(vals), "avg": sum(vals) / len(vals)}
assert distributed_group_by([[], []], "region", "amt") == {}

04Staff-level follow-up

What about COUNT(DISTINCT) or MEDIAN, which aren't algebraic? Either shuffle raw values (expensive) or use mergeable sketches - HLL for distinct count, t-digest/KLL for quantiles - as the partial aggregate.