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

Lookups for IDs that do not exist, sent straight to SQLite, then through a
Bloom filter loaded with every valid ID and sized for a 1% false-positive rate.
"""
import argparse, hashlib, json, math, os, random, sqlite3, tempfile, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--ids", type=int, default=1_000_000, help="Valid article IDs in the table.")
p.add_argument("--lookups", type=int, default=100_000, help="Lookups for IDs that do not exist.")
p.add_argument("--fp-rate", type=float, default=0.01, help="Target false-positive rate of the Bloom filter.")
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)

# Valid IDs are even numbers; absent IDs are odd, so they are certain to be missing
valid = [2 * i for i in range(args.ids)]
absent = [2 * random.randrange(args.ids * 4) + 1 for _ in range(args.lookups)]

path = os.path.join(tempfile.gettempdir(), "bloom-lab.db")
import atexit; atexit.register(lambda: os.path.exists(path) and os.remove(path))   # leave nothing behind
if os.path.exists(path): os.remove(path)
db = sqlite3.connect(path)
db.execute("CREATE TABLE articles (id INTEGER PRIMARY KEY, title TEXT)")
db.executemany("INSERT INTO articles VALUES (?, ?)", ((i, f"article {i}") for i in valid))
db.commit()
db_queries = 0
def lookup_db(i):
    global db_queries
    db_queries += 1
    return db.execute("SELECT title FROM articles WHERE id = ?", (i,)).fetchone()


class Bloom:
    """m bits, k hash functions from two halves of one digest (double hashing)."""
    def __init__(self, n, fp):
        self.m = math.ceil(-n * math.log(fp) / math.log(2) ** 2)   # ≈ 9.6 bits per key at 1%
        self.k = round(self.m / n * math.log(2))                     # ≈ 7 at 1%
        self.bits = bytearray((self.m + 7) // 8)
    def _idx(self, x):
        d = hashlib.blake2b(str(x).encode(), digest_size=16).digest()
        a, b = int.from_bytes(d[:8], "little"), int.from_bytes(d[8:], "little")
        return ((a + i * b) % self.m for i in range(self.k))
    def add(self, x):
        for j in self._idx(x): self.bits[j >> 3] |= 1 << (j & 7)
    def __contains__(self, x):
        return all(self.bits[j >> 3] >> (j & 7) & 1 for j in self._idx(x))


# 1. Absent lookups straight to the database
t = time.perf_counter()
for i in absent: lookup_db(i)
t_db = time.perf_counter() - t

# 2. A Bloom filter loaded with all valid IDs
bloom = Bloom(args.ids, args.fp_rate)
for i in valid: bloom.add(i)

# 3. The same absent lookups, filter first
db_queries = 0
t = time.perf_counter()
for i in absent:
    if i in bloom: lookup_db(i)
t_bloom = time.perf_counter() - t
passed = db_queries

# No false negatives: every valid ID in a sample is reported present; a new ID is found once added to both
sample = random.sample(valid, 100_000)
false_negatives = sum(1 for i in sample if i not in bloom)
new_id = 2 * args.ids
db.execute("INSERT INTO articles VALUES (?, ?)", (new_id, "a new article")); db.commit(); bloom.add(new_id)
found_new = new_id in bloom and lookup_db(new_id) is not None

print(f"Bloom filter: {bloom.m / args.ids:.1f} bits per key, k = {bloom.k}, {bloom.m / 8 / 1e6:.2f} MB")
print(f"absent lookups straight to the database: {args.lookups:,} queries in {t_db:.2f} s")
print(f"through the filter: {passed:,} reached the database ({passed / args.lookups:.2%}) in {t_bloom:.2f} s "
      f"-> {t_db / t_bloom:.0f}x less time")
print(f"false negatives among {len(sample):,} valid IDs: {false_negatives}   new ID found after insert: {found_new}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"summary": [{
            "bits_per_key": bloom.m / args.ids, "k": bloom.k, "filter_mb": bloom.m / 8 / 1e6,
            "db_seconds": t_db, "bloom_seconds": t_bloom, "speedup": t_db / t_bloom, "passed": passed,
            "passed_pct": passed / args.lookups * 100, "false_negatives": false_negatives,
            "found_new": "yes" if found_new else "no"}]}}, f)