# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Four pairs of equal queries.

Pairs of PostgreSQL queries that return the same answer and can run very
differently: NOT IN against NOT EXISTS (with NULLs in the subquery column,
then after making it NOT NULL); a function on a partition key against a
half-open range (partitions scanned); a CTE referenced twice, MATERIALIZED
against NOT MATERIALIZED.
Needs the PostgreSQL service from the compose file (port 55705).
"""
import argparse, json, re
import psycopg

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

c = psycopg.connect(args.dsn, autocommit=True)
def run(sql):
    txt = "\n".join(r[0] for r in c.execute("EXPLAIN (ANALYZE) " + sql).fetchall())
    return float(re.search(r"Execution Time: ([\d.]+)", txt).group(1)), txt
def answer(sql): return c.execute(sql).fetchone()[0]
c.execute("DROP TABLE IF EXISTS customers, refunds, sales CASCADE")
c.execute("CREATE TABLE customers AS SELECT g AS id FROM generate_series(1, 100000) g")
c.execute("CREATE TABLE refunds AS SELECT CASE WHEN g % 100 = 0 THEN NULL ELSE g * 2 END AS customer_id FROM generate_series(1, 50000) g")
c.execute("ANALYZE customers, refunds")
rows = []
q_in = "SELECT count(*) FROM customers WHERE id NOT IN (SELECT customer_id FROM refunds)"
q_ex = "SELECT count(*) FROM customers c WHERE NOT EXISTS (SELECT 1 FROM refunds r WHERE r.customer_id = c.id)"
for label, q in (("NOT IN, NULLs present", q_in), ("NOT EXISTS, NULLs present", q_ex)):
    t, _ = run(q); rows.append({"pair": "1. anti-join", "query": label, "ms": t, "answer": answer(q), "partitions": None})
c.execute("DELETE FROM refunds WHERE customer_id IS NULL"); c.execute("ALTER TABLE refunds ALTER customer_id SET NOT NULL"); c.execute("ANALYZE refunds")
t, _ = run(q_in); rows.append({"pair": "1. anti-join", "query": "NOT IN, column NOT NULL", "ms": t, "answer": answer(q_in), "partitions": None})

c.execute("CREATE TABLE sales (ts timestamp, amount int) PARTITION BY RANGE (ts)")
for y in (2023, 2024, 2025):
    for m in range(1, 13):
        nm = (y + (m == 12), m % 12 + 1)
        c.execute(f"CREATE TABLE sales_{y}_{m:02d} PARTITION OF sales FOR VALUES FROM ('{y}-{m:02d}-01') TO ('{nm[0]}-{nm[1]:02d}-01')")
c.execute("INSERT INTO sales SELECT TIMESTAMP '2023-01-01' + (g % (3 * 365 * 24)) * INTERVAL '1 hour', g % 100 FROM generate_series(1, 1000000) g"); c.execute("ANALYZE sales")
for label, q in (("extract(year FROM ts) = 2024", "SELECT sum(amount) FROM sales WHERE extract(year FROM ts) = 2024"),
                 ("ts >= '2024-01-01' AND ts < '2025-01-01'", "SELECT sum(amount) FROM sales WHERE ts >= '2024-01-01' AND ts < '2025-01-01'")):
    t, txt = run(q); parts = len(re.findall(r"Scan on sales_\d+_\d+", txt))
    rows.append({"pair": "2. partition key", "query": label, "ms": t, "answer": answer(q), "partitions": parts})

for label, mat in (("CTE twice, MATERIALIZED", "MATERIALIZED"), ("CTE twice, NOT MATERIALIZED", "NOT MATERIALIZED")):
    q = f"""WITH big AS {mat} (SELECT id, id % 10 AS g FROM customers)
            SELECT count(*) FROM big a JOIN big b ON a.id = b.id WHERE a.g = 3"""
    t, _ = run(q); rows.append({"pair": "3. CTE", "query": label, "ms": t, "answer": answer(q), "partitions": None})
for r in rows: print(f"{r['pair']:<17} {r['query']:<42} {r['ms']:9.1f} ms   answer {r['answer']}" + (f"   partitions scanned {r['partitions']}" if r["partitions"] is not None else ""))

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