# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Row order against the codec.

The same 2 million rows (depot, status, event_time) written to Parquet
shuffled and sorted by depot then time, with no codec, Snappy and Zstd; then
the depot column's min/max per row group in each file.
"""
import argparse, json, os, tempfile
import numpy as np, pyarrow as pa, pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=2_000_000, help="Rows.")
p.add_argument("--row-group", type=int, default=128_000, help="Rows per row group.")
p.add_argument("--seed", type=int, default=1, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed); N = args.rows
tbl = pa.table({"depot": rng.integers(0, 200, N), "status": rng.integers(0, 10, N), "event_time": np.cumsum(rng.integers(1, 5_000, N))})
d = tempfile.mkdtemp()
def size(table, name, codec):
    path = os.path.join(d, name); pq.write_table(table, path, compression=codec, use_dictionary=True, row_group_size=args.row_group)
    return path, os.path.getsize(path) / 1e6
shuffled = tbl.take(rng.permutation(N)); sorted_ = tbl.sort_by([("depot", "ascending"), ("event_time", "ascending")])
rows = []
for codec in ("NONE", "SNAPPY", "ZSTD"):
    _, a = size(shuffled, f"a-{codec}.parquet", codec); _, b = size(sorted_, f"b-{codec}.parquet", codec)
    rows.append({"codec": codec, "shuffled_mb": a, "sorted_mb": b, "ratio": a / b})
    print(f"{codec:<7} shuffled {a:6.1f} MB   sorted {b:6.1f} MB   ({a / b:.1f}x)")
widths = {}
for label, name in (("shuffled", "a-ZSTD.parquet"), ("sorted", "b-ZSTD.parquet")):
    md = pq.read_metadata(os.path.join(d, name))
    w = [md.row_group(i).column(0).statistics for i in range(md.num_row_groups)]
    widths[label] = sum(s.max - s.min for s in w) / len(w)
    print(f"{label:<9} depot min/max per row group: average width {widths[label]:.0f} of 0-199  (first: {w[0].min}-{w[0].max})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"codecs": rows, "summary": [{"shuffled_width": widths["shuffled"], "sorted_width": widths["sorted"]}]}}, f)