# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The Zipf ceiling on data-parallel speedup.

Zipf-distributed keys hashed to partitions; the best possible speedup is
total work divided by the largest partition.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--key-counts", type=lambda s: [int(v) for v in s.split(",")], default=[1_000, 10_000, 100_000, 1_000_000],
               help="Distinct keys to try, comma-separated.")
p.add_argument("--rows", type=int, default=20_000_000, help="Rows generated per key count.")
p.add_argument("--parts", type=int, default=1000, help="Partitions the keys are hashed to.")
p.add_argument("--s", type=float, default=1.0, help="Zipf exponent: below 1 is flatter, above 1 is steeper.")
p.add_argument("--seed", type=int, default=1, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)


def ceiling(n_keys, rows, parts, s):
    w = 1.0 / np.arange(1, n_keys + 1) ** s       # Zipf weights
    w /= w.sum()
    keys = rng.choice(n_keys, size=rows, p=w)
    part = np.bincount(keys % parts, minlength=parts)   # rows per partition
    return {"n_keys": n_keys, "top_key_share_pct": float(w[0] * 100), "one_over_f": float(1 / w[0]),
            "largest_partition_pct": float(part.max() / rows * 100), "speedup_cap": float(rows / part.max())}


rows = []
for n in args.key_counts:
    r = ceiling(n, args.rows, args.parts, args.s)
    rows.append(r)
    print(f"n_keys={n:>10,}  top key share={r['top_key_share_pct']:.3f}%  1/f={r['one_over_f']:.1f}  "
          f"largest partition={r['largest_partition_pct']:.3f}%  speedup cap={r['speedup_cap']:.1f}")

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