# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Read the statistics inside a Parquet file.

Writes 10 Parquet files with sorted timestamp values and 10 with random ones,
reads each footer with parquet_metadata, and counts the files a range filter
can skip from min/max alone. Then merges two files' statistics by hand (min
of mins, max of maxes, sum of counts) and checks them against the engine.
"""
import argparse, json, os, tempfile
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--files", type=int, default=10, help="Files per layout.")
p.add_argument("--rows", type=int, default=100_000, help="Rows per file.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

work = tempfile.mkdtemp(); c = duckdb.connect(); N = args.files * args.rows
c.execute(f"CREATE TABLE ev AS SELECT TIMESTAMP '2025-01-01' + INTERVAL (i) SECOND AS ts, i AS v FROM range({N}) t(i)")
for layout, order in (("sorted", "ts"), ("random", "random()")):
    c.execute(f"CREATE OR REPLACE TABLE l AS SELECT *, row_number() OVER (ORDER BY {order}) - 1 AS rn FROM ev")
    for f in range(args.files):
        c.execute(f"COPY (SELECT ts, v FROM l WHERE rn // {args.rows} = {f}) TO '{work}/{layout}-{f}.parquet' (FORMAT parquet)")
lo, hi = "2025-01-01 10:00:00", "2025-01-01 11:00:00"                   # a one-hour range
rows = []
for layout in ("sorted", "random"):
    st = c.execute(f"""SELECT file_name, min(stats_min)::TIMESTAMP, max(stats_max)::TIMESTAMP FROM parquet_metadata('{work}/{layout}-*.parquet')
                       WHERE path_in_schema = 'ts' GROUP BY file_name""").fetchall()
    skip = sum(1 for _, mn, mx in st if str(mx) < lo or str(mn) > hi)
    rows.append({"layout": layout, "files": len(st), "skippable": skip, "must_read": len(st) - skip})
md = c.execute(f"""SELECT stats_min_value, stats_max_value, num_values FROM parquet_metadata('{work}/sorted-[01].parquet')
                   WHERE path_in_schema = 'v'""").fetchall()
merged = {"min": min(int(m[0]) for m in md), "max": max(int(m[1]) for m in md), "count": sum(m[2] for m in md)}
truth = c.execute(f"SELECT min(v), max(v), count(v) FROM read_parquet('{work}/sorted-[01].parquet')").fetchone()
summary = {**{f"merged_{k}": v for k, v in merged.items()}, "engine_min": truth[0], "engine_max": truth[1], "engine_count": truth[2],
           "merge_matches": "yes" if (merged["min"], merged["max"], merged["count"]) == tuple(truth) else "no"}
for r in rows: print(f"{r['layout']:<7} {r['files']} files: {r['skippable']} skippable from min/max for a one-hour range, {r['must_read']} must be read")
print(f"merged statistics of two files by hand: min {merged['min']}, max {merged['max']}, count {merged['count']}; engine: {truth}; match: {summary['merge_matches']}")

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