# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure key skew and the effect of the partition key.

1,000,000 order events for 500 sellers, one of which owns 40% of them,
partitioned 12 ways by seller_id, by order_id, and by seller_id with a salt
for the hot seller; then two-phase aggregation of the salted counts.
"""
import argparse, json, random, zlib, collections

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--events", type=int, default=1_000_000, help="Order events.")
p.add_argument("--partitions", type=int, default=12, help="Partitions.")
p.add_argument("--hot-share", type=float, default=0.40, help="Share of events owned by the hot seller.")
p.add_argument("--salts", type=int, default=4, help="Salts for the hot seller.")
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)
N = args.partitions
def part(key): return zlib.crc32(key.encode()) % N

sellers = ["FreshBox" if random.random() < args.hot_share else f"s{random.randrange(499)}" for _ in range(args.events)]
keyings = {"seller_id": lambda s, i: s, "order_id": lambda s, i: f"order-{i}",
           "seller_id + salt for the hot seller": lambda s, i: f"{s}#{random.randrange(args.salts)}" if s == "FreshBox" else s}
rows, salted_counts = [], collections.Counter()
for name, fn in keyings.items():
    c = collections.Counter()
    for i, s in enumerate(sellers):
        k = fn(s, i); c[part(k)] += 1
        if name.startswith("seller_id +") and s == "FreshBox": salted_counts[k] += 1
    mean = args.events / N
    rows.append({"key": name, "max_over_mean": max(c.values()) / mean, "largest_share_pct": max(c.values()) / args.events * 100})
    print(f"{name:<38} max/mean = {rows[-1]['max_over_mean']:.2f}   largest partition {rows[-1]['largest_share_pct']:.1f}% of events")
true = sellers.count("FreshBox"); two_phase = sum(salted_counts.values())
print(f"FreshBox count: true {true:,}; two-phase (sum of {len(salted_counts)} salted partials) {two_phase:,}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"keys": rows, "summary": [{"true": true, "two_phase": two_phase}]}}, f)