Support boolean arrays, and list of array of indices for datasketch
from __future__ import absolute_import
+
+import numpy as np
+
from datasketch import MinHashLSHForest, MinHash
from ann_benchmarks.algorithms.base import BaseANN
@@ -17,6 +20,10 @@ class DataSketch(BaseANN):
self._index = MinHashLSHForest(num_perm=self._n_perm, l=self._n_rep)
for i, x in enumerate(X):
m = MinHash(num_perm=self._n_perm)
+ if x.dtype == np.bool_:
+ for e in np.flatnonzero(x):
+ m.update(str(e).encode('utf8'))
+ else:
for e in x:
m.update(str(e).encode('utf8'))
self._index.add(str(i), m)
@@ -24,6 +31,10 @@ class DataSketch(BaseANN):
def query(self, v, n):
m = MinHash(num_perm=self._n_perm)
+ if v.dtype == np.bool_:
+ for e in np.flatnonzero(v):
+ m.update(str(e).encode('utf8'))
+ else:
for e in v:
m.update(str(e).encode('utf8'))
return map(int, self._index.query(m, n))
