# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Many small reads against few large reads: the transfer floor.

Random reads of one chunk size at a time from a large file, with the page
cache bypassed (otherwise you measure memory, not the disk); then the shuffle
width ceiling that follows from the floor.
"""
import argparse, json, math, mmap, os, platform, random, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--path", default=os.path.join(tempfile.gettempdir(), "transfer-floor.bin"), help="Test file (removed at exit).")
p.add_argument("--size-mb", type=int, default=1024, help="Size of the test file, in MB.")
p.add_argument("--chunks", type=lambda s: [int(v) for v in s.split(",")],
               default=[4_096, 16_384, 65_536, 262_144, 1 << 20, 4 << 20, 16 << 20],
               help="Read sizes in bytes, comma-separated (multiples of 4096).")
p.add_argument("--reads", type=int, default=2000, help="Random reads per chunk size (capped at file size / chunk).")
p.add_argument("--cache", choices=["bypass", "page"], default="bypass",
               help="bypass: read around the page cache (measures the disk); page: normal reads (measures memory once cached).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
import atexit, shutil
atexit.register(lambda: shutil.rmtree(args.path, ignore_errors=True) if os.path.isdir(args.path) else os.path.exists(args.path) and os.remove(args.path))   # leave nothing behind
SIZE = args.size_mb << 20
DARWIN, LINUX = platform.system() == "Darwin", platform.system() == "Linux"
F_NOCACHE = 48                                           # macOS fcntl: do not keep this file's pages in the cache


def no_cache(fd):
    if DARWIN:
        import fcntl; fcntl.fcntl(fd, F_NOCACHE, 1)
    elif LINUX:
        os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)


if not os.path.exists(args.path) or os.path.getsize(args.path) != SIZE:
    fd = os.open(args.path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC)
    if args.cache == "bypass": no_cache(fd)
    for _ in range(args.size_mb): os.write(fd, os.urandom(1 << 20))
    os.fsync(fd); os.close(fd)

flags = os.O_RDONLY | (os.O_DIRECT if LINUX and args.cache == "bypass" else 0)
fd = os.open(args.path, flags)
if args.cache == "bypass": no_cache(fd)

rows = []
for chunk in args.chunks:
    buf = mmap.mmap(-1, chunk)                           # page-aligned buffer, as O_DIRECT requires
    n = min(args.reads, SIZE // chunk); total = 0
    t = time.perf_counter()
    for _ in range(n):
        off = random.randrange(0, SIZE - chunk) // 4096 * 4096
        total += os.preadv(fd, [buf], off)
    dt = time.perf_counter() - t
    rows.append({"chunk_bytes": chunk, "MBps": total / dt / 1e6})
    print(f"{chunk:>12,} B reads: {total / dt / 1e6:8.1f} MB/s")
os.close(fd)

# b_min: the smallest chunk that reaches about half of the peak rate
peak = max(r["MBps"] for r in rows)
b_min = next(r["chunk_bytes"] for r in rows if r["MBps"] >= peak / 2)
print(f"\npeak {peak:,.0f} MB/s · b_min (half of peak) = {b_min:,} B")
ceilings = []
for label, S in (("100 GB", 100e9), ("1 TB", 1e12), ("10 TB", 10e12)):
    m = math.sqrt(S / b_min)
    ceilings.append({"data": label, "max_width": m})
    print(f"S = {label:>6}: M = R <= sqrt(S / b_min) = {m:,.0f}")
print(f"10x the data -> {ceilings[1]['max_width'] / ceilings[0]['max_width']:.1f}x the width")
print(f"b = S / (M x R) for 1 TB at 10,000 x 10,000: {1e12 / (10_000 * 10_000) / 1e3:.0f} KB")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"reads": rows, "ceilings": ceilings, "floor": [{"peak_MBps": peak, "b_min_bytes": b_min}]}}, f)