# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The sort that is not an external-sort problem.

DuckDB with memory_limit = 500 MB sorts 20 million rows of (random key, 100
bytes of padding), about 2.4 GB, and records the peak size of its temp files,
the time and the CPU time; then the same table already sorted by the key.
Finally the one-pass merge capacity M²/B for M = 500 MB and B = 256 KB.
"""
import argparse, json, os, resource, shutil, tempfile, threading, time
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=20_000_000, help="Rows.")
p.add_argument("--memory-limit", default="500MB", help="DuckDB memory_limit.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

tmp = tempfile.mkdtemp(prefix="duck_spill"); dbdir = tempfile.mkdtemp()
con = duckdb.connect(os.path.join(dbdir, "s.duckdb"))
con.execute(f"PRAGMA memory_limit='{args.memory_limit}'"); con.execute(f"PRAGMA temp_directory='{tmp}'")
con.execute(f"CREATE TABLE t AS SELECT random() AS k, repeat('x', 100) AS pad FROM range({args.rows})")
con.execute("CREATE TABLE t_sorted AS SELECT * FROM t ORDER BY k")
data_mb = con.execute("SELECT sum(length(pad) + 8) / 1048576 FROM t").fetchone()[0]
def watch(stop, peak):
    while not stop.is_set():
        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)
rows = []
for table in ("t", "t_sorted"):
    stop, peak = threading.Event(), [0]; th = threading.Thread(target=watch, args=(stop, peak)); th.start()
    c0 = resource.getrusage(resource.RUSAGE_SELF); t = time.perf_counter()
    con.execute(f"CREATE OR REPLACE TEMP TABLE out AS SELECT * FROM {table} ORDER BY k")   # the sorted result must be produced in full
    wall = time.perf_counter() - t; c1 = resource.getrusage(resource.RUSAGE_SELF); stop.set(); th.join()
    cpu = (c1.ru_utime - c0.ru_utime) + (c1.ru_stime - c0.ru_stime)
    rows.append({"input": "random order" if table == "t" else "already sorted", "seconds": wall, "cpu_s": cpu, "cpu_per_wall": cpu / wall,
                 "temp_peak_mb": peak[0] / 2**20})
M, B = 500 * 2**20, 256 * 2**10
summary = {"data_mb": data_mb, "one_pass_tb": M * M / B / 2**40, "fan_in": M // B}
for r in rows: print(f"{r['input']:<15} {r['seconds']:6.2f} s wall, {r['cpu_s']:6.2f} s CPU ({r['cpu_per_wall']:.1f} cores busy), temp files peaked at {r['temp_peak_mb']:,.0f} MB")
print(f"data ~{data_mb:,.0f} MB; one merge pass handles M²/B = {summary['one_pass_tb']:.2f} TB with fan-in {summary['fan_in']:,}")
con.close(); shutil.rmtree(tmp, ignore_errors=True); shutil.rmtree(dbdir, ignore_errors=True)

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