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

A sorted file of fixed-size records and a sparse index with one entry per G
records, for G from 64 to 65,536: index size, and p50/p99 of random point
lookups that binary-search the index in memory and then read one granule
from disk with the page cache bypassed.
"""
import argparse, json, os, platform, random, tempfile, time
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--records", type=int, default=100_000_000, help="Records in the file (16 bytes each).")
p.add_argument("--granules", type=lambda s: [int(v) for v in s.split(",")], default=[64, 512, 8192, 65536], help="Records per index entry (G).")
p.add_argument("--lookups", type=int, default=20_000, help="Random point lookups per G.")
p.add_argument("--path", default=os.path.join(tempfile.gettempdir(), "sparse-index.bin"), help="The data file.")
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()
import atexit, shutil
atexit.register(lambda: shutil.rmtree(args.path, ignore_errors=True) if os.path.isdir(args.path) else os.path.exists(args.path) and os.remove(args.path))   # leave nothing behind
random.seed(args.seed)
REC = 16
F_NOCACHE = 48

if not os.path.exists(args.path) or os.path.getsize(args.path) != args.records * REC:
    with open(args.path, "wb") as f:
        step = 10_000_000
        for s in range(0, args.records, step):
            keys = np.arange(s, min(s + step, args.records), dtype=np.int64) * 3       # sorted keys, gaps of 3
            f.write(np.stack([keys, keys ^ 0x5555], axis=1).tobytes())

fd = os.open(args.path, os.O_RDONLY)
if platform.system() == "Darwin":
    import fcntl; fcntl.fcntl(fd, F_NOCACHE, 1)
else:
    os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)

rows = []
for G in args.granules:
    first_keys = np.arange(0, args.records, G, dtype=np.int64) * 3               # the sparse index: first key of each granule
    lat = []
    for _ in range(args.lookups):
        key = random.randrange(args.records) * 3
        t = time.perf_counter()
        g = int(np.searchsorted(first_keys, key, side="right")) - 1
        data = os.pread(fd, G * REC, g * G * REC)
        recs = np.frombuffer(data, dtype=np.int64).reshape(-1, 2)
        found = int(np.searchsorted(recs[:, 0], key))
        lat.append(time.perf_counter() - t)
    lat.sort()
    rows.append({"G": G, "index_kb": first_keys.nbytes * 2 / 1e3, "granule_kb": G * REC / 1e3,
                 "p50_us": lat[len(lat) // 2] * 1e6, "p99_us": lat[int(len(lat) * 0.99)] * 1e6})
    r = rows[-1]
    print(f"G = {G:>6,}: index {r['index_kb']:>10,.0f} KB, granule {r['granule_kb']:>7,.1f} KB   lookup p50 {r['p50_us']:7.0f} µs  p99 {r['p99_us']:7.0f} µs")
os.close(fd)

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