# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "lz4"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""When does shuffle compression pay? And a non-deterministic partition function.

Part A: codec throughput C and ratio r on shuffle-like data, against the rule
C > BW * r/(r-1). Part B: a partitioner with a random term, retried.
"""
import argparse, json, random, time, zlib
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--codecs", type=lambda s: s.split(","), default=["zlib", "lz4"], help="Codecs to measure: zlib and/or lz4.")
p.add_argument("--values", type=int, default=50_000_000, help="Part A: 4-byte values in the shuffle-like block.")
p.add_argument("--distinct", type=int, default=1000, help="Part A: distinct values (fewer compress better).")
p.add_argument("--bandwidths", type=lambda s: [float(v) for v in s.split(",")], default=[1, 10, 100],
               help="Part A: network speeds in Gbps, comma-separated.")
p.add_argument("--rows", type=int, default=1000, help="Part B: rows to partition.")
p.add_argument("--parts", type=int, default=4, help="Part B: partitions.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

# Part A
data = np.random.default_rng(0).integers(0, args.distinct, args.values, dtype=np.int32).tobytes()

import lz4.frame                                         # imported up front: not part of the timing

def compress(codec, b):
    if codec == "zlib": return zlib.compress(b, 1)
    return lz4.frame.compress(b)

codecs, decisions = [], []
for codec in args.codecs:
    t = time.perf_counter(); z = compress(codec, data); dt = time.perf_counter() - t
    C, r = len(data) / dt / 1e6, len(data) / len(z)     # MB/s, ratio
    codecs.append({"codec": codec, "C_MBps": C, "ratio": r, "factor": r / (r - 1)})
    print(f"{codec:>5}: C={C:.0f} MB/s  r={r:.2f}  threshold factor r/(r-1)={r / (r - 1):.2f}")
    for gbps in args.bandwidths:
        bw = gbps * 1000 / 8                             # Gbps -> MB/s
        need = bw * r / (r - 1)
        verdict = "compress wins" if C > need else "send raw"
        decisions.append({"codec": codec, "gbps": gbps, "needed_MBps": need, "verdict": verdict})
        print(f"       {gbps:>5g} Gbps: needs C > {need:,.0f} MB/s -> {verdict}")

# Part B: first attempt and retry of the same map task
def partition(rows, deterministic):
    out = [set() for _ in range(args.parts)]
    for row_id in rows:
        # hash((row_id, random)): CPython's hash(row_id + random.random()) % parts equals row_id % parts,
        # so the random term must go through a tuple hash to actually move rows
        h = hash(row_id) if deterministic else hash((row_id, random.random()))
        out[h % args.parts].add(row_id)
    return out

rows = range(args.rows)
half = args.parts // 2
print()
partitioners = []
for deterministic in (False, True):
    first, retry = partition(rows, deterministic), partition(rows, deterministic)
    moved = sum(len(f ^ r) for f, r in zip(first, retry)) // 2
    # reducers 0..half-1 fetched from the first attempt, the rest from the retry
    seen = [row for p in first[:half] + retry[half:] for row in p]
    lost, duplicated = args.rows - len(set(seen)), len(seen) - len(set(seen))
    name = "hash(row_id)" if deterministic else "hash((row_id, random()))"
    partitioners.append({"partitioner": name, "rows_moved": moved, "lost": lost, "duplicated": duplicated})
    print(f"{name:<24} rows in a different partition on retry: {moved:>4}   "
          f"after mixed fetch: lost {lost:>3}, duplicated {duplicated:>3}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"codecs": codecs, "decisions": decisions, "partitioners": partitioners}}, f)