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

Hash partitioning over 8, 32 and 128 instances when one key carries 30% of
the traffic and the rest is spread over 100,000 keys: the busiest instance's
share of the total. Then the hot key salted into 8 sub-keys. Finally a count
and a median over the salted sub-keys, combined without a second pass.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--events", type=int, default=2_000_000, help="Events.")
p.add_argument("--hot-share", type=float, default=0.3, help="Share of traffic on the hot key.")
p.add_argument("--keys", type=int, default=100_000, help="Other keys.")
p.add_argument("--salts", type=int, default=8, help="Sub-keys for the hot key.")
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()

rng = np.random.default_rng(args.seed)
hot = rng.random(args.events) < args.hot_share
keys = np.where(hot, 0, rng.integers(1, args.keys + 1, args.events)).astype(np.uint64)
salt = np.where(hot, rng.integers(0, args.salts, args.events), 0).astype(np.uint64)
def h(x):
    z = x * np.uint64(0x9E3779B97F4A7C15); z ^= z >> np.uint64(31); z *= np.uint64(0xBF58476D1CE4E5B9); return z ^ (z >> np.uint64(29))
rows = []
with np.errstate(over="ignore"):
    for n in (8, 32, 128):
        plain = np.bincount((h(keys) % np.uint64(n)).astype(np.int64), minlength=n)
        salted = np.bincount((h(keys * np.uint64(1000) + salt) % np.uint64(n)).astype(np.int64), minlength=n)
        rows.append({"instances": n, "busiest_pct": plain.max() / args.events * 100, "salted_busiest_pct": salted.max() / args.events * 100,
                     "even_pct": 100 / n, "max_speedup": args.events / plain.max(), "salted_max_speedup": args.events / salted.max()})
values = rng.lognormal(3, 1, args.events)
hv, hs = values[hot], salt[hot]
parts = [hv[hs == s] for s in range(args.salts)]
cnt_ok = sum(len(x) for x in parts) == len(hv)
true_med = float(np.median(hv)); med_of_meds = float(np.median([np.median(x) for x in parts]))
summary = {"count_combines": "yes" if cnt_ok else "no", "true_median": true_med, "median_of_medians": med_of_meds,
           "median_err_pct": (med_of_meds / true_med - 1) * 100, "cap": 1 / args.hot_share}
for r in rows: print(f"{r['instances']:>3} instances: busiest {r['busiest_pct']:5.1f}% (even {r['even_pct']:.1f}%), max speedup {r['max_speedup']:.1f}x; salted ×{args.salts}: busiest {r['salted_busiest_pct']:5.1f}%")
print(f"salted count adds up exactly: {summary['count_combines']}; median of the {args.salts} partial medians {med_of_meds:.2f} vs true {true_med:.2f}")

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