# /// script
# requires-python = ">=3.10"
# dependencies = ["zstandard"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Batch size is a compression parameter.

100,000 small JSON events of about 200 bytes with realistic repetition,
compressed with Zstd in batches of 1 to 10,000 events; then single events
compressed with a trained dictionary.
"""
import argparse, json, random
import zstandard as zstd

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--events", type=int, default=100_000, help="Events.")
p.add_argument("--level", type=int, default=3, help="Zstd level.")
p.add_argument("--sizes", type=lambda s: [int(v) for v in s.split(",")], default=[1, 10, 100, 1000, 10000], help="Batch sizes.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(1)
events = [json.dumps({"order_id": 1_000_000 + i, "seller_id": random.randint(1, 40),
                      "sku": f"SKU-{random.randint(1, 60)}", "qty": random.randint(1, 5),
                      "status": "OrderPlaced", "country": random.choice(["IN", "UK", "US"])}).encode()
          for i in range(args.events)]
raw = sum(len(e) for e in events)
c = zstd.ZstdCompressor(level=args.level)
rows = []
for batch in args.sizes:
    out = sum(len(c.compress(b"\n".join(events[i:i + batch]))) for i in range(0, len(events), batch))
    rows.append({"batch": batch, "ratio": raw / out, "mode": "plain"}); print(f"batch={batch:6d}  ratio={raw / out:5.2f}x")
d = zstd.train_dictionary(16384, events[:5000])
cd = zstd.ZstdCompressor(level=args.level, dict_data=d)
out = sum(len(cd.compress(e)) for e in events[5000:]) * len(events) / (len(events) - 5000)
rows.append({"batch": 1, "ratio": raw / out, "mode": "dictionary"})
print(f"batch=     1  with a 16 KB trained dictionary  ratio={raw / out:5.2f}x   (average event {raw / len(events):.0f} bytes)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"ratios": [r for r in rows if r["mode"] == "plain"], "summary": [{
            "single": rows[0]["ratio"], "k1000": next(r["ratio"] for r in rows if r["batch"] == 1000 and r["mode"] == "plain"),
            "dict": rows[-1]["ratio"], "avg_bytes": raw / len(events)}]}}, f)