# /// script
# requires-python = ">=3.10"
# dependencies = ["fastavro", "zstandard"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""How big is JSON really, after compression?

The same 1,000 parcel-scan records as JSON lines and as Avro (no field names,
binary numbers), raw and after Zstd level 9; then with sequential parcel ids.
"""
import argparse, io, json, random
import zstandard as zstd
from fastavro import schemaless_writer, parse_schema

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=1000, help="Records.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
schema = parse_schema({"type": "record", "name": "Scan", "fields": [
    {"name": "parcel_id", "type": "long"}, {"name": "scanner_id", "type": "int"},
    {"name": "event_time", "type": "long"}, {"name": "status", "type": "string"},
    {"name": "weight_grams", "type": "int"}, {"name": "depot", "type": "string"}]})
z = lambda b: len(zstd.ZstdCompressor(level=9).compress(b))
rows = []
for label, seq in (("random parcel ids", False), ("sequential parcel ids", True)):
    random.seed(1)
    recs = [{"parcel_id": (10**8 + i) if seq else random.randrange(10**9), "scanner_id": random.randrange(500),
             "event_time": 1_700_000_000_000_000 + i * 1000, "status": random.choice(["LOADED", "SCANNED", "DELIVERED"]),
             "weight_grams": random.randrange(100, 30000), "depot": random.choice(["PUNE-01", "MUMBAI-02", "DELHI-01"])}
            for i in range(args.records)]
    js = "\n".join(json.dumps(r) for r in recs).encode()
    buf = io.BytesIO()
    for r in recs: schemaless_writer(buf, schema, r)
    av = buf.getvalue()
    rows.append({"ids": label, "json_raw": len(js), "avro_raw": len(av), "raw_ratio": len(js) / len(av),
                 "json_zstd": z(js), "avro_zstd": z(av), "zstd_ratio": z(js) / z(av)})
    r = rows[-1]
    print(f"{label:<22} raw: JSON {r['json_raw']:,} B, Avro {r['avro_raw']:,} B ({r['raw_ratio']:.1f}x)   "
          f"zstd: JSON {r['json_zstd']:,} B, Avro {r['avro_zstd']:,} B ({r['zstd_ratio']:.1f}x)")

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