# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""B-tree key order: random UUIDs against a time-ordered key.

Two PostgreSQL tables with the same 200-byte payload, one keyed by a random
UUID (v4), one by an increasing bigint (time-ordered, like UUIDv7). 5 million
rows each: insert time, primary-key index size, leaf-page fill, and WAL bytes
written.
Needs the PostgreSQL service from the compose file (port 55490).
"""
import argparse, json, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--dsn", default="postgresql://postgres:lab@localhost:55490/postgres", help="Connection string.")
p.add_argument("--rows", type=int, default=5_000_000, help="Rows per table.")
p.add_argument("--batch", type=int, default=50_000, help="Rows per INSERT statement (so keys arrive in many batches).")
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("CREATE EXTENSION IF NOT EXISTS pgstattuple")
rows = []
for name, ddl, key in (("random UUID (v4)", "id uuid PRIMARY KEY", "gen_random_uuid()"),
                       ("time-ordered bigint", "id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY", None)):
    tbl = "t_uuid" if key else "t_seq"
    con.execute(f"DROP TABLE IF EXISTS {tbl}"); con.execute(f"CREATE TABLE {tbl} ({ddl}, payload text)")
    con.execute("CHECKPOINT")
    wal0 = con.execute("SELECT pg_current_wal_lsn()").fetchone()[0]
    t = time.perf_counter()
    for _ in range(args.rows // args.batch):
        if key: con.execute(f"INSERT INTO {tbl} SELECT {key}, repeat('x', 200) FROM generate_series(1, {args.batch})")
        else: con.execute(f"INSERT INTO {tbl} (payload) SELECT repeat('x', 200) FROM generate_series(1, {args.batch})")
    dt = time.perf_counter() - t
    wal = con.execute("SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), %s)", (wal0,)).fetchone()[0]
    idx = con.execute(f"SELECT pg_relation_size('{tbl}_pkey')").fetchone()[0]
    fill = con.execute(f"SELECT avg_leaf_density FROM pgstatindex('{tbl}_pkey')").fetchone()[0]
    rows.append({"key": name, "insert_s": dt, "index_mb": idx / 1e6, "leaf_fill_pct": float(fill), "wal_mb": float(wal) / 1e6})
    print(f"{name:<22} insert {dt:6.1f} s   index {idx / 1e6:7.1f} MB   leaf pages {float(fill):4.1f}% full   WAL {float(wal) / 1e6:8.1f} MB")
ratio = rows[0]["index_mb"] / rows[1]["index_mb"]
print(f"random-key index / ordered-key index = {ratio:.2f}x (a uuid key is also twice the width of a bigint)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"keys": rows, "summary": [{"index_ratio": ratio, "wal_ratio": rows[0]["wal_mb"] / rows[1]["wal_mb"],
                   "time_ratio": rows[0]["insert_s"] / rows[1]["insert_s"]}]}}, f)