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

DuckDB, small tables. NOT IN against NOT EXISTS before and after one NULL is
inserted into the subquery's table; a LEFT JOIN with the filter in WHERE
against the filter in ON; and SUM(amount) before and after a join to a table
with two rows per order.
"""
import argparse, json
import duckdb

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

c = duckdb.connect()
c.execute("CREATE TABLE customers AS SELECT i AS customer_id, i % 5 AS tier FROM range(1, 1001) t(i)")
c.execute("CREATE TABLE orders AS SELECT i AS order_id, 1 + i % 1000 AS customer_id, 10 AS amount FROM range(0, 5000) t(i)")
c.execute("CREATE TABLE blocked AS SELECT * FROM (VALUES (1), (2), (3)) t(customer_id)")
c.execute("CREATE TABLE shipments AS SELECT order_id, s FROM orders, range(2) t(s)")          # two rows per order
q = lambda s: c.execute(s).fetchone()[0]
NOT_IN = "SELECT count(*) FROM customers WHERE customer_id NOT IN (SELECT customer_id FROM blocked)"
NOT_EX = "SELECT count(*) FROM customers c WHERE NOT EXISTS (SELECT 1 FROM blocked b WHERE b.customer_id = c.customer_id)"
rows = [{"trap": "NOT IN", "case": "blocked has 3 ids", "result": q(NOT_IN)}, {"trap": "NOT EXISTS", "case": "blocked has 3 ids", "result": q(NOT_EX)}]
c.execute("INSERT INTO blocked VALUES (NULL)")
rows += [{"trap": "NOT IN", "case": "plus one NULL", "result": q(NOT_IN)}, {"trap": "NOT EXISTS", "case": "plus one NULL", "result": q(NOT_EX)}]
rows += [{"trap": "LEFT JOIN, filter in WHERE", "case": "c.tier = 3", "result": q("SELECT count(*) FROM orders o LEFT JOIN customers c USING (customer_id) WHERE c.tier = 3")},
         {"trap": "LEFT JOIN, filter in ON", "case": "AND c.tier = 3", "result": q("SELECT count(*) FROM orders o LEFT JOIN customers c ON c.customer_id = o.customer_id AND c.tier = 3")}]
rows += [{"trap": "SUM(amount)", "case": "orders alone", "result": q("SELECT sum(amount) FROM orders")},
         {"trap": "SUM(amount)", "case": "after joining two shipments per order", "result": q("SELECT sum(amount) FROM orders JOIN shipments USING (order_id)")}]
for r in rows: print(f"{r['trap']:<28} {r['case']:<40} {r['result']:>8,}")

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