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

A 5-million-row queue table with 1,000 pending rows and a partial index
WHERE status = 'pending'. Plans for the literal query, for a prepared
statement with a parameter over repeated executions, and for a full-column
index; and the sizes of both indexes.
Needs the PostgreSQL service from the compose file (port 55521).
"""
import argparse, json
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55521/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=5_000_000, help="Rows.")
p.add_argument("--pending", type=int, default=1000, help="Pending 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 q")
con.execute(f"CREATE TABLE q AS SELECT g AS id, CASE WHEN g % ({args.rows} / {args.pending}) = 0 THEN 'pending' ELSE 'done' END AS status, md5(g::text) AS payload FROM generate_series(1, {args.rows}) g")
con.execute("CREATE INDEX q_pending ON q (id) WHERE status = 'pending'"); con.execute("ANALYZE q")
plan = lambda sql: [r[0] for r in con.execute("EXPLAIN " + sql)][0].split("(")[0].strip()
rows = [{"query": "literal: WHERE status = 'pending'", "plan": plan("SELECT * FROM q WHERE status = 'pending'")}]
con.execute("PREPARE p(text) AS SELECT * FROM q WHERE status = $1")
for i in range(1, 8):
    rows.append({"query": f"prepared, execution {i}", "plan": plan("EXECUTE p('pending')")})
con.execute("SET plan_cache_mode = force_generic_plan")
rows.append({"query": "prepared, generic plan forced", "plan": plan("EXECUTE p('pending')")})
con.execute("RESET plan_cache_mode")
partial = con.execute("SELECT pg_relation_size('q_pending')").fetchone()[0]
con.execute("CREATE INDEX q_status ON q (status)"); con.execute("ANALYZE q")
full = con.execute("SELECT pg_relation_size('q_status')").fetchone()[0]
for r in rows: print(f"{r['query']:<36} {r['plan']}")
print(f"partial index {partial / 1e3:,.0f} KB   full-column index {full / 1e6:,.1f} MB   ({full / partial:,.0f}x)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"plans": rows, "summary": [{"partial_kb": partial / 1e3, "full_mb": full / 1e6, "ratio": full / partial}]}}, f)