# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy", "pandas", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Row against column on the same table.

10 million rows and 25 columns written as CSV (a stand-in for a row format)
and as Parquet. File sizes, the time to read 2 of the 25 columns, and the time
to read one row by position.
"""
import argparse, json, os, shutil, tempfile, time
import numpy as np, pandas as pd, pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=10_000_000, help="Rows.")
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(0); n = args.rows
df = pd.DataFrame({f"c{i}": rng.integers(0, 200, n) for i in range(24)}); df["weight_g"] = rng.integers(100, 30000, n)
csv, pqf = os.path.join(d, "t.csv"), os.path.join(d, "t.parquet")
df.to_csv(csv, index=False); df.to_parquet(pqf, row_group_size=640_000); del df
row = n // 2
t = time.perf_counter(); pd.read_parquet(pqf, columns=["c0", "weight_g"]).mean(); pq_scan = time.perf_counter() - t
t = time.perf_counter(); pd.read_csv(csv, usecols=["c0", "weight_g"]).mean(); csv_scan = time.perf_counter() - t
t = time.perf_counter(); f = pq.ParquetFile(pqf); rg = row // 640_000; f.read_row_group(rg).slice(row % 640_000, 1).to_pylist(); pq_point = time.perf_counter() - t
offsets = [0]                                                             # a row store knows where each row starts: build the offsets once
with open(csv, "rb") as fh:
    for line in fh: offsets.append(offsets[-1] + len(line))
t = time.perf_counter()
with open(csv, "rb") as fh: fh.seek(offsets[row + 1]); fh.readline()
csv_point = time.perf_counter() - t
rows = [{"format": "CSV (rows)", "size_mb": os.path.getsize(csv) / 2**20, "two_columns_s": csv_scan, "one_row_ms": csv_point * 1000},
        {"format": "Parquet (columns)", "size_mb": os.path.getsize(pqf) / 2**20, "two_columns_s": pq_scan, "one_row_ms": pq_point * 1000}]
shutil.rmtree(d, ignore_errors=True)
for r in rows: print(f"{r['format']:<18} {r['size_mb']:7.0f} MB   2 of 25 columns {r['two_columns_s']:6.2f} s   one row by position {r['one_row_ms']:8.3f} ms")

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