Used mask

Author: Pringled

vicinity/backends/annoy.py

@@ -130,5 +130,6 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query
         """Threshold the backend."""
         out: QueryResult = []
         for x, y in self.query(vectors, max_k):
-            out.append((x[y < threshold], y[y < threshold]))
+            mask = y < threshold
+            out.append((x[mask], y[mask]))
         return out

vicinity/backends/faiss.py

@@ -179,13 +179,15 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query
                 dist = D[start:end]
                 if self.arguments.metric == "cosine":
                     dist = 1 - dist
-                out.append((idx[dist < threshold], dist[dist < threshold]))
+                mask = dist < threshold
+                out.append((idx[mask], dist[mask]))
         else:
             distances, indices = self.index.search(vectors, max_k)
             for dist, idx in zip(distances, indices):
                 if self.arguments.metric == "cosine":
                     dist = 1 - dist
-                out.append((idx[dist < threshold], dist[dist < threshold]))
+                mask = dist < threshold
+                out.append((idx[mask], dist[mask]))
 
         return out
 

vicinity/backends/hnsw.py

@@ -108,6 +108,7 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query
         """Threshold the backend."""
         out: QueryResult = []
         for x, y in self.query(vectors, max_k):
-            out.append((x[y < threshold], y[y < threshold]))
+            mask = y < threshold
+            out.append((x[mask], y[mask]))
 
         return out

vicinity/backends/pynndescent.py

@@ -86,7 +86,8 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query
         indices, distances = self.index.query(normalized_vectors, k=max_k)
         out: QueryResult = []
         for idx, dist in zip(indices, distances):
-            out.append((idx[dist < threshold], dist[dist < threshold]))
+            mask = dist < threshold
+            out.append((idx[mask], dist[mask]))
         return out
 
     def save(self, base_path: Path) -> None:

vicinity/backends/voyager.py

@@ -98,7 +98,8 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query
         """Threshold the backend."""
         out: list[tuple[npt.NDArray, npt.NDArray]] = []
         for x, y in self.query(vectors, max_k):
-            out.append((x[y < threshold], y[y < threshold]))
+            mask = y < threshold
+            out.append((x[mask], y[mask]))
 
         return out