# /// script
# requires-python = ">=3.10"
# dependencies = ["zstandard"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find your own break-even point for compression.

Compression wins on a read only if D / B > r / (r - 1): decode speed D over
storage bandwidth B must beat a factor set by the ratio r. Zstd levels 1 to
19 on a JSON log, against the measured local disk and a slow object-store link.
"""
import argparse, json, os, platform, random, tempfile, time
import zstandard as zstd

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--mb", type=int, default=300, help="Size of the generated JSON log, MB.")
p.add_argument("--levels", type=lambda s: [int(v) for v in s.split(",")], default=[1, 3, 10, 19], help="Zstd levels.")
p.add_argument("--remote-gbps", type=float, default=0.08, help="An object-store connection, GB/s.")
p.add_argument("--path", default=os.path.join(tempfile.gettempdir(), "breakeven.jsonl"), help="The sample file.")
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
random.seed(0)
if not os.path.exists(args.path) or os.path.getsize(args.path) < args.mb * 1e6 * 0.99:
    with open(args.path, "w") as f:
        i = 0
        while f.tell() < args.mb * 1e6:
            f.write(json.dumps({"ts": 1_700_000_000 + i, "user": f"u{random.randint(1, 50000)}", "event": random.choice(["view", "click", "buy"]),
                                "path": f"/p/{random.randint(1, 3000)}", "ms": random.randint(1, 900)}) + "\n"); i += 1
data = open(args.path, "rb").read()

# B: read the file with the page cache bypassed
fd = os.open(args.path, os.O_RDONLY)
if platform.system() == "Darwin":
    import fcntl; fcntl.fcntl(fd, 48, 1)                      # F_NOCACHE
else:
    os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
t = time.perf_counter(); n = 0
while (chunk := os.read(fd, 8 << 20)): n += len(chunk)
disk = n / (time.perf_counter() - t) / 1e9; os.close(fd)
print(f"local disk read, page cache bypassed: B = {disk:.2f} GB/s; object store: B = {args.remote_gbps} GB/s")

rows = []
for level in args.levels:
    t0 = time.time(); comp = zstd.ZstdCompressor(level=level).compress(data); tc = time.time() - t0
    t0 = time.time(); zstd.ZstdDecompressor().decompress(comp); td = time.time() - t0
    r = len(data) / len(comp); D = len(data) / td / 1e9; need = r / (r - 1)
    rows.append({"level": level, "ratio": r, "write_s": tc, "decode_gbps": D, "need": need,
                 "local": "wins" if D / disk > need else "loses", "remote": "wins" if D / args.remote_gbps > need else "loses",
                 "d_over_b_local": D / disk})
    print(f"level {level:2d}  ratio {r:4.1f}  write {tc:6.2f}s  decode {D:5.2f} GB/s  D/B local {D / disk:5.2f} "
          f"need > {need:4.2f} -> local {rows[-1]['local']}, object store {rows[-1]['remote']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"levels": rows, "summary": [{"disk_gbps": disk}]}}, f)