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

DuckDB reads two Parquet files of 10 million orders (random order and sorted
by store_id, 100,000 rows per row group, with a wide text column). Times
SELECT sum(qty) WHERE store_id = 17 on each, with the rows the scan returned
(EXPLAIN ANALYZE); then all columns against one column; then the filter
wrapped in an expression (store_id + 0 = 17), to see whether the scan still
receives it.
"""
import argparse, json, os, re, tempfile, time
import duckdb, 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("--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)
store = rng.integers(1, 401, args.n); qty = rng.integers(1, 10, args.n)
note = pa.array(np.char.add("note-", rng.integers(0, 10**12, args.n).astype(str)))           # the wide text column
for name, order in (("random", np.arange(args.n)), ("sorted", np.argsort(store, kind="stable"))):
    pq.write_table(pa.table({"store_id": store[order], "qty": qty[order], "note": note.take(pa.array(order))}),
                   os.path.join(d, f"{name}.parquet"), row_group_size=100_000)
c = duckdb.connect(); c.execute("SET threads = 4")
def run(sql, reps=5):
    c.execute(sql).to_arrow_table(); ts = []
    for _ in range(reps):
        t = time.perf_counter(); c.execute(sql).to_arrow_table(); ts.append(time.perf_counter() - t)   # materialize every row read
    prof = c.execute("EXPLAIN ANALYZE " + sql).fetchall()[0][1]
    scanned = sum(int(x.replace(",", "")) for x in re.findall(r"(?:TABLE_SCAN|PARQUET_SCAN|READ_PARQUET)[\s\S]*?([\d,]+) [Rr]ows", prof)[:1]) if prof else None
    return min(ts) * 1000, scanned, prof
rows = []
for name in ("random", "sorted"):
    f = os.path.join(d, f"{name}.parquet")
    for label, sql in (("sum(qty) WHERE store_id = 17", f"SELECT sum(qty) FROM '{f}' WHERE store_id = 17"),
                       ("same, filter as store_id + 0 = 17", f"SELECT sum(qty) FROM '{f}' WHERE store_id + 0 = 17"),
                       ("SELECT * (all columns)", f"SELECT * FROM '{f}'"),
                       ("SELECT qty (one column)", f"SELECT qty FROM '{f}'")):
        ms, scanned, prof = run(sql)
        pushed = "yes" if re.search(r"Filters:\s*store_id\s*=\s*17", prof) else "no"
        rows.append({"file": name, "query": label, "ms": ms, "rows_scanned": scanned, "filter_in_scan": pushed if "WHERE" in sql else None})
sizes = {c_: pq.ParquetFile(os.path.join(d, "random.parquet")).metadata.row_group(0).column(i).total_compressed_size for i, c_ in enumerate(["store_id", "qty", "note"])}
summary = {"qty_share_pct": sizes["qty"] / sum(sizes.values()) * 100, "duckdb": duckdb.__version__}
for r in rows: print(f"{r['file']:<7} {r['query']:<36} {r['ms']:8.1f} ms   rows from scan {r['rows_scanned']}" + (f"   filter pushed into scan: {r['filter_in_scan']}" if r["filter_in_scan"] else ""))
print(f"qty is {summary['qty_share_pct']:.1f}% of the file's bytes")
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({"summary": summary, "tables": {"queries": rows}}, f)