Enforce MinHash construction compatibility

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 = []

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()

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):

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.

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: