# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Top-k with a bounded heap.

A file of random numbers; the top 10 by `sort -nr | head -10` and by a
bounded heap (heapq.nlargest): time and peak memory of each. Then DuckDB's
plan for ORDER BY x LIMIT 10.
"""
import argparse, heapq, json, os, random, resource, subprocess, sys, tempfile, time
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--numbers", type=int, default=20_000_000, help="Numbers in the file (the chapter uses 100 million).")
p.add_argument("--k", type=int, default=10, help="k.")
p.add_argument("--path", default=os.path.join(tempfile.gettempdir(), "topk-lab.txt"), help="The numbers file.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
import atexit, shutil
atexit.register(lambda: shutil.rmtree(args.path, ignore_errors=True) if os.path.isdir(args.path) else os.path.exists(args.path) and os.remove(args.path))   # leave nothing behind
random.seed(0)
if not os.path.exists(args.path) or os.path.getsize(args.path) < args.numbers * 5:
    with open(args.path, "w") as f:
        for _ in range(args.numbers // 100_000):
            f.write("\n".join(str(random.randrange(10**12)) for _ in range(100_000)) + "\n")

def child_peak_mb():
    r = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
    return r / 1e6 if sys.platform == "darwin" else r / 1e3                     # bytes on macOS, KB on Linux

t = time.perf_counter()
out = subprocess.run(f"sort -nr {args.path} | head -{args.k}", shell=True, capture_output=True, text=True).stdout.split()
t_sort, m_sort = time.perf_counter() - t, child_peak_mb()

heap_code = f"import heapq,resource,sys;print(heapq.nlargest({args.k}, map(int, open({args.path!r}))));" \
            "r=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss;print(r/1e6 if sys.platform=='darwin' else r/1e3)"
t = time.perf_counter()
res = subprocess.run([sys.executable, "-c", heap_code], capture_output=True, text=True).stdout.split("\n")
t_heap, m_heap = time.perf_counter() - t, float(res[1])
base = subprocess.run([sys.executable, "-c", "import resource,sys;r=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss;print(r/1e6 if sys.platform=='darwin' else r/1e3)"],
                      capture_output=True, text=True).stdout
m_heap_extra = m_heap - float(base)

con = duckdb.connect()
plan = con.execute(f"EXPLAIN SELECT * FROM read_csv('{args.path}', columns={{'x': 'BIGINT'}}, header=false) ORDER BY x DESC LIMIT {args.k}").fetchall()
plan_text = " ".join(r[1] for r in plan)
op = "TOP_N" if "TOP_N" in plan_text else ("ORDER_BY" if "ORDER" in plan_text else "?")
t = time.perf_counter(); con.execute(f"SELECT * FROM read_csv('{args.path}', columns={{'x': 'BIGINT'}}, header=false) ORDER BY x DESC LIMIT {args.k}").fetchall()
t_duck = time.perf_counter() - t
rows = [{"method": "sort -nr | head", "seconds": t_sort, "peak_mb": m_sort},
        {"method": "bounded heap (heapq.nlargest)", "seconds": t_heap, "peak_mb": m_heap_extra}]
for r in rows: print(f"{r['method']:<30} {r['seconds']:7.1f} s   peak memory {r['peak_mb']:8.1f} MB")
print(f"DuckDB plan for ORDER BY ... LIMIT {args.k}: {op} operator ({t_duck:.1f} s)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"methods": rows, "summary": [{"duckdb_op": op, "duckdb_s": t_duck, "numbers": args.numbers}]}}, f)