# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find the index-versus-scan crossover on your hardware.

A 5-million-row PostgreSQL table with a random value column and an index on
it. For predicates returning 0.01% to 20% of the rows, the time of an index
scan (sequential scan disabled) against a sequential scan (index disabled);
with random_page_cost 4 and 1.1; then after CLUSTER on the index.
Needs the PostgreSQL service from the compose file (port 55520).
"""
import argparse, json, re
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55520/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=5_000_000, help="Rows.")
p.add_argument("--fractions", type=lambda s: [float(v) for v in s.split(",")], default=[0.0001, 0.001, 0.01, 0.05, 0.2], help="Shares of rows returned.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
con = psycopg.connect(args.dsn, autocommit=True)
con.execute("DROP TABLE IF EXISTS t")
con.execute(f"CREATE TABLE t AS SELECT g AS id, (random()*1000000)::int AS v, md5(g::text) AS pad FROM generate_series(1, {args.rows}) g")
con.execute("CREATE INDEX t_v ON t (v)"); con.execute("ANALYZE t")

def timed(sql, scan):
    con.execute("SET enable_seqscan = %s" % ("on" if scan == "seq" else "off"))
    con.execute("SET enable_indexscan = %s" % ("on" if scan == "index" else "off"))
    con.execute("SET enable_bitmapscan = %s" % ("on" if scan == "index" else "off"))
    plan = "\n".join(r[0] for r in con.execute("EXPLAIN (ANALYZE) " + sql))
    return float(re.search(r"Execution Time: ([\d.]+)", plan).group(1)), plan.splitlines()[0].split("  ")[0].strip()

rows, crossings = [], []
for layout in ("scattered", "clustered on the index"):
    if layout.startswith("clustered"): con.execute("CLUSTER t USING t_v"); con.execute("ANALYZE t")
    for rpc in (4, 1.1):
        con.execute(f"SET random_page_cost = {rpc}")
        for f in args.fractions:
            sql = f"SELECT * FROM t WHERE v < {int(f * 1_000_000)}"
            ti, _ = timed(sql, "index"); ts, _ = timed(sql, "seq")
            con.execute("RESET enable_seqscan; RESET enable_indexscan; RESET enable_bitmapscan")
            chosen = "\n".join(r[0] for r in con.execute("EXPLAIN " + sql)).splitlines()[0].split("(")[0].strip()
            rows.append({"layout": layout, "random_page_cost": rpc, "pct": f * 100, "index_ms": ti, "seq_ms": ts, "planner_chose": chosen})
            print(f"{layout:<24} rpc {rpc:<4} {f:7.2%}: index {ti:9.1f} ms   seq scan {ts:9.1f} ms   planner: {chosen}")
        cross = next((r["pct"] for r in rows if r["layout"] == layout and r["random_page_cost"] == rpc and r["index_ms"] > r["seq_ms"]), None)
        crossings.append({"layout": layout, "random_page_cost": rpc, "index_stops_winning_pct": cross if cross is not None else -1})

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