# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The spill cliff: 5% spilled, how much slower?

Measures memory bandwidth (a sum over an int64 array) and the disk's
write-then-read bandwidth, giving r = memory / disk. Then spills a fraction f
of the array to a file (written and read back) while the rest is summed in
memory, and compares the slowdown with time ∝ (1 − f) + 2·f·r. Bypasses the
OS page cache for the spill file (F_NOCACHE on macOS, fadvise DONTNEED on
Linux) so the read-back comes from the device.
"""
import argparse, fcntl, json, os, sys, tempfile, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=100_000_000, help="int64 values (800 MB at the default).")
p.add_argument("--fractions", default="0,0.05,0.2", help="Spilled fractions f.")
p.add_argument("--r-network", type=float, default=500, help="A second r to compare with (a network disk).")
p.add_argument("--dir", default=tempfile.gettempdir(), help="Directory for the spill file.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

a = np.arange(args.n, dtype=np.int64)
a.sum(); t = time.perf_counter(); a.sum(); t_mem = time.perf_counter() - t
bw_mem = a.nbytes / t_mem / 1e9
path = os.path.join(args.dir, "spill.bin")

def nocache(fd):
    if sys.platform == "darwin": fcntl.fcntl(fd, fcntl.F_NOCACHE, 1)

def spill(k):
    """Write a[:k] and read it back, bypassing the page cache. Returns (write s, read s, data)."""
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC); nocache(fd)
    t = time.perf_counter(); out = memoryview(a[:k]).cast("B"); sent = 0
    while sent < len(out): sent += os.write(fd, out[sent:sent + (64 << 20)])
    os.fsync(fd); w = time.perf_counter() - t
    if hasattr(os, "posix_fadvise"): os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
    os.close(fd)
    b = np.empty(k, dtype=np.int64); fd = os.open(path, os.O_RDONLY); nocache(fd)
    t = time.perf_counter(); view = memoryview(b).cast("B"); got = 0
    while got < len(view): got += os.readv(fd, [view[got:]])
    r = time.perf_counter() - t; os.close(fd)
    return w, r, b

k_cal = args.n // 5; w, rd, _ = spill(k_cal)
bw_disk = 2 * k_cal * 8 / (w + rd) / 1e9                               # bytes moved once out and once in
r_ratio = bw_mem / bw_disk
rows = []
for f in map(float, args.fractions.split(",")):
    k = int(args.n * f)
    t = time.perf_counter(); w, rd, b = spill(k); total = a[k:].sum() + b.sum(); el = time.perf_counter() - t
    assert total == a.sum()
    rows.append({"f": f, "measured_x": el / t_mem, "formula_x": (1 - f) + 2 * f * r_ratio, "formula_network_x": (1 - f) + 2 * f * args.r_network,
                 "write_s": w, "read_s": rd})
os.remove(path)
summary = {"mem_gb_s": bw_mem, "disk_gb_s": bw_disk, "r": r_ratio, "r_network": args.r_network}
print(f"memory {bw_mem:.1f} GB/s, disk (write + read) {bw_disk:.2f} GB/s, r = {r_ratio:.0f}")
for r in rows: print(f"f = {r['f']:<5g} measured {r['measured_x']:6.1f}x   formula {r['formula_x']:6.1f}x   with r = {args.r_network:g}: {r['formula_network_x']:6.1f}x")

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