# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Zone maps and clustering.

50 million timestamped rows written to Parquet in random order and sorted by
time; a one-hour range query on each: how many row groups do their min/max
statistics let a reader skip?
"""
import argparse, datetime, json, os, tempfile
import duckdb
import pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=50_000_000, help="Rows.")
p.add_argument("--row-group", type=int, default=1_000_000, help="Rows per row group.")
p.add_argument("--dir", default=os.path.join(tempfile.gettempdir(), "zone-maps-indexing"), help="Where the files go.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
import atexit, shutil
atexit.register(lambda: shutil.rmtree(args.dir, ignore_errors=True) if os.path.isdir(args.dir) else os.path.exists(args.dir) and os.remove(args.dir))   # leave nothing behind
os.makedirs(args.dir, exist_ok=True)
con = duckdb.connect()
con.execute(f"CREATE TABLE t AS SELECT TIMESTAMP '2025-01-01' + INTERVAL (i * 2) SECOND AS ts, i AS id FROM range({args.rows}) r(i)")
paths = {"random order": f"{args.dir}/random.parquet", "sorted by time": f"{args.dir}/sorted.parquet"}
con.execute(f"COPY (SELECT * FROM t ORDER BY random()) TO '{paths['random order']}' (ROW_GROUP_SIZE {args.row_group})")
con.execute(f"COPY (SELECT * FROM t ORDER BY ts) TO '{paths['sorted by time']}' (ROW_GROUP_SIZE {args.row_group})")
lo, hi = con.execute("SELECT min(ts), max(ts) FROM t").fetchone()
q0 = lo + (hi - lo) / 2; q1 = q0 + datetime.timedelta(hours=1)
rows = []
for label, path in paths.items():
    md = pq.ParquetFile(path).metadata
    st = [md.row_group(i).column(0).statistics for i in range(md.num_row_groups)]
    read = sum(1 for s in st if not (s.max < q0 or s.min > q1))
    rows.append({"file": label, "row_groups": md.num_row_groups, "read": read})
    print(f"{label:<16} one-hour query reads {read} of {md.num_row_groups} row groups")

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