# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "datasketches"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Do not average percentiles.

100 shards of latencies from different log-normal distributions and sizes
that differ up to 100 times. The p99 of the whole fleet three ways: the
average of the shards' p99s, the exact p99 of the pooled values, and merged
sketches: KLL, t-digest, and REQ in high-rank-accuracy mode.
"""
import argparse, json
import numpy as np
from datasketches import kll_floats_sketch, req_floats_sketch, tdigest_double

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--shards", type=int, default=100, help="Shards.")
p.add_argument("--base-size", type=int, default=10_000, help="Values per shard when sizes are equal.")
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)

rows = []
for label, sizes in (("equal sizes", [args.base_size] * args.shards),
                     ("sizes differ up to 100x", list(np.geomspace(1_000, 100_000, args.shards).astype(int)))):
    shards = [rng.lognormal(np.log(20) + rng.normal(0, 0.5), rng.uniform(0.3, 1.2), n) for n in sizes]
    avg = float(np.mean([np.percentile(s, 99) for s in shards]))
    exact = float(np.percentile(np.concatenate(shards), 99))
    est = {}
    for name, make in (("kll", lambda: kll_floats_sketch(200)), ("tdigest", lambda: tdigest_double(200)),
                       ("req", lambda: req_floats_sketch(30, True))):
        merged = make()
        for s in shards:
            k = make()
            for x in s: k.update(float(x))
            merged.merge(k)
        est[name] = merged.get_quantile(0.99)
    rows.append({"shards": label, "average_of_p99s": avg, "exact_p99": exact, **{f"{k}_p99": v for k, v in est.items()},
                 "avg_err_pct": (avg / exact - 1) * 100, **{f"{k}_err_pct": (v / exact - 1) * 100 for k, v in est.items()}})
    print(f"{label:<24} exact p99 {exact:7.1f} ms   average of p99s {avg:7.1f} ({avg / exact - 1:+.0%})   merged: "
          + "   ".join(f"{k} {v:7.1f} ({v / exact - 1:+.2%})" for k, v in est.items()))

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