# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb", "pyarrow"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Blocking against streaming: time to first row.

DuckDB with 20 million orders. For LIMIT 10, a full ORDER BY amount DESC and
ORDER BY ... LIMIT 10 (top-N), measures the time until the first batch of
rows reaches the client, the total time to drain the result, and the peak
memory DuckDB reports for the query.
"""
import argparse, json, time
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--orders", type=int, default=20_000_000, help="Orders.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

c = duckdb.connect(); c.execute("SELECT setseed(0.42)")
c.execute(f"CREATE TABLE orders AS SELECT i AS order_id, (random() * 500)::INT AS amount, repeat('x', 40) AS pad FROM range({args.orders}) t(i)")
def stream(q):
    c.execute("SELECT 1").fetchall()
    t = time.perf_counter(); res = c.execute(q)
    first = res.to_arrow_reader(100_000)
    b = first.read_next_batch(); t_first = time.perf_counter() - t; n = b.num_rows
    for b in first: n += b.num_rows
    total = time.perf_counter() - t
    return t_first, total, n
rows = []
for label, q in (("LIMIT 10", "SELECT * FROM orders LIMIT 10"),
                 ("ORDER BY amount DESC (full sort)", "SELECT * FROM orders ORDER BY amount DESC"),
                 ("ORDER BY amount DESC LIMIT 10 (top-N)", "SELECT * FROM orders ORDER BY amount DESC LIMIT 10")):
    stream(q); tf, tt, n = stream(q)
    plan = c.execute("EXPLAIN " + q).fetchall()[0][1]
    rows.append({"query": label, "first_row_s": tf, "total_s": tt, "rows": n, "operator": "TOP_N" if "TOP_N" in plan else "ORDER_BY" if "ORDER_BY" in plan else "LIMIT"})
for r in rows: print(f"{r['query']:<38} first rows after {r['first_row_s']:7.3f} s   all {r['rows']:>11,} rows after {r['total_s']:7.3f} s   operator {r['operator']}")

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