# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Elias-Fano and the posting list.

1,000,000 sorted random doc ids from a universe of 10^9: size as raw int32,
as delta + varint, and as Elias-Fano; then 10^6 random-access reads of
element i from each (varint decodes from a checkpoint every 128 values;
Elias-Fano answers with a select over its high bits).
"""
import argparse, json, math, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--n", type=int, default=1_000_000, help="Doc ids.")
p.add_argument("--universe", type=int, default=10**9, help="Universe of doc ids.")
p.add_argument("--reads", type=int, default=1_000_000, help="Random-access reads (the varint path uses a tenth, then scales).")
p.add_argument("--checkpoint", type=int, default=128, help="Values between varint checkpoints.")
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)
ids = np.unique(rng.integers(0, args.universe, int(args.n * 1.01)))[: args.n].astype(np.int64)
n, U = len(ids), args.universe

# delta + varint, with checkpoints
gaps = np.diff(ids, prepend=0)
def varint(x):
    out = bytearray()
    while True:
        b = x & 0x7F; x >>= 7
        if x: out.append(b | 0x80)
        else: out.append(b); return out
blob, offsets = bytearray(), []
for i, g in enumerate(gaps.tolist()):
    if i % args.checkpoint == 0: offsets.append(len(blob)); g = int(ids[i])     # restart from the absolute value
    blob += varint(g)

# Elias-Fano
l = max(0, int(math.floor(math.log2(U / n))))
low = ids & ((1 << l) - 1)
high = ids >> l
hb = np.zeros(n + (U >> l) + 1, dtype=bool); hb[high + np.arange(n)] = True
ef_bits = n * l + len(hb)
ones = np.flatnonzero(hb)                                            # select(i) = position of the i-th one
def ef_get(i): return ((int(ones[i]) - i) << l) | int(low[i])

def varint_get(i):
    c = i // args.checkpoint; pos = offsets[c]; val = 0
    for j in range(c * args.checkpoint, i + 1):
        x, shift = 0, 0
        while True:
            b = blob[pos]; pos += 1; x |= (b & 0x7F) << shift; shift += 7
            if b < 0x80: break
        val = x if j == c * args.checkpoint else val + x
    return val

q = rng.integers(0, n, args.reads)
t = time.perf_counter(); r_ef = [ef_get(int(i)) for i in q]; t_ef = (time.perf_counter() - t) / len(q)
qv = q[: len(q) // 10]
t = time.perf_counter(); r_v = [varint_get(int(i)) for i in qv]; t_v = (time.perf_counter() - t) / len(qv)
t = time.perf_counter(); r_raw = [int(ids[i]) for i in q]; t_raw = (time.perf_counter() - t) / len(q)
assert r_ef[: len(qv)] == r_v == r_raw[: len(qv)]
rows = [{"format": "raw int32", "mb": n * 4 / 1e6, "ns_per_get": t_raw * 1e9},
        {"format": "delta + varint", "mb": len(blob) / 1e6, "ns_per_get": t_v * 1e9},
        {"format": "Elias-Fano", "mb": ef_bits / 8 / 1e6, "ns_per_get": t_ef * 1e9}]
for r in rows: print(f"{r['format']:<16} {r['mb']:6.2f} MB   get(i): {r['ns_per_get']:9.0f} ns")
print(f"Elias-Fano: {l} low bits per id + {len(hb) / n:.2f} high bits per id = {ef_bits / n:.2f} bits per id")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"formats": rows, "summary": [{"low_bits": l, "bits_per_id": ef_bits / n}]}}, f)