# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Join algorithms and the memory threshold.

DuckDB joins a 20-million-row table to a 5-million-row table (with a text
payload) on an integer key, under memory limits from 4 GB down to 100 MB, each in a fresh process,
recording the time and the temp space used. Then threads=1 against threads=8
at the largest limit.
"""
import argparse, json, os, shutil, subprocess, sys, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--probe", type=int, default=20_000_000, help="Rows in the probe (large) table.")
p.add_argument("--build", type=int, default=5_000_000, help="Rows in the build (smaller) table.")
p.add_argument("--limits", default="4GB,1GB,500MB,300MB,200MB,100MB", help="memory_limit values.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

CHILD = r"""
import duckdb, json, os, sys, threading, time
db, tmp, limit, threads = sys.argv[1:5]
c = duckdb.connect(db, read_only=True); c.execute(f"SET memory_limit='{limit}'"); c.execute(f"SET temp_directory='{tmp}'"); c.execute(f"SET threads={threads}"); c.execute("SET enable_progress_bar=false")
peak = [0]; stop = [False]
def watch():
    while not stop[0]:
        peak[0] = max(peak[0], sum(os.path.getsize(os.path.join(d, f)) for d, _, fs in os.walk(tmp) for f in fs if os.path.exists(os.path.join(d, f))))
        time.sleep(0.05)
threading.Thread(target=watch, daemon=True).start()
q = "SELECT count(*), sum(b.v), sum(length(b.pad)) FROM big a JOIN small b ON a.k = b.k"   # the build side carries a payload
t = time.perf_counter()
try: c.execute(q).fetchall(); ok = "ok"
except Exception as e: ok = str(e).splitlines()[0][:100]
el = time.perf_counter() - t; stop[0] = True
print(json.dumps({"seconds": el, "temp_mb": peak[0] / 2**20, "outcome": ok}))
"""
work = tempfile.mkdtemp(); db = os.path.join(work, "j.duckdb")
import duckdb
c = duckdb.connect(db); c.execute("SELECT setseed(0.1)")
c.execute(f"CREATE TABLE small AS SELECT i * 7919 AS k, i * 2 AS v, md5(i::VARCHAR) AS pad FROM range({args.build}) t(i)")   # sparse keys: no perfect hash
c.execute(f"CREATE TABLE big AS SELECT (random() * {args.build})::BIGINT * 7919 AS k, i AS w FROM range({args.probe}) t(i)"); c.close()
def run(limit, threads):
    tmp = os.path.join(work, "tmp"); shutil.rmtree(tmp, ignore_errors=True); os.makedirs(tmp)
    return json.loads(subprocess.run([sys.executable, "-c", CHILD, db, tmp, limit, str(threads)], capture_output=True, text=True, check=True).stdout.strip().splitlines()[-1])
rows = [{"memory_limit": l, "threads": 8, **run(l, 8)} for l in args.limits.split(",")]
first = args.limits.split(",")[0]
threads = [{"threads": 1, **run(first, 1)}, {"threads": 8, **next({k: v for k, v in r.items() if k not in ("memory_limit", "threads")} for r in rows if r["memory_limit"] == first)}]
shutil.rmtree(work, ignore_errors=True)
for r in rows: print(f"memory_limit {r['memory_limit']:>6}: {r['seconds']:6.2f} s   temp {r['temp_mb']:7.0f} MB   {r['outcome']}")
for r in threads: print(f"threads={r['threads']}: {r['seconds']:.2f} s")

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