# /// 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
"""The composite-key rule: C_a against U.

20 million rows with depot_id (500 values), user_id (5 million values) and
scan_time (730 days), written as 20 files sorted by (depot_id, scan_time) and
again sorted by (user_id, scan_time). From the scan_time statistics, counts
the files a one-day query must open, alone and with a depot given. When the
leading column's cardinality C_a is far above the number of files U, the
second column's ranges cover everything in every file.
"""
import argparse, json, os, shutil, tempfile
import numpy as np, pandas as pd, pyarrow as pa, pyarrow.parquet as pq

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=20_000_000, help="Rows.")
p.add_argument("--files", type=int, default=20, help="Files U.")
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()

rng = np.random.default_rng(args.seed); base = tempfile.mkdtemp(); n = args.rows; per = n // args.files
df = pd.DataFrame({"depot_id": rng.integers(0, 500, n), "user_id": rng.integers(0, 5_000_000, n), "scan_time": rng.integers(0, 730, n)})
def stats(frame, folder):
    os.makedirs(folder, exist_ok=True); out = []
    for i in range(args.files):
        part = frame.iloc[i * per:(i + 1) * per]; path = f"{folder}/f{i}.parquet"
        pq.write_table(pa.Table.from_pandas(part, preserve_index=False), path, row_group_size=per)
        md = pq.ParquetFile(path).metadata.row_group(0)
        col = {md.column(j).path_in_schema: md.column(j).statistics for j in range(md.num_columns)}
        out.append(col)
    return out
day, depot = 365, 17
rows = []
for label, key, ca in (("sorted by (depot_id, scan_time)", ["depot_id", "scan_time"], 500), ("sorted by (user_id, scan_time)", ["user_id", "scan_time"], 5_000_000)):
    st = stats(df.sort_values(key), f"{base}/{key[0]}")
    day_only = sum(1 for s in st if s["scan_time"].min <= day <= s["scan_time"].max)
    with_depot = sum(1 for s in st if s["scan_time"].min <= day <= s["scan_time"].max and s["depot_id"].min <= depot <= s["depot_id"].max)
    rows.append({"layout": label, "C_a": ca, "U": args.files, "C_a_over_U": ca / args.files, "files_day_only": day_only, "files_day_and_depot": with_depot})
shutil.rmtree(base, ignore_errors=True)
for r in rows: print(f"{r['layout']:<32} C_a/U = {r['C_a_over_U']:>9,.0f}: one day opens {r['files_day_only']:>2} of {r['U']} files; one day + one depot opens {r['files_day_and_depot']}")

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