# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Sort versus hash under a hot key.

50 million (meter_id, kwh) rows, 30% of them with meter_id = 0. DuckDB with a
500 MB memory limit runs an algebraic aggregation (SUM, COUNT) and a holistic
one (median) by meter_id; each query runs in its own process, which reports
its peak memory and the most temp-directory space it saw. Then a grace-style
hash partitioner splits the rows into 16 partitions under 5 hash seeds and
prints the largest partition's share.
"""
import argparse, json, os, shutil, subprocess, sys, tempfile
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=50_000_000, help="Rows.")
p.add_argument("--hot-share", type=float, default=0.3, help="Share of rows with meter_id = 0.")
p.add_argument("--memory-limit", default="500MB", help="DuckDB memory_limit.")
p.add_argument("--partitions", type=int, default=16, help="Partitions in the grace partitioner.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

CHILD = r"""
import duckdb, json, os, resource, sys, threading, time
db, tmp, limit, sql = sys.argv[1:5]
peak_tmp = [0]; stop = [False]
def watch():
    while not stop[0]:
        size = sum(os.path.getsize(os.path.join(dp, f)) for dp, _, fs in os.walk(tmp) for f in fs if os.path.exists(os.path.join(dp, f)))
        peak_tmp[0] = max(peak_tmp[0], size); time.sleep(0.05)
c = duckdb.connect(db, read_only=True)
c.execute(f"SET memory_limit='{limit}'"); c.execute(f"SET temp_directory='{tmp}'"); c.execute("SET threads=4"); c.execute("SET enable_progress_bar=false")
threading.Thread(target=watch, daemon=True).start()
t = time.perf_counter(); out = {"ok": True}
try: n = len(c.execute(sql).fetchall())
except Exception as e: out = {"ok": False, "error": str(e).splitlines()[0][:120]}; n = 0
out.update(seconds=time.perf_counter() - t, groups=n); stop[0] = True; time.sleep(0.06)
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
out.update(peak_mb=rss / (2**20 if sys.platform == "darwin" else 2**10), temp_mb=peak_tmp[0] / 2**20)
print(json.dumps(out))
"""
work = tempfile.mkdtemp(); db = os.path.join(work, "t.duckdb")
import duckdb
c = duckdb.connect(db)
c.execute(f"""CREATE TABLE t AS SELECT CASE WHEN random() < {args.hot_share} THEN 0 ELSE 1 + (random() * 999999)::INT END AS meter_id,
              random() * 10 AS kwh FROM range({args.rows})""")
hot = c.execute("SELECT count(*) FROM t WHERE meter_id = 0").fetchone()[0]; c.close()

rows = []
for kind, sql in (("algebraic: SUM, COUNT", "SELECT meter_id, SUM(kwh), COUNT(*) FROM t GROUP BY meter_id"),
                  ("holistic: median", "SELECT meter_id, median(kwh) FROM t GROUP BY meter_id")):
    tmp = os.path.join(work, "tmp"); shutil.rmtree(tmp, ignore_errors=True); os.makedirs(tmp)
    r = json.loads(subprocess.run([sys.executable, "-c", CHILD, db, tmp, args.memory_limit, sql], capture_output=True, text=True, check=True).stdout.strip().splitlines()[-1])
    rows.append({"query": kind, **r, "outcome": "ok" if r["ok"] else r["error"]})

# grace-style partitioner: partition = hash(meter_id, seed) mod P
ids = duckdb.connect(db, read_only=True).execute("SELECT meter_id FROM t").fetchnumpy()["meter_id"].astype(np.uint64)
def h(x, seed):
    z = x + np.uint64(seed) * np.uint64(0x9E3779B97F4A7C15)
    z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9); z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
    return z ^ (z >> np.uint64(31))
parts = []
with np.errstate(over="ignore"):
    for seed in range(5):
        sizes = np.bincount((h(ids, seed) % np.uint64(args.partitions)).astype(np.int64), minlength=args.partitions)
        parts.append({"seed": seed, "largest_pct": sizes.max() / len(ids) * 100, "smallest_pct": sizes.min() / len(ids) * 100,
                      "even_pct": 100 / args.partitions})
shutil.rmtree(work, ignore_errors=True)
summary = {"hot_rows": hot, "hot_pct": hot / args.rows * 100, "memory_limit": args.memory_limit}
print(f"{args.rows:,} rows, {hot:,} ({summary['hot_pct']:.1f}%) with meter_id = 0; DuckDB memory_limit {args.memory_limit}")
for r in rows: print(f"{r['query']:<24} {r['seconds']:6.2f} s   peak memory {r['peak_mb']:7.0f} MB   temp space {r['temp_mb']:7.0f} MB   {r['outcome']}")
for r in parts: print(f"hash seed {r['seed']}: largest of {args.partitions} partitions holds {r['largest_pct']:.1f}% of the rows (even split {r['even_pct']:.1f}%)")

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