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

A 100-million-element array and a random mask at selectivities 0.01 to 0.9.
Compares (a) arr[mask], a sequential pass over everything, with (b)
arr[np.flatnonzero(mask)], building a position list and gathering from it,
and finds the selectivity where the gather starts to lose.
"""
import argparse, json, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=100_000_000, help="Array length.")
p.add_argument("--selectivities", default="0.001,0.01,0.05,0.1,0.2,0.5,0.9", help="Selectivities.")
p.add_argument("--repeats", type=int, default=3, help="Runs per point (best is kept).")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(args.seed); arr = rng.random(args.n); u = rng.random(args.n)
def best(fn):
    ts = []
    for _ in range(args.repeats):
        t = time.perf_counter(); fn(); ts.append(time.perf_counter() - t)
    return min(ts) * 1000
rows = []
for s in map(float, args.selectivities.split(",")):
    mask = u < s
    early = best(lambda: arr[mask]); pos = np.flatnonzero(mask)
    late = best(lambda: arr[pos]); late_total = best(lambda: arr[np.flatnonzero(mask)])
    rows.append({"selectivity": s, "mask_ms": early, "gather_ms": late, "positions_plus_gather_ms": late_total})
cross = next((r["selectivity"] for r in rows if r["positions_plus_gather_ms"] > r["mask_ms"]), None)
summary = {"crossover": cross}
for r in rows: print(f"s = {r['selectivity']:<5g}  arr[mask] {r['mask_ms']:7.1f} ms   positions + gather {r['positions_plus_gather_ms']:7.1f} ms   (gather alone {r['gather_ms']:7.1f} ms)")
print(f"positions + gather loses from s = {cross}")

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