# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Composite index order and the range column.

5 million random rows t(a, b, c) with indexes (a, b, c) and (a, c, b), each
tried alone for a = 1 AND b > 5 AND c = 3: what goes into Index Cond and what
is left to Filter, the buffers read, and the time.
Needs the PostgreSQL service from the compose file (port 55522).
"""
import argparse, json, re
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55522/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()
con = psycopg.connect(args.dsn, autocommit=True)
con.execute("DROP TABLE IF EXISTS t")
con.execute(f"CREATE TABLE t AS SELECT (random()*100)::int AS a, (random()*100)::int AS b, (random()*100)::int AS c FROM generate_series(1, {args.rows})")
SQL = "SELECT * FROM t WHERE a = 1 AND b > 5 AND c = 3"
rows = []
for cols in ("a, b, c", "a, c, b"):
    con.execute("DROP INDEX IF EXISTS ix"); con.execute(f"CREATE INDEX ix ON t ({cols})"); con.execute("ANALYZE t")
    con.execute("SET enable_seqscan = off; SET enable_bitmapscan = off")
    plan = "\n".join(r[0] for r in con.execute("EXPLAIN (ANALYZE, BUFFERS) " + SQL))
    cond = re.search(r"Index Cond: \((.*)\)", plan); filt = re.search(r"Filter: \((.*)\)", plan)
    removed = re.search(r"Rows Removed by Filter: (\d+)", plan); bufs = re.search(r"Buffers: shared hit=(\d+)(?: read=(\d+))?", plan)
    rows.append({"index": f"({cols})", "index_cond": cond.group(1) if cond else "", "filter": filt.group(1) if filt else "(none)",
                 "rows_removed_by_filter": int(removed.group(1)) if removed else 0,
                 "buffers": int(bufs.group(1)) + int(bufs.group(2) or 0) if bufs else 0,
                 "ms": float(re.search(r"Execution Time: ([\d.]+)", plan).group(1))})
    r = rows[-1]
    print(f"index ({cols}): Index Cond [{r['index_cond']}]  Filter [{r['filter']}]  removed by filter {r['rows_removed_by_filter']:,}  "
          f"buffers {r['buffers']:,}  {r['ms']:.1f} ms")

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