# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Inverted index intersection: order and skips.

An inverted index over a million short documents with a Zipf-like word
distribution. Queries mixing a rare and a common term, intersected in query
order (common first) and smallest list first, with and without skips every
128 entries.
"""
import argparse, bisect, json, time
from collections import defaultdict
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--docs", type=int, default=1_000_000, help="Documents.")
p.add_argument("--words-per-doc", type=int, default=8, help="Words per document.")
p.add_argument("--vocab", type=int, default=50_000, help="Vocabulary size.")
p.add_argument("--queries", type=int, default=200, help="Queries.")
p.add_argument("--skip", type=int, default=128, help="Skip step.")
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)
w = 1 / np.arange(1, args.vocab + 1) ** 1.1; w /= w.sum()
words = rng.choice(args.vocab, size=(args.docs, args.words_per_doc), p=w)
postings = defaultdict(list)
for d, ws in enumerate(words.tolist()):
    for t in set(ws): postings[t].append(d)

def intersect(a, b, skip=0):
    i = j = 0; out = []; steps = 0
    while i < len(a) and j < len(b):
        steps += 1
        if a[i] == b[j]: out.append(a[i]); i += 1; j += 1
        elif a[i] < b[j]:
            if skip and i + skip < len(a) and a[i + skip] <= b[j]: i += skip
            else: i += 1
        else:
            if skip and j + skip < len(b) and b[j + skip] <= a[i]: j += skip
            else: j += 1
    return out, steps

common = [t for t in range(5)]                                     # the most frequent terms
rare = [t for t in range(2000, 40000) if 20 <= len(postings[t]) <= 200][: args.queries]
qs = [(common[i % 5], rare[i]) for i in range(len(rare))]
rows = []
for label, order, skip in (("query order (common term first)", "query", 0), ("smallest list first", "small", 0),
                           ("query order + skips", "query", args.skip), ("smallest first + skips", "small", args.skip)):
    t = time.perf_counter(); steps = 0
    for c, r in qs:
        a, b = postings[c], postings[r]
        if order == "small": a, b = (a, b) if len(a) <= len(b) else (b, a)
        _, s = intersect(a, b, skip); steps += s
    dt = (time.perf_counter() - t) / len(qs)
    rows.append({"strategy": label, "ms_per_query": dt * 1000, "steps_per_query": steps / len(qs)})
    print(f"{label:<34} {dt * 1000:8.2f} ms per query   {steps / len(qs):10,.0f} comparison steps")
print(f"common term lists: {len(postings[0]):,} docs; rare term lists: {min(len(postings[r]) for r in rare)}-{max(len(postings[r]) for r in rare)} docs")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"strategies": rows, "summary": [{"common_len": len(postings[0])}]}}, f)