# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The footer pays for column skipping.

A Parquet file of 200 columns and 1,000,000 rows: reading 3 columns against
all 200 (time and bytes); the footer's per-row-group statistics; and an
attempt to append to the file.
"""
import argparse, json, os, tempfile, time
import numpy as np, pyarrow as pa, pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=1_000_000, help="Rows.")
p.add_argument("--cols", type=int, default=200, help="Columns.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
path = os.path.join(tempfile.mkdtemp(), "wide.parquet")
table = pa.table({f"c{i}": np.random.default_rng(i).integers(0, 1000, args.rows) for i in range(args.cols)})
pq.write_table(table, path, row_group_size=100_000)
md = pq.ParquetFile(path).metadata
def chunk_bytes(cols): return sum(md.row_group(g).column(c).total_compressed_size for g in range(md.num_row_groups) for c in cols)
pq.read_table(path, columns=["c1"])                                     # warm the file into the cache for both reads
t = time.perf_counter(); pq.read_table(path, columns=["c1", "c2", "c3"]); t3 = time.perf_counter() - t
t = time.perf_counter(); pq.read_table(path); tall = time.perf_counter() - t
st = md.row_group(0).column(0).statistics
try:
    with open(path, "ab") as f: f.write(b"more rows")
    pq.ParquetFile(path); appended = "the file still reads"
except Exception as e:
    appended = f"the file no longer reads: {type(e).__name__}"
rows = [{"read": "3 columns", "seconds": t3, "mb": chunk_bytes([1, 2, 3]) / 1e6}, {"read": f"all {args.cols} columns", "seconds": tall, "mb": chunk_bytes(range(args.cols)) / 1e6}]
for r in rows: print(f"{r['read']:<12} {r['seconds']:6.3f} s   {r['mb']:7.1f} MB of column chunks")
print(f"footer statistics, row group 0, c0: min {st.min}, max {st.max}, nulls {st.null_count}")
print(f"appending bytes after the footer: {appended}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"reads": rows, "summary": [{"time_ratio": tall / t3, "byte_ratio": rows[1]["mb"] / rows[0]["mb"], "append": appended,
                   "stat_min": st.min, "stat_max": st.max}]}}, f)