# /// script
# requires-python = ">=3.10"
# dependencies = ["zstandard"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Block size: compression ratio against point-lookup cost.

200,000 JSON rows compressed with Zstd in blocks of 1 KB to 1 MB: the ratio,
and the time to decode the one block a point lookup needs.
"""
import argparse, json, random, time
import zstandard as zstd

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=200_000, help="Rows.")
p.add_argument("--blocks", type=lambda s: [int(v) for v in s.split(",")], default=[1024, 4096, 16384, 65536, 262144, 1048576], help="Block sizes, bytes.")
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)
rows_ = [json.dumps({"parcel": f"P{random.randint(0, 10**6)}", "depot": random.choice(["RTM", "LYN", "HAM"]),
                     "status": random.choice(["SCANNED", "LOADED", "DELIVERED"]), "t": 1700000000 + i}).encode() for i in range(args.rows)]
data = b"\n".join(rows_)
c, d = zstd.ZstdCompressor(level=3), zstd.ZstdDecompressor()
rows = []
for block in args.blocks:
    comp = [c.compress(data[i:i + block]) for i in range(0, len(data), block)]
    ratio = len(data) / sum(len(x) for x in comp)
    t0 = time.perf_counter()
    for _ in range(200): d.decompress(comp[len(comp) // 2])
    us = (time.perf_counter() - t0) / 200 * 1e6
    rows.append({"block": block, "ratio": ratio, "lookup_us": us})
    print(f"block {block:>8} B  ratio {ratio:5.2f}  one lookup decodes {block} B in {us:7.1f} us")

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