# /// 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.

A PostgreSQL table of 5 million rows with an indexed column k whose values
select 0.01% to 50% of the rows. For each selectivity, forces an index scan
and then a sequential scan and times both (EXPLAIN ANALYZE), with the data in
physical order of k (clustered) and shuffled. Also records which plan the
planner picks on its own, at the default random_page_cost (4) and at 1.1.
Needs the PostgreSQL service from the compose file (port 55703).
"""
import argparse, json, re
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55703/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=5_000_000, help="Rows.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

SELECTIVITY = [0.0001, 0.001, 0.01, 0.05, 0.2, 0.5]
c = psycopg.connect(args.dsn, autocommit=True)
def ms(sql):
    plan = "\n".join(r[0] for r in c.execute("EXPLAIN (ANALYZE, BUFFERS) " + sql).fetchall())
    return float(re.search(r"Execution Time: ([\d.]+)", plan).group(1)), plan
def k_where(sel):                                                       # k is a percentile bucket: k < sel * 10,000 selects sel of the rows
    return f"k < {int(sel * 10000)}"
rows, choices = [], []
for layout, order in (("clustered", "g"), ("shuffled", "random()")):
    c.execute("DROP TABLE IF EXISTS t")
    c.execute(f"CREATE TABLE t AS SELECT (g * 10000::bigint / {args.rows})::int AS k, md5(g::text) AS pad FROM generate_series(0, {args.rows - 1}) g ORDER BY {order}")
    c.execute("CREATE INDEX ON t (k)"); c.execute("VACUUM ANALYZE t")
    for sel in SELECTIVITY:
        q = f"SELECT count(pad) FROM t WHERE {k_where(sel)}"
        for setting in ("index", "seq"):
            c.execute("RESET ALL")
            if setting == "index": c.execute("SET enable_seqscan = off"); c.execute("SET enable_bitmapscan = off")
            else: c.execute("SET enable_indexscan = off"); c.execute("SET enable_bitmapscan = off")
            ms(q); t, _ = ms(q)                                          # warm, then measure
            rows.append({"layout": layout, "selectivity_pct": sel * 100, "method": "index scan" if setting == "index" else "sequential scan", "ms": t})
        for rpc in (4, 1.1):
            c.execute("RESET ALL"); c.execute(f"SET random_page_cost = {rpc}")
            plan = "\n".join(r[0] for r in c.execute("EXPLAIN " + q).fetchall())
            pick = "bitmap" if "Bitmap" in plan else "index" if "Index" in plan else "seq"
            choices.append({"layout": layout, "selectivity_pct": sel * 100, "random_page_cost": rpc, "planner_pick": pick})
cross = []
for layout in ("clustered", "shuffled"):
    pts = [(r["selectivity_pct"], r["ms"]) for r in rows if r["layout"] == layout and r["method"] == "index scan"]
    seq = {r["selectivity_pct"]: r["ms"] for r in rows if r["layout"] == layout and r["method"] == "sequential scan"}
    first = next((s for s, t in pts if t > seq[s]), None)
    cross.append({"layout": layout, "index_loses_from_pct": first})
for r in rows: print(f"{r['layout']:<9} {r['selectivity_pct']:>6g}%  {r['method']:<15} {r['ms']:9.1f} ms")
for r in cross: print(f"{r['layout']}: index scan slower than the sequential scan from {r['index_loses_from_pct']}% selectivity")
for r in choices: print(f"  planner at random_page_cost {r['random_page_cost']:<4}: {r['layout']:<9} {r['selectivity_pct']:>6g}% -> {r['planner_pick']}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"times": rows, "crossover": cross, "planner": choices}}, f)