# /// script
# requires-python = ">=3.10"
# dependencies = ["pyarrow", "numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Min/max skipping on sorted against unsorted data.

10 million rows with a store_id from 1 to 400, written to two Parquet files
with 100,000 rows per row group: one in random order, one sorted by store_id.
Reads the footer statistics and counts the row groups a filter store_id = 17
must read. Then adds one row with store_id = 0 to the sorted file's first
row group and reads that group's range again.
"""
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("--n", type=int, default=10_000_000, help="Rows.")
p.add_argument("--row-group", type=int, default=100_000, help="Rows per row group.")
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()

d = tempfile.mkdtemp(); rng = np.random.default_rng(args.seed)
ids = rng.integers(1, 401, args.n)
def must_read(path):
    md = pq.ParquetFile(path).metadata
    stats = [md.row_group(i).column(0).statistics for i in range(md.num_row_groups)]
    return sum(1 for s in stats if s.min <= 17 <= s.max), md.num_row_groups, (stats[0].min, stats[0].max)
rows = []
for name, col in (("random order", ids), ("sorted by store_id", np.sort(ids))):
    path = os.path.join(d, name.split()[0] + ".parquet")
    pq.write_table(pa.table({"store_id": col}), path, row_group_size=args.row_group)
    keep, total, first = must_read(path)
    rows.append({"file": name, "row_groups": total, "must_read": keep, "skipped": total - keep, "first_group_range": f"{first[0]}-{first[1]}"})
col = np.sort(ids); col[0] = 0                                          # one outlier in the first group
path = os.path.join(d, "outlier.parquet"); pq.write_table(pa.table({"store_id": col}), path, row_group_size=args.row_group)
keep, total, first = must_read(path)
rows.append({"file": "sorted, one row with store_id = 0", "row_groups": total, "must_read": keep, "skipped": total - keep, "first_group_range": f"{first[0]}-{first[1]}"})
for r in rows: print(f"{r['file']:<36} row groups that must be read: {r['must_read']} of {r['row_groups']}   first group's range {r['first_group_range']}")
for f in os.listdir(d): os.remove(os.path.join(d, f))

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