# /// script
# requires-python = ">=3.10"
# dependencies = ["boto3", "psutil"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The one ratio: bytes touched against bytes needed.

Local part: a 1 GB file of random data, 20,000 records of 200 bytes read at
random offsets, as plain 200-byte reads and as 4 KiB aligned reads, bypassing
the page cache (F_NOCACHE on macOS, posix_fadvise elsewhere); the bytes the
OS reports read from the disk against the bytes asked for. Object-store part:
1,000 range GETs of 4 KiB against 10 range GETs of 4 MB from SeaweedFS.
Needs the S3 service from the compose file (port 59200) for the second part.
"""
import argparse, json, os, random, sys, tempfile, time
import boto3, psutil
from botocore.config import Config

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--file-mb", type=int, default=1024, help="Local test file (MB).")
p.add_argument("--reads", type=int, default=20_000, help="Random record reads.")
p.add_argument("--endpoint", default="http://localhost:59200", help="S3 endpoint.")
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()
random.seed(args.seed)

path = os.path.join(tempfile.mkdtemp(), "ratio.bin")
with open(path, "wb") as f:
    for _ in range(args.file_mb): f.write(os.urandom(1 << 20))
size = os.path.getsize(path)
def nocache(fd):
    if sys.platform == "darwin":
        import fcntl; fcntl.fcntl(fd, fcntl.F_NOCACHE, 1)
    elif hasattr(os, "posix_fadvise"): os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
offsets = [random.randrange(0, size - 4096) for _ in range(args.reads)]
local = []
for label, aligned in (("200-byte reads", False), ("4 KiB aligned reads", True)):
    fd = os.open(path, os.O_RDONLY); nocache(fd)
    before = psutil.disk_io_counters().read_bytes; t = time.perf_counter(); got = 0
    for off in offsets:
        if aligned: got += len(os.pread(fd, 4096, off - off % 4096))
        else: got += len(os.pread(fd, 200, off))
    el = time.perf_counter() - t; disk = psutil.disk_io_counters().read_bytes - before; os.close(fd)
    local.append({"reads": label, "asked_mb": args.reads * 200 / 2**20, "returned_mb": got / 2**20, "disk_mb": disk / 2**20,
                  "disk_over_asked": disk / (args.reads * 200), "seconds": el})
os.remove(path)
s3 = boto3.client("s3", endpoint_url=args.endpoint, aws_access_key_id="lab", aws_secret_access_key="lablablab",
                  config=Config(s3={"addressing_style": "path"}), region_name="us-east-1")
try: s3.create_bucket(Bucket="ratio")
except Exception: pass
s3.put_object(Bucket="ratio", Key="obj", Body=os.urandom(64 << 20))
remote = []
for label, n, sz in (("1,000 GETs of 4 KiB", 1000, 4096), ("10 GETs of 4 MB", 10, 4 << 20)):
    t = time.perf_counter()
    for i in range(n):
        off = random.randrange(0, (64 << 20) - sz); s3.get_object(Bucket="ratio", Key="obj", Range=f"bytes={off}-{off + sz - 1}")["Body"].read()
    el = time.perf_counter() - t
    remote.append({"gets": label, "mb": n * sz / 2**20, "seconds": el, "ms_per_get": el / n * 1000, "mb_per_s": n * sz / 2**20 / el})
for r in local: print(f"{r['reads']:<20} asked {r['asked_mb']:6.1f} MB, OS read from disk {r['disk_mb']:7.1f} MB ({r['disk_over_asked']:.1f}x), {r['seconds']:.2f} s")
for r in remote: print(f"{r['gets']:<20} {r['mb']:6.1f} MB in {r['seconds']:6.2f} s ({r['ms_per_get']:.1f} ms per GET, {r['mb_per_s']:.0f} MB/s)")

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