# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "matplotlib"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Triage plot, salting and the join product.

Three kinds of long tail on a duration-against-bytes plot; speculation against
salting for the hot key; and why join skew multiplies.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--tasks", type=int, default=1000, help="Tasks in the stage.")
p.add_argument("--hot-share", type=float, default=0.07, help="Case C: share of all bytes in the one hot task.")
p.add_argument("--slow-factor", type=float, default=10.0, help="Case B: how much slower the one slow task is.")
p.add_argument("--salts", type=int, default=16, help="k: salts the hot key is split into.")
p.add_argument("--join-skew", type=float, default=10.0, help="A key's multiple of its fair share, on each side of the join.")
p.add_argument("--plot", help="Save the triage scatter plots to this PNG file.")
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)
n = args.tasks
MB_PER_S = 2.0                                          # processing rate: duration = bytes / rate

# 1. Three long tails
# A: equal bytes, random durations
a_bytes = np.full(n, 100.0)
a_secs = rng.lognormal(np.log(50), 0.25, n)
# B: equal bytes, one task many times slower (a slow machine)
b_bytes = np.full(n, 100.0)
b_secs = b_bytes / MB_PER_S * rng.normal(1, 0.05, n)
b_secs[0] *= args.slow_factor
# C: durations proportional to bytes, one task holding hot_share of all bytes (a hot key)
c_bytes = rng.normal(100, 10, n).clip(10)
c_bytes[0] = args.hot_share * c_bytes[1:].sum() / (1 - args.hot_share)
c_secs = c_bytes / MB_PER_S * rng.normal(1, 0.05, n)

cases = {"A · random durations": (a_bytes, a_secs), "B · one slow task": (b_bytes, b_secs),
         "C · one hot key": (c_bytes, c_secs)}
tasks, tails = [], []
for name, (bts, secs) in cases.items():
    med = np.median(secs)
    print(f"{name:<22} slowest task {secs.max():6.0f} s = {secs.max() / med:4.1f}x median, "
          f"holding {bts[secs.argmax()]:,.0f} MB ({bts[secs.argmax()] / np.median(bts):.1f}x median bytes)")
    tasks += [{"case": name, "mb": float(b), "seconds": float(s)} for b, s in zip(bts, secs)]
    tails.append(float(secs.max() / med))

# 2. Case C: speculation against salting
stage = c_secs.max()
speculated = max(c_secs.max(), c_secs[0])               # the duplicate reads the same bytes: as slow as the original
salted = c_secs.copy()
salted[0] = c_secs[0] / args.salts                       # the hot key split into k salts ...
salted = np.append(salted, [salted[0]] * (args.salts - 1))   # ... k smaller tasks
print(f"\nhot task: {c_secs[0]:.0f} s   with speculation: stage {speculated:.0f} s (unchanged)   "
      f"salted k={args.salts}: largest task {c_secs[0] / args.salts:.0f} s ({c_secs[0] / (c_secs[0] / args.salts):.0f}x smaller), "
      f"stage {salted.max():.0f} s")

# 3. The join product: one key's work is L x R
s = args.join_skew
print(f"\njoin: a key at {s:g}x its fair share on each side -> each side {s:g}x too big, "
      f"work L x R = {s * s:g}x the fair work")

if args.plot:
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 3, figsize=(12, 3.5))
    for ax, (name, (bts, secs)) in zip(axes, cases.items()):
        ax.scatter(bts, secs, s=6); ax.set_title(name); ax.set_xlabel("MB read"); ax.set_ylabel("seconds")
    fig.tight_layout(); fig.savefig(args.plot)

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "tasks": tasks,
            "fix": [{"hot_task_s": float(c_secs[0]), "stage_s": float(stage), "speculated_stage_s": float(speculated),
                     "salted_largest_task_s": float(c_secs[0] / args.salts), "salted_stage_s": float(salted.max())}],
            "join": [{"skew_each_side": s, "work_multiple": s * s}],
            "summary": [{"a_tail": tails[0], "b_tail": tails[1], "c_tail": tails[2], "hot_mb": float(c_bytes[0])}],
        }}, f)