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

Joins 8 to 16 small PostgreSQL tables in a chain and in a star and records the planning
time with the exhaustive (dynamic-programming) search and with the genetic
optimizer (GEQO, used above geqo_threshold). At 12 or more relations, plans
the same query 20 times and counts the distinct plans each search produces.
Needs the PostgreSQL service from the compose file (port 55704).
"""
import argparse, hashlib, json, re
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55704/postgres", help="Connection string.")
p.add_argument("--sizes", default="8,10,12,14,16", help="Numbers of joined tables.")
p.add_argument("--repeats", type=int, default=20, help="Plans per query for the distinct-plan count.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

c = psycopg.connect(args.dsn, autocommit=True)
sizes = [int(x) for x in args.sizes.split(",")]
for i in range(max(sizes)):
    c.execute(f"DROP TABLE IF EXISTS j{i}"); c.execute(f"CREATE TABLE j{i} (id int PRIMARY KEY, nxt int, v int)")
    c.execute(f"INSERT INTO j{i} SELECT g, (g * 7 + {i}) % {100 + i * 13}, g % 10 FROM generate_series(1, {100 + i * 13}) g"); c.execute(f"ANALYZE j{i}")
def query(n, shape):
    if shape == "chain": joins = " ".join(f"JOIN j{i} ON j{i}.id = j{i - 1}.nxt" for i in range(1, n))
    else: joins = " ".join(f"JOIN j{i} ON j{i}.id = j0.nxt + {i}" for i in range(1, n))      # star: every table joins the hub
    return f"SELECT count(*) FROM j0 {joins} WHERE j0.v = 3"
def plan(q):
    txt = [r[0] for r in c.execute("EXPLAIN (SUMMARY ON, COSTS OFF) " + q).fetchall()]
    t = float(re.search(r"Planning Time: ([\d.]+)", "\n".join(txt)).group(1))
    shape = hashlib.md5("\n".join(l for l in txt if "Planning Time" not in l).encode()).hexdigest()
    return t, shape
rows = []
for shape in ("chain", "star"):
    for n in sizes:
        q = query(n, shape)
        for mode in ("exhaustive (geqo off)", "GEQO, fixed seed", "GEQO, seed varied"):
            c.execute("RESET ALL"); c.execute("SET join_collapse_limit = 100"); c.execute("SET from_collapse_limit = 100")
            if mode.startswith("GEQO"): c.execute("SET geqo = on"); c.execute("SET geqo_threshold = 2")
            else: c.execute("SET geqo = off")
            times, shapes = [], set()
            for r_ in range(args.repeats if n >= 12 else 3):
                if mode == "GEQO, seed varied": c.execute(f"SET geqo_seed = {r_ / args.repeats}")
                t, sh = plan(q); times.append(t); shapes.add(sh)
            rows.append({"shape": shape, "tables": n, "search": mode, "planning_ms": sorted(times)[len(times) // 2], "distinct_plans": len(shapes) if n >= 12 else None})
for r in rows: print(f"{r['shape']:<5} {r['tables']:>2} tables, {r['search']:<22} planning {r['planning_ms']:9.2f} ms" + (f"   distinct plans in {args.repeats}: {r['distinct_plans']}" if r["distinct_plans"] else ""))

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