From ac2fb8f110e789873f5d5b6699375327fc2e7d7e Mon Sep 17 00:00:00 2001 From: Dipesh Babu Date: Sat, 18 Jul 2026 23:26:55 -0400 Subject: [PATCH] Enforce MinHash construction compatibility --- datasketch/aio/lsh.py | 19 +++--- datasketch/lsh.py | 36 +++++++----- datasketch/lsh_bloom.py | 16 +++-- datasketch/lshensemble.py | 4 +- datasketch/lshforest.py | 14 +++-- datasketch/minhash.py | 110 ++++++++++++++++++++++++++++++----- test/test_minhash.py | 32 ++++++++++ test/test_minhash_schemes.py | 44 ++++++++++++++ 8 files changed, 223 insertions(+), 52 deletions(-) diff --git a/datasketch/aio/lsh.py b/datasketch/aio/lsh.py index 634464dc..5e6e78d2 100644 --- a/datasketch/aio/lsh.py +++ b/datasketch/aio/lsh.py @@ -14,7 +14,7 @@ async_unordered_storage, ) from datasketch.lsh import _optimal_param -from datasketch.minhash import _check_scheme_consistency +from datasketch.minhash import _check_minhash_compatibility from datasketch.storage import _random_name, unordered_storage @@ -89,10 +89,10 @@ def __init__( self.hashranges = [(i * self.r, (i + 1) * self.r) for i in range(self.b)] self.hashtables = None self.keys = None - # The permutation scheme of the indexed MinHash, learned from the - # first insert. Note that an index attached to pre-existing external - # storage re-learns the scheme on its first insert. + # MinHash construction metadata learned from the first insert. An + # index attached to pre-existing external storage re-learns it. self._minhash_scheme: Optional[str] = None + self._minhash_config = None self._lock = asyncio.Lock() self._initialized = False @@ -253,7 +253,10 @@ async def main(): async def _insert(self, key, minhash, check_duplication=True, buffer=False): if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - self._minhash_scheme = _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + self._minhash_config = _check_minhash_compatibility(known, minhash) + if self._minhash_config is not None: + self._minhash_scheme = self._minhash_config[0] if self._require_bytes_keys and not isinstance(key, bytes): raise TypeError( f"prepickle=False requires bytes keys for non-dict storage, got {type(key).__name__}. " @@ -278,7 +281,8 @@ async def query(self, minhash): """See :class:`datasketch.MinHashLSH`.""" if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) fs = ( hashtable.get(self._H(minhash.hashvalues[start:end])) @@ -329,7 +333,8 @@ def _H(hs): async def _query_b(self, minhash, b): if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) if b > len(self.hashtables): raise ValueError("b must be less or equal to the number of hash tables") fs = [] diff --git a/datasketch/lsh.py b/datasketch/lsh.py index bdfff167..d99929ad 100644 --- a/datasketch/lsh.py +++ b/datasketch/lsh.py @@ -7,7 +7,7 @@ from scipy.integrate import quad as integrate -from datasketch.minhash import MinHash, _check_scheme_consistency +from datasketch.minhash import MinHash, _check_minhash_compatibility from datasketch.storage import ( OrderedStorage, UnorderedStorage, @@ -198,10 +198,10 @@ def __init__( ] self.hashranges = [(i * self.r, (i + 1) * self.r) for i in range(self.b)] self.keys: OrderedStorage = ordered_storage(storage_config, name=b"".join([basename, b"_keys"])) - # The permutation scheme of the indexed MinHash, learned from the - # first insert. Note that an index attached to pre-existing external - # storage (e.g. Redis) re-learns the scheme on its first insert. + # MinHash construction metadata learned from the first insert. An + # index attached to pre-existing external storage re-learns it. self._minhash_scheme: Optional[str] = None + self._minhash_config = None @property def buffer_size(self) -> int: @@ -336,7 +336,10 @@ def _insert( ): if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - self._minhash_scheme = _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + self._minhash_config = _check_minhash_compatibility(known, minhash) + if self._minhash_config is not None: + self._minhash_scheme = self._minhash_config[0] if self._require_bytes_keys and not isinstance(key, bytes): raise TypeError( f"prepickle=False requires bytes keys for non-dict storage, got {type(key).__name__}. " @@ -360,16 +363,14 @@ def __equivalent(self, other: MinHashLSH) -> bool: def _merge(self, other: MinHashLSH, check_overlap: bool = False, buffer: bool = False) -> None: if self.__equivalent(other): - known, other_known = getattr(self, "_minhash_scheme", None), getattr(other, "_minhash_scheme", None) - if known is not None and other_known is not None and known != other_known: - raise ValueError( - "Cannot merge MinHashLSH indexed with MinHash scheme %r into one indexed with scheme %r" - % (other_known, known) - ) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + other_known = getattr(other, "_minhash_config", getattr(other, "_minhash_scheme", None)) + merged_config = _check_minhash_compatibility(known, other_known) if check_overlap and set(self.keys).intersection(set(other.keys)): raise ValueError("The keys are overlapping, duplicate key exists.") - if known is None: - self._minhash_scheme = other_known + self._minhash_config = merged_config + if merged_config is not None: + self._minhash_scheme = merged_config[0] for key in other.keys: Hs = other.keys.get(key) self.keys.insert(key, *Hs, buffer=buffer) @@ -435,7 +436,8 @@ def query(self, minhash) -> list[Hashable]: """ if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) candidates = set() for (start, end), hashtable in zip(self.hashranges, self.hashtables): H = self._H(minhash.hashvalues[start:end]) @@ -462,7 +464,8 @@ def add_to_query_buffer(self, minhash: Union[MinHash, WeightedMinHash]) -> None: """ if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) for (start, end), hashtable in zip(self.hashranges, self.hashtables): H = self._H(minhash.hashvalues[start:end]) hashtable.add_to_select_buffer([H]) @@ -560,7 +563,8 @@ def _hashed_byteswap(self, hs): def _query_b(self, minhash, b): if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) if b > len(self.hashtables): raise ValueError("b must be less or equal to the number of hash tables") candidates = set() diff --git a/datasketch/lsh_bloom.py b/datasketch/lsh_bloom.py index 2add6272..c4d45871 100644 --- a/datasketch/lsh_bloom.py +++ b/datasketch/lsh_bloom.py @@ -8,7 +8,7 @@ import numpy as np from scipy.integrate import quad as integrate -from datasketch.minhash import MinHash, _check_scheme_consistency +from datasketch.minhash import MinHash, _check_minhash_compatibility try: import pybloomfilter @@ -298,10 +298,10 @@ def __init__( for i in range(self.b) ] self.hashranges = [(i * self.r, (i + 1) * self.r) for i in range(self.b)] - # The permutation scheme of the indexed MinHash, learned from the - # first insert. Note that an index restored from Bloom filter files - # re-learns the scheme on its first insert. + # MinHash construction metadata learned from the first insert. An + # index restored from Bloom filter files re-learns it. self._minhash_scheme: Optional[str] = None + self._minhash_config = None def insert(self, minhash: MinHash): """Insert the MinHash or Weighted MinHash @@ -316,7 +316,10 @@ def insert(self, minhash: MinHash): def _insert(self, minhash: MinHash): if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - self._minhash_scheme = _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + self._minhash_config = _check_minhash_compatibility(known, minhash) + if self._minhash_config is not None: + self._minhash_scheme = self._minhash_config[0] Hs = [minhash.hashvalues[start:end] for start, end in self.hashranges] @@ -371,7 +374,8 @@ def query(self, minhash) -> bool: """ if len(minhash) != self.h: raise ValueError("Expecting minhash with length %d, got %d" % (self.h, len(minhash))) - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) # if we match in any band, this is a candidate pair for (start, end), hashtable in zip(self.hashranges, self.hashtables): diff --git a/datasketch/lshensemble.py b/datasketch/lshensemble.py index b25cca4a..22304f58 100644 --- a/datasketch/lshensemble.py +++ b/datasketch/lshensemble.py @@ -245,8 +245,8 @@ def query(self, minhash: MinHash, size: int) -> Generator[Hashable, None, None]: if u is None: continue b, r = self._get_optimal_param(u, size) - # The permutation-scheme guard lives in the inner MinHashLSH - # (_query_b calls _check_scheme_consistency), so it is checked + # The MinHash compatibility guard lives in the inner MinHashLSH, + # so it is checked # lazily as this generator is iterated and only for partitions # that have had an insert. A query reaching only empty partitions # returns nothing regardless, so it cannot yield wrong results. diff --git a/datasketch/lshforest.py b/datasketch/lshforest.py index 71ff4ea9..29452480 100644 --- a/datasketch/lshforest.py +++ b/datasketch/lshforest.py @@ -3,7 +3,7 @@ import numpy as np -from datasketch.minhash import MinHash, _check_scheme_consistency +from datasketch.minhash import MinHash, _check_minhash_compatibility class MinHashLSHForest: @@ -42,8 +42,9 @@ def __init__(self, num_perm: int = 128, l: int = 8) -> None: self.keys = dict() # This is the sorted array implementation for the prefix trees self.sorted_hashtables = [[] for _ in range(self.l)] - # The permutation scheme of the indexed MinHash, learned on first add. + # MinHash construction metadata learned on first add. self._minhash_scheme = None + self._minhash_config = None def add(self, key: Hashable, minhash: MinHash) -> None: """Add a unique key, together @@ -60,8 +61,10 @@ def add(self, key: Hashable, minhash: MinHash) -> None: """ if len(minhash) < self.k * self.l: raise ValueError("The num_perm of MinHash out of range") - self._minhash_scheme = _check_scheme_consistency( - getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + self._minhash_config = _check_minhash_compatibility(known, minhash) + if self._minhash_config is not None: + self._minhash_scheme = self._minhash_config[0] if key in self.keys: raise ValueError("The given key has already been added") self.keys[key] = [self._H(minhash.hashvalues[start:end]) @@ -121,7 +124,8 @@ def query(self, minhash: MinHash, k: int) -> list[Hashable]: raise ValueError("k must be positive") if len(minhash) < self.k * self.l: raise ValueError("The num_perm of MinHash out of range") - _check_scheme_consistency(getattr(self, "_minhash_scheme", None), minhash) + known = getattr(self, "_minhash_config", getattr(self, "_minhash_scheme", None)) + _check_minhash_compatibility(known, minhash) results = set() r = self.k while r > 0: diff --git a/datasketch/minhash.py b/datasketch/minhash.py index d1870b48..15f4ec18 100644 --- a/datasketch/minhash.py +++ b/datasketch/minhash.py @@ -1,6 +1,8 @@ from __future__ import annotations import copy +import hashlib +import marshal import warnings from collections.abc import Generator, Iterable from typing import TYPE_CHECKING, Callable, Optional, Union @@ -108,26 +110,93 @@ def _fmix(hv, width: int): return hv ^ (hv >> s3) -def _check_scheme_consistency(known: Optional[str], minhash) -> Optional[str]: - """Check the permutation scheme of a MinHash given to an LSH index - against the scheme of the MinHash previously given to it, and return - the scheme the index should remember. +def _permutations_fingerprint(permutations) -> str: + """Return a compact fingerprint of a MinHash permutation matrix.""" + values = np.ascontiguousarray(permutations) + digest = hashlib.sha256() + digest.update(values.dtype.str.encode("ascii")) + digest.update(repr(values.shape).encode("ascii")) + digest.update(values.tobytes()) + return digest.hexdigest() - Hash values carry no trace of the scheme that produced them, so mixing - schemes in one index silently returns wrong results instead of failing; - this check makes it fail. Sketch types without a scheme attribute - (WeightedMinHash) are exempt. + +def _hashfunc_fingerprint(hashfunc: Callable) -> str: + """Return a stable-enough fingerprint for a hash callable. + + Python functions include their bytecode, defaults, and closure values; + callable objects additionally include their instance state. This is a + compatibility guard, not a serialization format or security boundary. """ + digest = hashlib.sha256() + callable_type = type(hashfunc) + parts = ( + getattr(hashfunc, "__module__", callable_type.__module__), + getattr(hashfunc, "__qualname__", callable_type.__qualname__), + ) + digest.update(repr(parts).encode("utf-8")) + code = getattr(hashfunc, "__code__", None) + if code is None and callable(hashfunc): + code = getattr(hashfunc.__call__, "__code__", None) + if code is not None: + digest.update(marshal.dumps(code)) + digest.update(repr(getattr(hashfunc, "__defaults__", None)).encode("utf-8")) + digest.update(repr(getattr(hashfunc, "__kwdefaults__", None)).encode("utf-8")) + closure = getattr(hashfunc, "__closure__", None) + if closure is not None: + digest.update(repr(tuple(cell.cell_contents for cell in closure)).encode("utf-8")) + digest.update(repr(getattr(hashfunc, "__dict__", None)).encode("utf-8")) + return digest.hexdigest() + + +def _minhash_compatibility(minhash): + """Return the comparable MinHash construction metadata, when available.""" scheme = getattr(minhash, "scheme", None) if scheme is None: + # WeightedMinHash has no permutation scheme and remains exempt. + return None + permutations = getattr(minhash, "permutations", None) + hashfunc = getattr(minhash, "hashfunc", None) + return ( + scheme, + _permutations_fingerprint(permutations) if permutations is not None else None, + _hashfunc_fingerprint(hashfunc) if hashfunc is not None else None, + ) + + +def _check_minhash_compatibility(known, minhash): + """Check MinHash construction metadata against metadata already known. + + LSH indexes use this to prevent mixing sketches made with different + permutation matrices or hash functions. A legacy string value is accepted + for indexes pickled before the richer compatibility metadata was added. + + Hash values carry no trace of the scheme that produced them, so mixing + incompatible sketches silently returns wrong results instead of failing. + Sketch types without a scheme attribute (WeightedMinHash) are exempt. + """ + if isinstance(known, str): + known = (known, None, None) + candidate = minhash if isinstance(minhash, tuple) and len(minhash) == 3 else _minhash_compatibility(minhash) + if candidate is None: return known if known is None: - return scheme - if scheme != known: + return candidate + known_scheme, known_permutations, known_hashfunc = known + scheme, permutations, hashfunc = candidate + if scheme != known_scheme: raise ValueError( - "MinHash scheme %r does not match scheme %r of the MinHash previously given to this index" % (scheme, known) + "MinHash scheme %r does not match scheme %r of the MinHash previously given to this index" + % (scheme, known_scheme) ) - return known + if known_permutations is not None and permutations is not None and permutations != known_permutations: + raise ValueError("MinHash permutations do not match those of the MinHash previously given to this index") + if known_hashfunc is not None and hashfunc is not None and hashfunc != known_hashfunc: + raise ValueError("MinHash hash function does not match that of the MinHash previously given to this index") + return ( + known_scheme, + known_permutations if known_permutations is not None else permutations, + known_hashfunc if known_hashfunc is not None else hashfunc, + ) _GPU_OK_CACHE: Optional[bool] = None @@ -550,7 +619,8 @@ def jaccard(self, other: MinHash) -> float: Raises: ValueError: If the two MinHashes have different numbers of permutation functions, different seeds, or different - permutation schemes. + permutation schemes, permutation parameters, or hash + functions. """ if getattr(other, "scheme", None) != self.scheme: @@ -565,6 +635,7 @@ def jaccard(self, other: MinHash) -> float: raise ValueError( "Cannot compute Jaccard given MinHash with different numbers of permutation functions" ) + _check_minhash_compatibility(_minhash_compatibility(self), other) return float(np.count_nonzero(self.hashvalues == other.hashvalues)) / float(len(self)) def count(self) -> float: @@ -588,7 +659,8 @@ def merge(self, other: MinHash) -> None: Raises: ValueError: If the two MinHashes have different numbers of permutation functions, different seeds, or different - permutation schemes. + permutation schemes, permutation parameters, or hash + functions. """ if getattr(other, "scheme", None) != self.scheme: @@ -603,6 +675,7 @@ def merge(self, other: MinHash) -> None: raise ValueError( "Cannot merge MinHash with different numbers of permutation functions" ) + _check_minhash_compatibility(_minhash_compatibility(self), other) self.hashvalues = np.minimum(other.hashvalues, self.hashvalues) def digest(self) -> np.ndarray: @@ -649,13 +722,14 @@ def __len__(self) -> int: def __eq__(self, other: MinHash) -> bool: """Returns: - bool: If their schemes, seeds and hash values are all equal then two are equivalent. + bool: If their construction metadata and hash values are all equal then two are equivalent. """ return ( type(self) is type(other) and self.scheme == other.scheme and self.seed == other.seed + and _minhash_compatibility(self) == _minhash_compatibility(other) and np.array_equal(self.hashvalues, other.hashvalues) ) @@ -674,7 +748,8 @@ def union(cls, *mhs: MinHash) -> MinHash: ValueError: If the number of MinHash objects passed as arguments is less than 2, or if the MinHash objects passed as arguments have different seeds, different numbers of permutation functions, or different - permutation schemes. + permutation schemes, permutation parameters, or hash + functions. Example: @@ -702,6 +777,9 @@ def union(cls, *mhs: MinHash) -> MinHash: raise ValueError( "The unioning MinHash must have the same seed, number of permutation functions and scheme" ) + compatibility = _minhash_compatibility(mhs[0]) + for minhash in mhs[1:]: + _check_minhash_compatibility(compatibility, minhash) hashvalues = np.minimum.reduce([m.hashvalues for m in mhs]) permutations = mhs[0].permutations return cls( diff --git a/test/test_minhash.py b/test/test_minhash.py index 61567f91..3681c1ef 100644 --- a/test/test_minhash.py +++ b/test/test_minhash.py @@ -8,6 +8,10 @@ from test.utils import fake_hash_func +def alternate_fake_hash_func(data): + return data + 1 + + class TestMinHash(unittest.TestCase): def test_init(self): m1 = minhash.MinHash(4, 1, hashfunc=fake_hash_func) @@ -66,6 +70,34 @@ def test_union(self): self.assertEqual(u._gpu_mode, "detect") u.update(13) + def test_operations_reject_different_hash_functions(self): + m1 = minhash.MinHash(4, 1, hashfunc=fake_hash_func) + m2 = minhash.MinHash(4, 1, hashfunc=alternate_fake_hash_func) + with self.assertRaisesRegex(ValueError, "hash function"): + m1.jaccard(m2) + with self.assertRaisesRegex(ValueError, "hash function"): + m1.merge(m2) + with self.assertRaisesRegex(ValueError, "hash function"): + minhash.MinHash.union(m1, m2) + + def test_operations_reject_different_explicit_permutations(self): + m1 = minhash.MinHash(4, 1, hashfunc=fake_hash_func) + permutations = m1.permutations.copy() + permutations[1, 0] += 1 + m2 = minhash.MinHash( + 4, + 1, + hashfunc=fake_hash_func, + permutations=permutations, + scheme=m1.scheme, + ) + with self.assertRaisesRegex(ValueError, "permutations"): + m1.jaccard(m2) + with self.assertRaisesRegex(ValueError, "permutations"): + m1.merge(m2) + with self.assertRaisesRegex(ValueError, "permutations"): + minhash.MinHash.union(m1, m2) + def test_pickle(self): m = minhash.MinHash(4, 1, hashfunc=fake_hash_func) m.update(123) diff --git a/test/test_minhash_schemes.py b/test/test_minhash_schemes.py index 4109e1d1..af322423 100644 --- a/test/test_minhash_schemes.py +++ b/test/test_minhash_schemes.py @@ -29,6 +29,10 @@ AFFINE_SCHEMES = ("affine32", "affine64") ALL_SCHEMES = ("affine32", "affine64", "legacy") + +def alternate_fake_hash_func(data): + return data + 1 + # -------------------------------------------------------------------------- # Frozen format vectors. These freeze the affine schemes' permutation # generation and hash computation: any change to these values is a breaking @@ -353,6 +357,46 @@ def test_lsh_pickle_preserves_learned_scheme(self): restored.query(self._mh("legacy")) self.assertIn("k", restored.query(self._mh("affine32"))) + def test_lsh_rejects_different_hash_function(self): + first = MinHash(128, hashfunc=fake_hash_func) + incompatible = MinHash(128, hashfunc=alternate_fake_hash_func) + lsh = MinHashLSH(threshold=0.5, num_perm=128) + lsh.insert("first", first) + with self.assertRaisesRegex(ValueError, "hash function"): + lsh.insert("second", incompatible) + with self.assertRaisesRegex(ValueError, "hash function"): + lsh.query(incompatible) + + other = MinHashLSH(threshold=0.5, num_perm=128) + other.insert("second", incompatible) + with self.assertRaisesRegex(ValueError, "hash function"): + lsh.merge(other) + + def test_lsh_and_forest_reject_different_permutations(self): + first = MinHash(128, hashfunc=fake_hash_func) + permutations = first.permutations.copy() + permutations[1, 0] += 1 + incompatible = MinHash( + 128, + seed=first.seed, + hashfunc=fake_hash_func, + permutations=permutations, + scheme=first.scheme, + ) + + lsh = MinHashLSH(threshold=0.5, num_perm=128) + lsh.insert("first", first) + with self.assertRaisesRegex(ValueError, "permutations"): + lsh.query(incompatible) + + forest = MinHashLSHForest(num_perm=128) + forest.add("first", first) + with self.assertRaisesRegex(ValueError, "permutations"): + forest.add("second", incompatible) + forest.index() + with self.assertRaisesRegex(ValueError, "permutations"): + forest.query(incompatible, 1) + def test_lshforest_mixed_scheme_raises(self): forest = MinHashLSHForest(num_perm=128) forest.add("k", self._mh("affine32"))