beam: add bloom filter with optimal sizing and double hashing

wizard-work/beam/bloom_filter.py

@@ -0,0 +1,102 @@
+"""Bloom filter with optimal sizing and double hashing.
+
+A probabilistic set membership structure: false negatives never occur,
+false positives occur with a configurable rate. Dependency-free.
+"""
+import hashlib
+import math
+import struct
+from typing import Any
+
+
+def _to_bytes(item: Any) -> bytes:
+    if isinstance(item, bytes):
+        return item
+    if isinstance(item, str):
+        return item.encode("utf-8")
+    return repr(item).encode("utf-8")
+
+
+def _hashes(item: Any):
+    b = _to_bytes(item)
+    d1 = hashlib.md5(b).digest()
+    d2 = hashlib.sha1(b).digest()
+    h1 = int.from_bytes(d1[:8], "little")
+    h2 = int.from_bytes(d2[:8], "little")
+    if h2 == 0:
+        h2 = 0x9e3779b97f4a7c15
+    return h1, h2
+
+
+def optimal_params(n: int, p: float):
+    if n <= 0:
+        raise ValueError("n must be positive")
+    if not (0.0 < p < 1.0):
+        raise ValueError("p must be in (0,1)")
+    m = - (n * math.log(p)) / (math.log(2) ** 2)
+    k = (m / n) * math.log(2)
+    m = max(1, int(math.ceil(m)))
+    k = max(1, int(round(k)))
+    return m, k
+
+
+class BloomFilter:
+    def __init__(self, m: int, k: int):
+        if not isinstance(m, int) or m < 1:
+            raise ValueError("m must be a positive int")
+        if not isinstance(k, int) or k < 1:
+            raise ValueError("k must be a positive int")
+        self.m = m
+        self.k = k
+        self._bits = bytearray((m + 7) // 8)
+        self._count = 0
+
+    @classmethod
+    def create(cls, n: int, p: float) -> "BloomFilter":
+        m, k = optimal_params(n, p)
+        return cls(m, k)
+
+    def _positions(self, item: Any):
+        h1, h2 = _hashes(item)
+        for i in range(self.k):
+            yield (h1 + i * h2) % self.m
+
+    def add(self, item: Any) -> None:
+        for pos in self._positions(item):
+            self._bits[pos >> 3] |= 1 << (pos & 7)
+        self._count += 1
+
+    def __contains__(self, item: Any) -> bool:
+        return all(
+            (self._bits[pos >> 3] >> (pos & 7)) & 1
+            for pos in self._positions(item)
+        )
+
+    def __len__(self) -> int:
+        return self._count
+
+    def clear(self) -> None:
+        for i in range(len(self._bits)):
+            self._bits[i] = 0
+        self._count = 0
+
+    def estimated_fpr(self) -> float:
+        if self._count == 0:
+            return 0.0
+        return (1.0 - math.exp(-self.k * self._count / self.m)) ** self.k
+
+    def to_bytes(self) -> bytes:
+        header = struct.pack("!II", self.m, self.k)
+        return header + bytes(self._bits)
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "BloomFilter":
+        if len(data) < 8:
+            raise ValueError("truncated bloom filter")
+        m, k = struct.unpack("!II", data[:8])
+        expected = (m + 7) // 8
+        if len(data) - 8 != expected:
+            raise ValueError("bit array length mismatch")
+        bf = cls(m, k)
+        bf._bits = bytearray(data[8:])
+        return bf