# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""As-of join against a range join.

DuckDB: prices(product_id, valid_from, price) with 1,000 products × 1,000
changes, and orders. The price at order time computed (a) with a join on
valid_from <= order_ts plus ROW_NUMBER() = 1, and (b) with ASOF JOIN: time,
the intermediate row count of (a), and whether both agree. (a) runs on a
smaller order table first and is extrapolated; (b) runs on 5 million orders.
"""
import argparse, json, re, time
import duckdb

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--orders", type=int, default=5_000_000, help="Orders for the ASOF join.")
p.add_argument("--range-orders", type=int, default=50_000, help="Orders for the range join (it is quadratic).")
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.3)")
c.execute("""CREATE TABLE prices AS SELECT p AS product_id, TIMESTAMP '2025-01-01' + (ch * 1000 + p % 997) * INTERVAL '1 minute' AS valid_from,
             (random() * 100)::DECIMAL(10, 2) AS price FROM range(1000) a(p), range(1000) b(ch)""")
c.execute(f"""CREATE TABLE orders AS SELECT i AS order_id, (random() * 999)::INT AS product_id,
              TIMESTAMP '2025-01-01' + (random() * 1000 * 1000)::INT * INTERVAL '1 minute' AS order_ts FROM range({args.orders}) t(i)""")
RANGE = """SELECT order_id, price FROM (SELECT o.order_id, p.price, ROW_NUMBER() OVER (PARTITION BY o.order_id ORDER BY p.valid_from DESC) rn
           FROM (SELECT * FROM orders WHERE order_id < {n}) o JOIN prices p ON p.product_id = o.product_id AND p.valid_from <= o.order_ts) WHERE rn = 1"""
ASOF = """SELECT o.order_id, p.price FROM (SELECT * FROM orders WHERE order_id < {n}) o
          ASOF JOIN prices p ON p.product_id = o.product_id AND o.order_ts >= p.valid_from"""
def timed(q):
    t = time.perf_counter(); r = c.execute(q).fetchall(); return time.perf_counter() - t, r
rows = []
t_range, r_range = timed(RANGE.format(n=args.range_orders))
prof = c.execute("EXPLAIN ANALYZE " + RANGE.format(n=args.range_orders)).fetchall()[0][1]
pairs = max(int(x.replace(",", "")) for x in re.findall(r"([\d,]+)\s+[Rr]ows", prof))
t_asof_small, r_asof = timed(ASOF.format(n=args.range_orders))
agree = "yes" if sorted(r_range) == sorted(r_asof) else "no"
t_asof, _ = timed(ASOF.format(n=args.orders))
rows.append({"method": "(a) range join + ROW_NUMBER", "orders": args.range_orders, "seconds": t_range, "intermediate_rows": pairs})
rows.append({"method": "(b) ASOF JOIN", "orders": args.range_orders, "seconds": t_asof_small, "intermediate_rows": None})
rows.append({"method": "(b) ASOF JOIN", "orders": args.orders, "seconds": t_asof, "intermediate_rows": None})
scale = args.orders / args.range_orders
summary = {"agree": agree, "pairs_per_order": pairs / args.range_orders, "pairs_extrapolated": pairs * scale, "range_extrapolated_s": t_range * scale}
for r in rows: print(f"{r['method']:<30} {r['orders']:>9,} orders: {r['seconds']:7.2f} s" + (f"   intermediate rows {r['intermediate_rows']:,}" if r["intermediate_rows"] else ""))
print(f"same prices on the small run: {agree}; (a) at {args.orders:,} orders would build ~{summary['pairs_extrapolated']:,.0f} pairs and take ~{summary['range_extrapolated_s']:,.0f} s")

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