# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""When does partial aggregation pay?

N rows over P partitions, grouped by a key with G distinct values. A local
pre-aggregation on one partition reduces its rows by per-partition rows ÷
local groups. Compares the measured reduction with the bound N/(P·G) and the
expected local groups G·(1 − e^(−(N/P)/G)); then times a Counter on one
partition against simply forwarding its rows.
"""
import argparse, json, math, random, time
from collections import Counter

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=2_000_000, help="Rows.")
p.add_argument("--p", type=int, default=20, help="Partitions.")
p.add_argument("--groups", default="100,10000,200000,4000000", help="Values of G.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)

per = args.n // args.p; rows = []
for G in map(int, args.groups.split(",")):
    keys = [random.randrange(G) for _ in range(per)]                   # one partition
    t = time.perf_counter(); c = Counter(keys); agg = time.perf_counter() - t
    t = time.perf_counter(); fwd = list(keys); fw = time.perf_counter() - t
    local = len(c); exp = G * (1 - math.exp(-per / G))
    rows.append({"G": G, "bound": args.n / (args.p * G), "measured": per / local, "local_groups": local, "expected_groups": exp,
                 "aggregate_ms": agg * 1000, "forward_ms": fw * 1000})
for r in rows: print(f"G={r['G']:>9,}  N/(P*G)={r['bound']:9.2f}  measured={r['measured']:8.2f}  local groups={r['local_groups']:,} (expected ~{r['expected_groups']:,.0f})  Counter {r['aggregate_ms']:.0f} ms vs forward {r['forward_ms']:.1f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"groups": rows}}, f)