# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Block-max skipping on a toy index.

Posting lists for 50,000 documents and 5 terms with Zipfian frequencies; a
document's score is the sum of per-term contributions. Exhaustive scoring
against top-10 with block-max skipping (blocks of 128 postings whose stored
maximum, plus the other terms' maximums, cannot beat the tenth-best score are
skipped). Postings read and whether the top 10 agree; then a score with a
per-document boost (about 1% of documents) that is not part of any term's
stored maximum, so the block bounds no longer hold.
"""
import argparse, heapq, json
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--docs", type=int, default=50_000, help="Documents.")
p.add_argument("--block", type=int, default=128, help="Postings per block.")
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)
shares = [0.6, 0.2, 0.05, 0.01, 0.002]                                      # Zipf-like document frequencies of the 5 terms
post = {}
for t, s in enumerate(shares):
    docs = np.sort(rng.choice(args.docs, int(args.docs * s), replace=False))
    post[t] = (docs, rng.random(len(docs)) * (1 + 3 * (1 - s)))             # rarer terms weigh more, like idf
def exhaustive(terms, bonus=False):
    score = {}; read = 0
    for t in terms:
        d, w = post[t]; read += len(d)
        for doc, val in zip(d.tolist(), w.tolist()): score[doc] = score.get(doc, 0) + val
    if bonus:
        for doc in score:
            if doc % 97 == 0: score[doc] += 4                            # a boost that is not part of any term's stored maximum
    return [d for _, d in heapq.nlargest(10, ((v, d) for d, v in score.items()))], read
def blockmax(terms, bonus=False):
    """Term at a time, rarest (highest-impact) term first. A later term's block is decoded only if its stored maximum, plus the
    best partial score any document in the block's id range already has, could beat the current tenth-best score."""
    order = sorted(terms, key=lambda t: len(post[t][0])); score, seen, read = {}, {}, 0
    for n, t in enumerate(order):
        d, w = post[t]; tenth = sorted(score.values(), reverse=True)[9] if len(score) >= 10 else -1
        remaining = sum(float(post[o][1].max()) for o in order[n + 1:])        # what terms still to come could add
        known = np.array(sorted(score)) if score else np.array([], dtype=np.int64)
        for lo in range(0, len(d), args.block):
            hi = min(lo + args.block, len(d)); bmax = float(w[lo:hi].max())
            a, b = np.searchsorted(known, d[lo]), np.searchsorted(known, d[hi - 1], side="right")
            best = max((score[k] for k in known[a:b].tolist()), default=0.0)
            if n > 0 and len(score) >= 10 and bmax + best + remaining <= tenth: continue
            read += hi - lo
            for doc, val in zip(d[lo:hi].tolist(), w[lo:hi].tolist()):
                score[doc] = score.get(doc, 0) + val; seen[doc] = seen.get(doc, 0) + 1
    if bonus:
        for doc in score:
            if doc % 97 == 0: score[doc] += 4                            # applied only to documents the reader decoded
    return [d for _, d in heapq.nlargest(10, ((v, d) for d, v in score.items()))], read
rows = []
for name, terms in (("rare + common (terms 0, 4)", [0, 4]), ("all five terms", [0, 1, 2, 3, 4])):
    ex, r_ex = exhaustive(terms); bm, r_bm = blockmax(terms)
    rows.append({"query": name, "score": "sum of per-term parts", "postings_exhaustive": r_ex, "postings_blockmax": r_bm, "read_pct": r_bm / r_ex * 100,
                 "same_top10": "yes" if set(ex) == set(bm) else "no"})
ex, r_ex = exhaustive([0, 4], bonus=True); bm, r_bm = blockmax([0, 4], bonus=True)
rows.append({"query": "rare + common (terms 0, 4)", "score": "plus a boost outside the bounds", "postings_exhaustive": r_ex, "postings_blockmax": r_bm, "read_pct": r_bm / r_ex * 100,
             "same_top10": "yes" if set(ex) == set(bm) else "no"})
for r in rows: print(f"{r['query']:<28} {r['score']:<28} exhaustive {r['postings_exhaustive']:>6,}  block-max {r['postings_blockmax']:>6,} ({r['read_pct']:5.1f}%)  same top 10: {r['same_top10']}")

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