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

A dry run estimates the bytes a query will scan from metadata alone: the
column chunks of the referenced columns, in the partitions and row groups the
filters cannot rule out. This script builds a date-partitioned Parquet table
and computes that estimate for the same question asked four ways.
"""
import argparse, json, os, tempfile
import numpy as np
import pyarrow as pa, pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--days", type=int, default=30, help="Daily partitions.")
p.add_argument("--rows-per-day", type=int, default=200_000, help="Rows per partition.")
p.add_argument("--columns", type=int, default=20, help="Columns in the table.")
p.add_argument("--row-group", type=int, default=20_000, help="Rows per Parquet row group.")
p.add_argument("--path", default=os.path.join(tempfile.gettempdir(), "dry-run-lab"), help="Where the table is written.")
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()
import atexit, shutil
atexit.register(lambda: shutil.rmtree(args.path, ignore_errors=True) if os.path.isdir(args.path) else os.path.exists(args.path) and os.remove(args.path))   # leave nothing behind
rng = np.random.default_rng(args.seed)

os.makedirs(args.path, exist_ok=True)
names = ["order_id", "amount", "customer_score"] + [f"col{i}" for i in range(3, args.columns)]
for d in range(args.days):
    n = args.rows_per_day
    cols = {"order_id": np.arange(d * n, (d + 1) * n), "amount": rng.gamma(2, 50, n),
            "customer_score": rng.uniform(0, 100, n)}                      # uncorrelated with everything: no clustering
    for c in names[3:]: cols[c] = rng.normal(0, 1, n)
    pq.write_table(pa.table(cols), f"{args.path}/day={d:02d}.parquet", row_group_size=args.row_group)

def estimate(columns, days=None, score_range=None):
    """Bytes a dry run would charge: referenced columns, in files and row groups not ruled out."""
    total = 0
    for d in range(args.days):
        if days is not None and d not in days: continue                   # partition pruning
        md = pq.ParquetFile(f"{args.path}/day={d:02d}.parquet").metadata
        for rg in range(md.num_row_groups):
            g = md.row_group(rg)
            cols = {g.column(i).path_in_schema: g.column(i) for i in range(g.num_columns)}
            if score_range:                                                # min/max statistics
                st = cols["customer_score"].statistics
                if st.max < score_range[0] or st.min > score_range[1]: continue
            total += sum(cols[c].total_compressed_size for c in columns)
    return total

queries = [
    ("SELECT * over the full table", estimate(names)),
    ("named columns (2)", estimate(["order_id", "amount"])),
    ("named columns + partition filter (3 days)", estimate(["order_id", "amount"], days={0, 1, 2})),
    ("named columns + filter on an uncorrelated column", estimate(["order_id", "amount", "customer_score"], score_range=(90, 100))),
]
rows = [{"query": q, "estimated_mb": b / 1e6, "vs_select_star": queries[0][1] / b} for q, b in queries]
for r in rows: print(f"{r['query']:<50} {r['estimated_mb']:9.1f} MB   ({r['vs_select_star']:5.1f}x less than SELECT *)")
print(f"(total columns / columns used = {args.columns} / 2 = {args.columns / 2:g})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"queries": rows, "summary": [{"cols_ratio": args.columns / 2, "named_x": rows[1]["vs_select_star"],
                   "partition_x": rows[2]["vs_select_star"], "uncorrelated_x": rows[3]["vs_select_star"]}]}}, f)