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

Prints the external merge sort's pass count ⌈log_{M/B − 1}(N/M)⌉ for 1 TB of
data across memory grants M, with 1 MB and with 100 MB buffers. Then sorts a
table in PostgreSQL across a range of work_mem settings and reads, from
EXPLAIN ANALYZE and the server's sort trace (trace_sort), the sort method,
the number of initial runs, the merge passes and the disk space used.
Needs the PostgreSQL service from the compose file (port 55602).
"""
import argparse, json, math, re, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55602/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=5_000_000, help="Rows in the table to sort.")
p.add_argument("--work-mem", default="64kB,128kB,256kB,512kB,1MB,4MB,16MB,64MB,1GB", help="work_mem settings to sort under.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

model = []
for B in (0.001, 0.1):
    for M in (1, 2, 4, 8, 16, 32, 64, 128, 256):
        N = 1000; runs = math.ceil(N / M); fan = M / B - 1
        passes = math.ceil(math.log(runs, fan) - 1e-9) if runs > 1 else 0
        model.append({"buffer": f"{B * 1000:g} MB buffers", "M_gb": M, "runs": runs, "fan_in": fan, "passes": passes})

c = psycopg.connect(args.dsn, autocommit=True)
c.execute("DROP TABLE IF EXISTS s"); c.execute("CREATE TABLE s (k double precision, pad text)")
c.execute(f"INSERT INTO s SELECT random(), md5(g::text) FROM generate_series(1, {args.rows}) g"); c.execute("VACUUM ANALYZE s")
table_mb = float(c.execute("SELECT pg_relation_size('s') / 1048576.0").fetchone()[0])
c.execute("SET max_parallel_workers_per_gather = 0"); c.execute("SET trace_sort = on"); c.execute("SET client_min_messages = log")
notices = []; c.add_notice_handler(lambda d: notices.append(d.message_primary or ""))
rows = []
for wm in args.work_mem.split(","):
    c.execute(f"SET work_mem = '{wm}'"); notices.clear()
    t = time.perf_counter()
    plan = "\n".join(r[0] for r in c.execute("EXPLAIN (ANALYZE, BUFFERS) SELECT k, pad FROM s ORDER BY k").fetchall())
    el = time.perf_counter() - t
    method = re.search(r"Sort Method: ([\w -]+?)\s+(Disk|Memory): (\d+)kB", plan)
    runs = max([int(x) for n in notices for x in re.findall(r"finished writing run (\d+)", n)] or [0])
    passes = sum(1 for n in notices if re.search(r"starting merge pass", n))
    final = any(re.search(r"(final merge|performsort .* done.*merge|starting final merge)", n) for n in notices)
    temp_w = sum(int(x) for x in re.findall(r"temp read=\d+ written=(\d+)", plan)) * 8 / 1024
    rows.append({"work_mem": wm, "method": method.group(1).strip() if method else "?", "space_kb": int(method.group(3)) if method else None,
                 "initial_runs": runs, "merge_passes": passes + (1 if runs > 1 else 0), "temp_written_mb": temp_w, "seconds": el})
summary = {"table_mb": table_mb, "rows": args.rows}
for r in model:
    if r["M_gb"] <= 16: print(f"{r['buffer']:<17} M = {r['M_gb']:>3} GB: {r['runs']:>4} runs, fan-in {r['fan_in']:>6.0f}, {r['passes']} pass{'es' if r['passes'] != 1 else ''}")
print(f"PostgreSQL: sorting {args.rows:,} rows ({table_mb:.0f} MB)")
for r in rows: print(f"  work_mem {r['work_mem']:>6}: {r['method']:<15} runs {r['initial_runs']:>5}  merge passes {r['merge_passes']}  temp written {r['temp_written_mb']:7.1f} MB  {r['seconds']:.2f} s")

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