Retrieve MinHash from LSHForest

datasketch/lshforest.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 from typing import Hashable, List
+import numpy as np
 
 from datasketch.minhash import MinHash
 
@@ -128,6 +129,29 @@ def query(self, minhash: MinHash, k: int) -> List[Hashable]:
             r -= 1
         return list(results)
 
+    def get_minhash_from_key(self, key: Hashable) -> MinHash:
+        """
+        Returns the MinHash value that corresponds to the given key in the LSHForest,
+        if it exists. This is useful for when we want to manually check the 
+        Jaccard Similarity for the top-k results from a query.
+
+        Args:
+            key (Hashable): The key whose MinHash we want to retrieve.
+
+        Returns:
+            MinHash: The corresponding MinHash value for the provided key.
+        """
+        byteslist = self.keys.get(key, None)
+        if byteslist is None:
+            raise KeyError(f"The provided key does not exist in the LSHForest: {key}")
+        hashvalues = np.array([], dtype=np.uint64)
+        for item in byteslist:
+            # unswap the bytes, as their representation is flipped during storage
+            hv_segment = np.frombuffer(item, dtype=np.uint64).byteswap()
+            hashvalues = np.append(hashvalues, hv_segment)
+        minhash = MinHash(hashvalues=hashvalues)
+        return minhash
+
     def _binary_search(self, n, func):
         """
         https://golang.org/src/sort/search.go?s=2247:2287#L49

test/test_lshforest.py

@@ -62,6 +62,13 @@ def test_query(self):
             results = forest.query(data[key], 10)
             self.assertIn(key, results)
 
+    def test_get_minhash(self):
+        forest, data = self._setup()
+        for key in data:
+            minhash_ori = data[key]
+            minhash_retrieved = forest.get_minhash_from_key(key)
+            self.assertEqual(minhash_retrieved.jaccard(minhash_ori), 1.0)
+
     def test_pickle(self):
         forest, _ = self._setup()
         forest2 = pickle.loads(pickle.dumps(forest))