# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Planning time against partition count.

An SQLite catalog table partitions(table_id, day, file_count) holds 1,000 to
1,000,000 partitions of one table. Times the unfiltered read of all of them
(the work of planning a full scan) and a read with a date filter served by an
index, then applies Little's Law: slots L = target plans/s × planning time W.
"""
import argparse, json, sqlite3, statistics, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--counts", default="1000,10000,100000,1000000", help="Partition counts.")
p.add_argument("--target", type=float, default=20, help="Target plans per second.")
p.add_argument("--repeats", type=int, default=5, help="Timings per query (median).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def timed(c, sql, params=()):
    ts = []
    for _ in range(args.repeats):
        t = time.perf_counter(); c.execute(sql, params).fetchall(); ts.append(time.perf_counter() - t)
    return statistics.median(ts)

rows = []
for n in map(int, args.counts.split(",")):
    c = sqlite3.connect(":memory:")
    c.execute("CREATE TABLE partitions (table_id int, day int, file_count int)")
    c.executemany("INSERT INTO partitions VALUES (1, ?, ?)", ((d, d % 17 + 1) for d in range(n)))
    c.execute("CREATE INDEX ON_day ON partitions (table_id, day)")
    full = timed(c, "SELECT * FROM partitions WHERE table_id = 1")
    filt = timed(c, "SELECT * FROM partitions WHERE table_id = 1 AND day BETWEEN ? AND ?", (n // 2, n // 2 + 30))
    rows.append({"partitions": n, "full_ms": full * 1000, "filtered_ms": filt * 1000, "slots": args.target * full})
for r in rows: print(f"{r['partitions']:>9,} partitions: full-scan planning {r['full_ms']:9.3f} ms   filtered {r['filtered_ms']:7.3f} ms   slots for {args.target:g} plans/s: {r['slots']:7.2f}")

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