# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Build a pruning cascade and count the work.

10 million rows (tracking_id, ts) written in time order into 100 files. Per
file: a zone map on ts, a Bloom filter on tracking_id (1% false positives),
and a sparse index with one entry per 8,192 rows (rows sorted by tracking_id
within each file). 1,000 queries for one tracking_id in a one-day range:
rows examined after each stage. Then the same with the rows shuffled.
"""
import argparse, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=10_000_000, help="Rows.")
p.add_argument("--files", type=int, default=100, help="Files.")
p.add_argument("--ids", type=int, default=2_000_000, help="Distinct tracking ids.")
p.add_argument("--days", type=int, default=100, help="Days of data.")
p.add_argument("--granule", type=int, default=8192, help="Rows per sparse-index entry.")
p.add_argument("--queries", type=int, default=1000, help="Queries.")
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)
ts = np.sort(rng.uniform(0, args.days, args.rows))
tid = rng.integers(0, args.ids, args.rows)

def build(order):
    t, k = ts[order], tid[order]
    files = []
    for f in np.array_split(np.arange(args.rows), args.files):
        ft, fk = t[f], k[f]
        srt = np.argsort(fk, kind="stable"); ft, fk = ft[srt], fk[srt]           # within a file, sort by tracking_id
        m = int(len(fk) * 9.6); kk = 7
        bits = np.zeros(m, dtype=bool)
        h = [((fk * (2 * i + 1) * 0x9E3779B1) ^ (fk >> 3)) % m for i in range(kk)]
        for x in h: bits[x] = True
        files.append({"tmin": ft.min(), "tmax": ft.max(), "bloom": (bits, m, kk), "keys": fk, "sparse": fk[:: args.granule]})
    return files

def bloom_has(b, key):
    bits, m, kk = b
    return all(bits[((key * (2 * i + 1) * 0x9E3779B1) ^ (key >> 3)) % m] for i in range(kk))

def run(files):
    qs = [(int(rng.integers(0, args.ids)), float(rng.uniform(0, args.days - 1))) for _ in range(args.queries)]
    work = {"no pruning": 0, "zone map": 0, "zone map + Bloom": 0, "zone map + Bloom + sparse index": 0}
    for key, day in qs:
        for f in files:
            n = len(f["keys"]); work["no pruning"] += n
            if f["tmax"] < day or f["tmin"] > day + 1: continue
            work["zone map"] += n
            if not bloom_has(f["bloom"], key): continue
            work["zone map + Bloom"] += n
            g = np.searchsorted(f["sparse"], key, side="right") - 1                # the granule that can hold the key
            work["zone map + Bloom + sparse index"] += min(args.granule, n - max(g, 0) * args.granule)
    return {k: v / args.queries for k, v in work.items()}

rows = []
for label, order in (("written in time order", np.arange(args.rows)), ("rows shuffled", rng.permutation(args.rows))):
    w = run(build(order))
    for stage, v in w.items(): rows.append({"layout": label, "stage": stage, "rows_examined": v})
    print(f"{label}: " + "   ".join(f"{s}: {v:,.0f}" for s, v in w.items()))

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