# /// script
# requires-python = ">=3.10"
# dependencies = ["duckdb"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Same answer, different cost.

DuckDB, 1 million customers and 20 million orders. The same report written
two ways: (A) join, then filter amount > 400, then group by tier; (B) filter
orders in a subquery first, then join. Then (C) the join placed under a window
function, which the optimizer cannot push the filter past. Times each and reads the row count each
operator produced (EXPLAIN ANALYZE), to find where the rows peak.
"""
import argparse, json, re, 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("CREATE TABLE customers AS SELECT i AS customer_id, i % 5 AS tier FROM range(1, 1000001) t(i)")
c.execute(f"""CREATE TABLE orders AS SELECT i AS order_id, 1 + (random() * 999999)::INT AS customer_id, (random() * 500)::INT AS amount
              FROM range(0, {args.orders}) t(i)""")
Q = {"A: join, then filter": "SELECT tier, sum(amount) FROM orders o JOIN customers c USING (customer_id) WHERE amount > 400 GROUP BY tier ORDER BY tier",
     "B: filter, then join": "SELECT tier, sum(amount) FROM (SELECT * FROM orders WHERE amount > 400) o JOIN customers c USING (customer_id) GROUP BY tier ORDER BY tier",
     "C: filter after a window": """SELECT tier, sum(amount) FROM (SELECT o.amount, c.tier, row_number() OVER () AS rn
                                       FROM orders o JOIN customers c USING (customer_id)) j
                                       WHERE amount > 400 GROUP BY tier ORDER BY tier"""}                      # a filter cannot move below the window
rows, answers = [], {}
for label, q in Q.items():
    answers[label] = c.execute(q).fetchall()
    ts = []
    for _ in range(3):
        t = time.perf_counter(); c.execute(q).fetchall(); ts.append(time.perf_counter() - t)
    prof = c.execute("EXPLAIN ANALYZE " + q).fetchall()[0][1]
    join_in = [int(x.replace(",", "")) for x in re.findall(r"HASH_JOIN[\s\S]*?([\d,]+)\s+[Rr]ows", prof)]
    counts = [int(x.replace(",", "")) for x in re.findall(r"([\d,]+)\s+[Rr]ows", prof)]
    rows.append({"query": label, "seconds": min(ts), "join_output_rows": join_in[0] if join_in else None, "peak_rows": max(counts) if counts else None})
summary = {"same_answer": "yes" if len({str(v) for v in answers.values()}) == 1 else "no"}
for r in rows: print(f"{r['query']:<22} {r['seconds']:6.3f} s   rows out of the join {r['join_output_rows']:>12,}   largest operator output {r['peak_rows']:>12,}")
print("all three return the same answer:", summary["same_answer"])

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