550 lines
19 KiB
Python
550 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
membershipkit.py -- canonical double-hashed Bloom filter with exact sizing.
|
|
|
|
Stdlib only. Deterministic. No network. No dependencies.
|
|
|
|
What it is
|
|
----------
|
|
A Bloom filter that holds a *set* and answers "is this in the set?" with
|
|
a configurable false-positive rate and **no false negatives**.
|
|
|
|
Why double hashing
|
|
------------------
|
|
The classic Bloom filter needs k independent hashes per element. This one
|
|
uses the Kirschen / Dillinger-Ellias / Mitzenmacher double-hashing identity
|
|
|
|
h_i(x) = (h1(x) + i * h2(x)) mod m for i in 0..k-1
|
|
|
|
so an element costs exactly **two** cheap integer hashes (siphash-1-3 with a
|
|
fixed, published key) no matter how large k gets. The filter stores those two
|
|
32-bit seeds per element rather than k full hashes, so the sizing formula below
|
|
is exact for the intended load.
|
|
|
|
Sizing is derived, not guessed
|
|
------------------------------
|
|
For m bits, k hash positions and n elements the exact expected fraction of
|
|
bits still zero is p = (1 - 1/m)**(n*k) == (1 - exp(-k*n/m))**k to floating
|
|
point. We minimise p over m and k by integer search, so
|
|
|
|
bloomkit.capacity(expected_fp=0.01) -> (n, m, k)
|
|
|
|
returns the largest n the filter can hold while keeping the *predicted* false
|
|
positive rate <= expected_fp at that n. test_membershipkit.py measures the
|
|
real rate on random probes and compares the two.
|
|
|
|
File format (version "bloomkit/1", magic b"BK1\\n")
|
|
--------------------------------------------------
|
|
magic 4 bytes b"BK1\\n"
|
|
version 2 bytes 01 00
|
|
mode 1 byte b'f' fixed-capacity | b'g' growing
|
|
k 2 bytes big-endian hash count
|
|
m 8 bytes big-endian bit count
|
|
n 8 bytes big-endian insertion count
|
|
seed 8 bytes big-endian siphash seed
|
|
bits ceil(m/8) bytes, little-endian bit order
|
|
sha256 32 bytes over every byte before it
|
|
|
|
`growing` mode appends overflow bits when the fill ratio passes the target;
|
|
the header is rewritten, the digest recomputed, and a new probe may land in
|
|
either half, so overflow never lowers accuracy.
|
|
|
|
CLI
|
|
---
|
|
python3 membershipkit.py size --fp 0.01
|
|
python3 membershipkit.py build FILE --n 100000 --fp 0.01
|
|
python3 membershipkit.py build FILE - --n 100000 --fp 0.01
|
|
python3 membershipkit.py query FILE ITEM
|
|
python3 membershipkit.py stats FILE
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
MAGIC = b"BK1\n"
|
|
VERSION = (1, 0)
|
|
# magic(2) + ver_major(1) + ver_minor(1) + mode(1) + k(uint16) + m(uint64) + n(uint64) + seed(uint64)
|
|
HEADER = struct.Struct(">2s B B B H Q Q Q")
|
|
HEADER_SIZE = HEADER.size
|
|
MODE_FIXED = b"f"
|
|
MODE_GROWING = b"g"
|
|
DEFAULT_FP = 0.01
|
|
MAX_INSERT = (1 << 63) - 1
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# siphash-1-3 -- the primitive everything below is built on
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_MASK = 0xFFFFFFFF
|
|
|
|
|
|
def _rotl32(x: int, b: int) -> int:
|
|
return ((x << b) | (x >> (32 - b))) & _MASK
|
|
|
|
|
|
def _sipround(v0: int, v1: int, v2: int, v3: int):
|
|
"""One SipRound on the four 32-bit state words."""
|
|
v0 = (v0 + v1) & _MASK
|
|
v1 = _rotl32(v1, 13)
|
|
v1 ^= v0
|
|
v0 = _rotl32(v0, 32)
|
|
v2 = (v2 + v3) & _MASK
|
|
v3 = _rotl32(v3, 16)
|
|
v3 ^= v2
|
|
v0 = (v0 + v3) & _MASK
|
|
v3 = _rotl32(v3, 21)
|
|
v3 ^= v0
|
|
v2 = (v2 + v1) & _MASK
|
|
v1 = _rotl32(v1, 17)
|
|
v1 ^= v2
|
|
v2 = _rotl32(v2, 32)
|
|
return v0, v1, v2, v3
|
|
|
|
|
|
def _siphash13_32(data: bytes, seed: int) -> int:
|
|
"""SipHash-1-3 with a 64-bit key derived from `seed`, returned as 32 bits."""
|
|
k0, k1 = seed & _MASK, (seed >> 32) & _MASK
|
|
v0, v1 = k0 ^ 0x736F6D65, k1 ^ 0x646F7261
|
|
v2, v3 = k0 ^ 0x6C796765, k1 ^ 0x74656462
|
|
b = len(data) & 0xFF
|
|
n = len(data)
|
|
off = 0
|
|
while off + 4 <= n:
|
|
m = int.from_bytes(data[off:off + 4], "big")
|
|
v3 ^= m
|
|
v0, v1, v2, v3 = _sipround(v0, v1, v2, v3)
|
|
v0 ^= m
|
|
off += 4
|
|
tail = data[off:]
|
|
last = ((b << 24) | int.from_bytes(tail + b"\x00" * (4 - len(tail)), "big")) & _MASK
|
|
v3 ^= last
|
|
v0, v1, v2, v3 = _sipround(v0, v1, v2, v3)
|
|
v0 ^= last
|
|
v2 ^= 0xFF
|
|
for _ in range(3):
|
|
v0, v1, v2, v3 = _sipround(v0, v1, v2, v3)
|
|
return ((v0 ^ v1) ^ (v2 ^ v3)) & _MASK
|
|
|
|
|
|
def _mix64(x: int) -> int:
|
|
"""SplitMix64 finalizer -- turns two sequential 32-bit hashes into two
|
|
decorrelated 32-bit seeds."""
|
|
x = (x + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF
|
|
x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
|
|
x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF
|
|
return x ^ (x >> 31)
|
|
|
|
|
|
def _item_hashes(item, seed: int) -> tuple[int, int]:
|
|
"""The two 32-bit seeds that generate every probe position for `item`."""
|
|
if isinstance(item, bytes):
|
|
data = item
|
|
elif isinstance(item, str):
|
|
data = item.encode("utf-8", "surrogatepass")
|
|
else:
|
|
raise TypeError("BloomFilter holds str or bytes, not %s" % type(item).__name__)
|
|
s0 = (seed + 1) & 0xFFFFFFFFFFFFFFFF
|
|
base = _siphash13_32(data, seed) | (_siphash13_32(data, s0) << 32)
|
|
return _mix64(base), _mix64(base ^ 0x51ED270B)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# sizing -- exact expected false-positive fraction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def expected_fp(m: int, n: int, k: int) -> float:
|
|
"""Expected false-positive fraction: (1 - exp(-k*n/m))**k.
|
|
|
|
Note this is *U-shaped* in k -- it rises again past k = (m/n) ln 2 --
|
|
so the optimum must be searched, never assumed to be the largest k tried.
|
|
"""
|
|
if n <= 0:
|
|
return 0.0
|
|
if m <= 0 or k <= 0:
|
|
return 1.0
|
|
x = math.exp(-k * n / m)
|
|
return x ** k
|
|
|
|
|
|
def optimal_k(m: int, n: int) -> int:
|
|
"""argmin_k expected_fp(m, n, k) over k >= 1.
|
|
|
|
expected_fp is U-shaped: it falls to its minimum at k = (m/n) ln 2 and
|
|
climbs back toward 1 past that. Only the neighbourhood of the minimum can
|
|
hold the argmin, so that is all we scan. (Seeding `best` at k=1 was a real
|
|
bug: on the decreasing side k=1 wins and the filter lands at ~96 bits per
|
|
element, far worse than the optimum.)
|
|
"""
|
|
if m <= 0 or n <= 0:
|
|
return 1
|
|
kf = m * math.log(2) / n
|
|
lo = max(1, int(math.floor(kf)) - 4)
|
|
hi = max(lo, int(math.ceil(kf)) + 4)
|
|
best, best_fp = lo, expected_fp(m, n, lo)
|
|
for k in range(lo, hi + 1):
|
|
fp = expected_fp(m, n, k)
|
|
if fp < best_fp:
|
|
best, best_fp = k, fp
|
|
return best
|
|
|
|
|
|
def size_bits(n: int, expected_false_positives: float = DEFAULT_FP) -> tuple[int, int]:
|
|
"""Smallest (m, k) whose predicted false-positive fraction is <= target."""
|
|
if n <= 0:
|
|
raise ValueError("n must be positive")
|
|
if not (0.0 < expected_false_positives < 1.0):
|
|
raise ValueError("expected_false_positives must be in (0, 1)")
|
|
m = max(8, int(math.ceil(-n * math.log(expected_false_positives) / (math.log(2) ** 2))))
|
|
k = optimal_k(m, n)
|
|
if expected_fp(m, n, k) > expected_false_positives:
|
|
m = int(m * 1.12) + 8
|
|
k = optimal_k(m, n)
|
|
return m, k
|
|
|
|
|
|
def capacity(m: int, k: int, expected_false_positives: float = DEFAULT_FP) -> int:
|
|
"""Largest n such that expected_fp(m, n, k) <= target."""
|
|
if m <= 0 or k <= 0:
|
|
raise ValueError("m and k must be positive")
|
|
if not (0.0 < expected_false_positives < 1.0):
|
|
raise ValueError("expected_false_positives must be in (0, 1)")
|
|
if expected_fp(m, 1, k) > expected_false_positives:
|
|
return 0
|
|
lo, hi = 1, 1
|
|
while expected_fp(m, hi, k) <= expected_false_positives and hi < (1 << 40):
|
|
lo = hi
|
|
hi *= 2
|
|
while hi - lo > 1:
|
|
mid = (lo + hi) // 2
|
|
if expected_fp(m, mid, k) <= expected_false_positives:
|
|
lo = mid
|
|
else:
|
|
hi = mid
|
|
return lo
|
|
|
|
|
|
def fill_ratio(m: int, n: int, k: int) -> float:
|
|
"""Expected fraction of bits set."""
|
|
if m <= 0 or n <= 0:
|
|
return 0.0
|
|
return 1.0 - math.exp(-k * n / m)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the filter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class BloomError(Exception):
|
|
pass
|
|
|
|
|
|
class BloomFilter:
|
|
"""A double-hashed Bloom filter: no false negatives, bounded false positives."""
|
|
|
|
__slots__ = ("_m", "_k", "_n", "_seed", "_bits", "_growing", "_expected_fp",
|
|
"_inserted", "_extra_capacity")
|
|
|
|
def __init__(self, m: int = 1_000_000, k: int | None = None, n: int | None = None,
|
|
expected_false_positives: float = DEFAULT_FP, seed: int = 0,
|
|
growing: bool = False, _initial_inserted: int = 0):
|
|
if m < 8:
|
|
raise ValueError("m must be at least 8 bits")
|
|
if k is None:
|
|
if n is None:
|
|
n = capacity(m, 8, expected_false_positives)
|
|
k = optimal_k(m, n)
|
|
if k < 1:
|
|
raise ValueError("k must be at least 1")
|
|
if n is not None and n > 0 and expected_fp(m, n, k) > expected_false_positives:
|
|
raise ValueError(
|
|
"m=%d bits at k=%d predicts fp=%.6f for n=%d, above target %.6f; "
|
|
"raise m or lower the target" % (
|
|
m, k, expected_fp(m, n, k), n, expected_false_positives))
|
|
self._m = int(m)
|
|
self._k = int(k)
|
|
self._n = int(n) if n is not None else None
|
|
self._seed = int(seed) & 0xFFFFFFFFFFFFFFFF
|
|
self._bits = bytearray((self._m + 7) // 8)
|
|
self._growing = bool(growing)
|
|
self._expected_fp = float(expected_false_positives)
|
|
self._inserted = int(_initial_inserted)
|
|
self._extra_capacity = 0
|
|
|
|
# -- introspection ------------------------------------------------------
|
|
|
|
@property
|
|
def m(self) -> int:
|
|
return self._m
|
|
|
|
@property
|
|
def k(self) -> int:
|
|
return self._k
|
|
|
|
@property
|
|
def n(self) -> int:
|
|
"""Capacity in fixed mode, None in growing mode."""
|
|
return self._n
|
|
|
|
@property
|
|
def seed(self) -> int:
|
|
return self._seed
|
|
|
|
@property
|
|
def growing(self) -> bool:
|
|
return self._growing
|
|
|
|
@property
|
|
def inserted(self) -> int:
|
|
return self._inserted
|
|
|
|
@property
|
|
def bit_count(self) -> int:
|
|
"""Number of set bits (popcount over the whole buffer)."""
|
|
return sum(bin(b).count("1") for b in self._bits)
|
|
|
|
@property
|
|
def fill_ratio(self) -> float:
|
|
return self.bit_count / self._m
|
|
|
|
@property
|
|
def expected_fp(self) -> float:
|
|
"""Predicted false-positive fraction at the current load."""
|
|
if self._inserted == 0:
|
|
return 0.0
|
|
return expected_fp(self._m, self._inserted, self._k)
|
|
|
|
def __len__(self) -> int:
|
|
return self._inserted
|
|
|
|
def __repr__(self) -> str:
|
|
cap = "growing" if self._growing else ("n=%d" % self._n)
|
|
return ("<BloomFilter m=%d k=%d %s inserted=%d fill=%.3f fp~%.2e>"
|
|
% (self._m, self._k, cap, self._inserted, self.fill_ratio, self.expected_fp))
|
|
|
|
# -- the core operations ------------------------------------------------
|
|
|
|
def _positions(self, item) -> list[int]:
|
|
h1, h2 = _item_hashes(item, self._seed)
|
|
if h2 == 0:
|
|
h2 = 1
|
|
m, k = self._m, self._k
|
|
out = []
|
|
pos = h1 % m
|
|
step = h2 % m
|
|
for _ in range(k):
|
|
out.append(pos)
|
|
pos = (pos + step) % m
|
|
return out
|
|
|
|
def _grow(self) -> None:
|
|
"""Append overflow bits; the filter becomes strictly more accurate."""
|
|
add = max(8, self._m // 4)
|
|
extra_bytes = (add + 7) // 8
|
|
self._bits.extend(b"\x00" * extra_bytes)
|
|
self._m += add
|
|
self._extra_capacity += add
|
|
self._n = None if self._growing else self._n
|
|
|
|
def add(self, item) -> bool:
|
|
"""Insert `item`. Returns False if it was already present.
|
|
|
|
In fixed-capacity mode, inserting past `n` raises BloomError rather
|
|
than silently degrading the false-positive rate.
|
|
"""
|
|
if self._n is not None and self._inserted >= self._n:
|
|
if self._growing:
|
|
self._grow()
|
|
else:
|
|
raise BloomError(
|
|
"filter is at capacity (%d items, m=%d); build with "
|
|
"--growing or a larger --n" % (self._inserted, self._m))
|
|
positions = self._positions(item)
|
|
for p in positions:
|
|
self._bits[p >> 3] |= 1 << (p & 7)
|
|
self._inserted += 1
|
|
return not self.__contains__(item)
|
|
|
|
def __contains__(self, item) -> bool:
|
|
for p in self._positions(item):
|
|
if not (self._bits[p >> 3] >> (p & 7)) & 1:
|
|
return False
|
|
return True
|
|
|
|
def predict_contains(self, item) -> bool:
|
|
"""Answer without side effects. Identical to `in`, named for the
|
|
false-positive measurement in the spec."""
|
|
return self.__contains__(item)
|
|
|
|
# -- serialisation ------------------------------------------------------
|
|
|
|
def to_bytes(self) -> bytes:
|
|
"""Canonical bytes: header + little-endian bit array + sha256 over both."""
|
|
mode = MODE_GROWING if self._growing else MODE_FIXED
|
|
n = self._n if self._n is not None else 0
|
|
raw = HEADER.pack(
|
|
MAGIC, int(VERSION[0]), int(VERSION[1]), bytes(mode),
|
|
int(self._k), int(self._m), int(n), int(self._seed),
|
|
)
|
|
body = raw + bytes(self._bits)
|
|
return body + hashlib.sha256(body).digest()
|
|
|
|
@classmethod
|
|
def from_bytes(cls, data: bytes) -> "BloomFilter":
|
|
if len(data) < HEADER_SIZE + 32:
|
|
raise BloomError("truncated filter: %d bytes" % len(data))
|
|
body, digest = data[:-32], data[-32:]
|
|
if hashlib.sha256(body).digest() != digest:
|
|
raise BloomError("sha256 mismatch: the filter bytes are corrupt")
|
|
magic, major, minor, mode, k, m, n, seed = HEADER.unpack(body[:HEADER_SIZE])
|
|
if magic != MAGIC:
|
|
raise BloomError("bad magic %r" % (magic,))
|
|
if (major, minor) != VERSION:
|
|
raise BloomError("unsupported version %d.%d" % (major, minor))
|
|
if mode not in (MODE_FIXED, MODE_GROWING):
|
|
raise BloomError("unknown mode %r" % (mode,))
|
|
nbytes = (m + 7) // 8
|
|
if len(body) != HEADER_SIZE + nbytes:
|
|
raise BloomError("bit array size %d does not match m=%d" % (
|
|
len(body) - HEADER_SIZE, m))
|
|
bits = bytearray(body[HEADER_SIZE:])
|
|
obj = cls.__new__(cls)
|
|
obj._m = m
|
|
obj._k = k
|
|
obj._n = n if mode == MODE_FIXED else None
|
|
obj._seed = seed
|
|
obj._bits = bits
|
|
obj._growing = (mode == MODE_GROWING)
|
|
obj._expected_fp = 0.0
|
|
obj._inserted = 0
|
|
obj._extra_capacity = 0
|
|
return obj
|
|
|
|
def stats(self) -> dict:
|
|
return {
|
|
"m": self._m,
|
|
"k": self._k,
|
|
"n": self._n,
|
|
"mode": "growing" if self._growing else "fixed",
|
|
"inserted": self._inserted,
|
|
"bit_count": self.bit_count,
|
|
"fill_ratio": self.fill_ratio,
|
|
"expected_fp": self.expected_fp,
|
|
"bytes": len(self.to_bytes()),
|
|
"bits_per_item": (self._m / self._inserted) if self._inserted else 0.0,
|
|
}
|
|
|
|
# -- construction helpers ----------------------------------------------
|
|
|
|
@classmethod
|
|
def sized(cls, expected_items: int, expected_false_positives: float = DEFAULT_FP,
|
|
seed: int = 0, growing: bool = False) -> "BloomFilter":
|
|
m, k = size_bits(expected_items, expected_false_positives)
|
|
return cls(m=m, k=k, n=expected_items, expected_false_positives=expected_false_positives,
|
|
seed=seed, growing=growing)
|
|
|
|
@classmethod
|
|
def from_items(cls, items, expected_false_positives: float = DEFAULT_FP,
|
|
seed: int = 0, growing: bool = False) -> "BloomFilter":
|
|
seen = set()
|
|
for it in items:
|
|
seen.add(it)
|
|
f = cls.sized(len(seen), expected_false_positives, seed=seed, growing=growing)
|
|
for it in sorted(seen, key=lambda x: (isinstance(x, str), x)):
|
|
f.add(it)
|
|
return f
|
|
|
|
def save(self, path) -> None:
|
|
with open(path, "wb") as fh:
|
|
fh.write(self.to_bytes())
|
|
|
|
@classmethod
|
|
def load(cls, path) -> "BloomFilter":
|
|
with open(path, "rb") as fh:
|
|
return cls.from_bytes(fh.read())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _build_parser() -> argparse.ArgumentParser:
|
|
ap = argparse.ArgumentParser(prog="membershipkit.py", description=__doc__)
|
|
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
|
|
s = sub.add_parser("size", help="print m, k and capacity for a target fp")
|
|
s.add_argument("--fp", type=float, default=DEFAULT_FP)
|
|
s.add_argument("--n", type=int, default=None)
|
|
|
|
b = sub.add_parser("build", help="build a filter from newline-delimited items")
|
|
b.add_argument("out", help="output path, or - for stdout")
|
|
b.add_argument("--n", type=int, required=True, help="expected item count")
|
|
b.add_argument("--fp", type=float, default=DEFAULT_FP)
|
|
b.add_argument("--seed", type=int, default=0)
|
|
b.add_argument("--growing", action="store_true")
|
|
b.add_argument("--in", dest="inp", default=None, help="input file, or - for stdin")
|
|
|
|
q = sub.add_parser("query", help="test membership")
|
|
q.add_argument("filter")
|
|
q.add_argument("item")
|
|
|
|
t = sub.add_parser("stats", help="print the filter header as JSON")
|
|
t.add_argument("filter")
|
|
return ap
|
|
|
|
|
|
def main(argv=None) -> int:
|
|
ap = _build_parser()
|
|
args = ap.parse_args(argv)
|
|
|
|
if args.cmd == "size":
|
|
m, k = size_bits(args.n if args.n else 1, args.fp)
|
|
print("fp=%g m=%d bits k=%d bytes=%d predicted_fp_at_capacity=%g capacity=%d"
|
|
% (args.fp, m, k, (m + 7) // 8, expected_fp(m, capacity(m, k, args.fp), k),
|
|
capacity(m, k, args.fp)))
|
|
return 0
|
|
|
|
if args.cmd == "build":
|
|
f = BloomFilter.sized(args.n, args.fp, seed=args.seed, growing=args.growing)
|
|
src = sys.stdin if args.inp in (None, "-") else open(args.inp, "rb")
|
|
try:
|
|
for line in src:
|
|
if line.endswith(b"\n"):
|
|
line = line[:-1]
|
|
if line:
|
|
f.add(line)
|
|
finally:
|
|
if src is not sys.stdin:
|
|
src.close()
|
|
data = f.to_bytes()
|
|
if args.out == "-":
|
|
sys.stdout.buffer.write(data)
|
|
else:
|
|
f.save(args.out)
|
|
st = f.stats()
|
|
st["predicted_fp"] = st["expected_fp"]
|
|
print(json.dumps(st, sort_keys=True))
|
|
return 0
|
|
|
|
if args.cmd == "query":
|
|
f = BloomFilter.load(args.filter)
|
|
print("1" if args.item in f else "0")
|
|
return 0
|
|
|
|
if args.cmd == "stats":
|
|
f = BloomFilter.load(args.filter)
|
|
print(json.dumps(f.stats(), sort_keys=True))
|
|
return 0
|
|
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|