# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Estimate against actual on a correlated filter.

A PostgreSQL table where postcode implies city. EXPLAIN ANALYZE of a filter
on both columns shows the planner's row estimate (it multiplies the two
selectivities as if they were independent) against the actual rows; then
CREATE STATISTICS (dependencies) and the same query again; then the same
with an equality filter, and with an MCV statistic over both columns.
Needs the PostgreSQL service from the compose file (port 55701).
"""
import argparse, json, re
import psycopg

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

c = psycopg.connect(args.dsn, autocommit=True)
c.execute("DROP TABLE IF EXISTS people")
c.execute(f"""CREATE TABLE people AS SELECT g AS id, (g % 100) AS city_id, (g % 100) * 1000 + (g % 7) AS postcode
              FROM generate_series(1, {args.rows}) g""")                # postcode implies city
c.execute("ANALYZE people")
Q = "SELECT * FROM people WHERE city_id = 5 AND postcode BETWEEN 5000 AND 5006"
EQ = "SELECT * FROM people WHERE city_id = 5 AND postcode = 5003"
def est(q=Q):
    plan = "\n".join(r[0] for r in c.execute("EXPLAIN ANALYZE " + q).fetchall())
    m = re.search(r"rows=(\d+) width=\d+\) \(actual time=[\d.]+\.\.[\d.]+ rows=([\d.]+)", plan)
    return int(m.group(1)), float(m.group(2))
rows = []
def record(stats):
    for filt, q in (("BETWEEN", Q), ("equality", EQ)):
        e, a = est(q); rows.append({"statistics": stats, "filter": filt, "estimate": e, "actual": a, "error_x": a / e})
record("per-column only")
c.execute("CREATE STATISTICS s (dependencies) ON city_id, postcode FROM people"); c.execute("ANALYZE people"); record("dependencies")
c.execute("DROP STATISTICS s"); c.execute("ALTER TABLE people ALTER postcode SET STATISTICS 1000")
c.execute("CREATE STATISTICS s2 (mcv) ON city_id, postcode FROM people"); c.execute("ANALYZE people"); record("mcv (statistics target 1000)")
for r in rows: print(f"{r['statistics']:<30} {r['filter']:<9} estimate {r['estimate']:>6,}   actual {r['actual']:>8,.0f}   actual / estimate {r['error_x']:.1f}x")

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