Monkey-style filter allocation

# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Monkey-style filter allocation.

A 3-level LSM (10,000, 100,000 and 1,000,000 keys) with one Bloom filter per
level and one total memory budget, split two ways: equal bits per key at
every level, or false-positive rates proportional to level size (about 4.8
bits per key fewer for each larger level). Wasted I/Os for lookups of
missing keys, then the same with a fourth level. Computed from the Bloom
false-positive formula, and checked by simulation.
"""
import argparse, json, math, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--bits-per-key", type=float, default=10, help="Average bits per key over all levels (the budget).")
p.add_argument("--lookups", type=int, default=1_000_000, help="Lookups of missing keys.")
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()
random.seed(args.seed)
LN2SQ = math.log(2) ** 2
fpr = lambda bits_per_key: math.exp(-bits_per_key * LN2SQ)

def allocate(levels, mode):
    total = args.bits_per_key * sum(levels)
    if mode == "equal":
        return [args.bits_per_key] * len(levels)
    # Monkey: FPR_i proportional to n_i; solve for the constant c with sum n_i * bits(c * n_i) = total
    lo, hi = 1e-12, 1.0
    for _ in range(100):
        c = math.sqrt(lo * hi)
        used = sum(n * (-math.log(min(1.0, c * n / max(levels))) / LN2SQ) for n in levels)
        lo, hi = (lo, c) if used < total else (c, hi)
    return [-math.log(min(1.0, c * n / max(levels))) / LN2SQ for n in levels]

rows = []
for levels in ([10_000, 100_000, 1_000_000], [10_000, 100_000, 1_000_000, 10_000_000]):
    for mode in ("equal", "monkey"):
        bits = allocate(levels, mode)
        expected = sum(fpr(b) for b in bits)                      # wasted I/Os per missing-key lookup
        simulated = sum(sum(random.random() < fpr(b) for b in bits) for _ in range(args.lookups // 10)) / (args.lookups // 10)
        rows.append({"levels": len(levels), "allocation": mode, "bits": ", ".join(f"{b:.1f}" for b in bits),
                     "wasted_per_lookup": expected, "simulated": simulated, "wasted_per_million": expected * 1e6})
        print(f"{len(levels)} levels, {mode:<7} bits/key per level [{rows[-1]['bits']}]  "
              f"wasted I/Os per million missing-key lookups {expected * 1e6:9,.0f} (simulated {simulated * 1e6:9,.0f})")

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