vicinity 0.3.3__tar.gz → 0.4.0__tar.gz

{vicinity-0.3.3 → vicinity-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vicinity
-Version: 0.3.3
+Version: 0.4.0
 Summary: Lightweight Nearest Neighbors with Flexible Backends
 Author-email: Stéphan Tulkens <stephantul@gmail.com>, Thomas van Dongen <thomas123@live.nl>
 License: MIT License

{vicinity-0.3.3 → vicinity-0.4.0}/tests/test_vicinity.py

@@ -58,6 +58,10 @@ def test_vicinity_query(vicinity_instance: Vicinity, query_vector: np.ndarray) -
 
     assert len(results) == 1
 
+    results = vicinity_instance.query(np.stack([query_vector, query_vector]), k=2)
+
+    assert results[0] == results[1]
+
 
 def test_vicinity_query_threshold(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:
     """
@@ -70,6 +74,10 @@ def test_vicinity_query_threshold(vicinity_instance: Vicinity, query_vector: np.
 
     assert len(results) >= 1
 
+    results = vicinity_instance.query_threshold(np.stack([query_vector, query_vector]), threshold=0.7)
+
+    assert results[0] == results[1]
+
 
 def test_vicinity_insert(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:
     """

{vicinity-0.3.3 → vicinity-0.4.0}/uv.lock

@@ -1087,7 +1087,7 @@ wheels = [
 
 [[package]]
 name = "vicinity"
-version = "0.3.1"
+version = "0.3.4"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/annoy.py

@@ -126,9 +126,10 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
         """Delete vectors from the backend."""
         raise NotImplementedError("Deletion is not supported in Annoy backend.")
 
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
         """Threshold the backend."""
-        out: list[npt.NDArray] = []
-        for x, y in self.query(vectors, 100):
-            out.append(x[y < threshold])
+        out: QueryResult = []
+        for x, y in self.query(vectors, max_k):
+            mask = y < threshold
+            out.append((x[mask], y[mask]))
         return out

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/base.py

@@ -92,7 +92,7 @@ class AbstractBackend(ABC, Generic[ArgType]):
         raise NotImplementedError()
 
     @abstractmethod
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
         """Threshold the backend."""
         raise NotImplementedError()
 

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/basic.py

@@ -150,22 +150,25 @@ class BasicBackend(BasicVectorStore, AbstractBackend[BasicArgs], ABC):
         self,
         vectors: npt.NDArray,
         threshold: float,
-    ) -> list[npt.NDArray]:
+        max_k: int,
+    ) -> QueryResult:
         """
         Batched distance thresholding.
 
         :param vectors: The vectors to threshold.
         :param threshold: The threshold to use.
-        :return: A list of lists of indices of vectors that are below the threshold
+        :param max_k: The maximum number of neighbors to consider.
+        :return: A list of tuples with the indices and distances.
         """
-        out: list[npt.NDArray] = []
+        out: QueryResult = []
         for i in range(0, len(vectors), 1024):
             batch = vectors[i : i + 1024]
             distances = self._dist(batch)
             for dists in distances:
-                indices = np.flatnonzero(dists <= threshold)
-                sorted_indices = indices[np.argsort(dists[indices])]
-                out.append(sorted_indices)
+                mask = dists <= threshold
+                indices = np.flatnonzero(mask)
+                filtered_distances = dists[mask]
+                out.append((indices, filtered_distances))
         return out
 
     def query(

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/faiss.py

@@ -146,6 +146,7 @@ class FaissBackend(AbstractBackend[FaissArgs]):
 
     def query(self, vectors: npt.NDArray, k: int) -> QueryResult:
         """Perform a k-NN search in the FAISS index."""
+        k = min(len(self), k)
         if self.arguments.metric == "cosine":
             vectors = normalize(vectors)
         distances, indices = self.index.search(vectors, k)
@@ -163,9 +164,9 @@ class FaissBackend(AbstractBackend[FaissArgs]):
         """Delete vectors from the backend."""
         raise NotImplementedError("Deletion is not supported in FAISS backends.")
 
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
         """Query vectors within a distance threshold, using range_search if supported."""
-        out: list[npt.NDArray] = []
+        out: QueryResult = []
         if self.arguments.metric == "cosine":
             vectors = normalize(vectors)
 
@@ -178,13 +179,15 @@ class FaissBackend(AbstractBackend[FaissArgs]):
                 dist = D[start:end]
                 if self.arguments.metric == "cosine":
                     dist = 1 - dist
-                out.append(idx[dist < threshold])
+                mask = dist < threshold
+                out.append((idx[mask], dist[mask]))
         else:
-            distances, indices = self.index.search(vectors, 100)
+            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])
+                mask = dist < threshold
+                out.append((idx[mask], dist[mask]))
 
         return out
 

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/hnsw.py

@@ -93,6 +93,7 @@ class HNSWBackend(AbstractBackend[HNSWArgs]):
 
     def query(self, vectors: npt.NDArray, k: int) -> QueryResult:
         """Query the backend."""
+        k = min(k, len(self))
         return list(zip(*self.index.knn_query(vectors, k)))
 
     def insert(self, vectors: npt.NDArray) -> None:
@@ -103,10 +104,11 @@ class HNSWBackend(AbstractBackend[HNSWArgs]):
         """Delete vectors from the backend."""
         raise NotImplementedError("Deletion is not supported in HNSW backend.")
 
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
         """Threshold the backend."""
-        out: list[npt.NDArray] = []
-        for x, y in self.query(vectors, 100):
-            out.append(x[y < threshold])
+        out: QueryResult = []
+        for x, y in self.query(vectors, max_k):
+            mask = y < threshold
+            out.append((x[mask], y[mask]))
 
         return out

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/pynndescent.py

@@ -80,15 +80,15 @@ class PyNNDescentBackend(AbstractBackend[PyNNDescentArgs]):
         """Delete vectors from the backend."""
         raise NotImplementedError("Deletion is not supported in PyNNDescent backend.")
 
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
         """Find neighbors within a distance threshold."""
         normalized_vectors = normalize_or_copy(vectors)
-        indices, distances = self.index.query(normalized_vectors, k=100)
-        result = []
+        indices, distances = self.index.query(normalized_vectors, k=max_k)
+        out: QueryResult = []
         for idx, dist in zip(indices, distances):
-            within_threshold = idx[dist < threshold]
-            result.append(within_threshold)
-        return result
+            mask = dist < threshold
+            out.append((idx[mask], dist[mask]))
+        return out
 
     def save(self, base_path: Path) -> None:
         """Save the vectors and configuration to a specified path."""

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/usearch.py

@@ -115,9 +115,10 @@ class UsearchBackend(AbstractBackend[UsearchArgs]):
 
     def query(self, vectors: npt.NDArray, k: int) -> QueryResult:
         """Query the backend and return results as tuples of keys and distances."""
+        k = min(k, len(self))
         results = self.index.search(vectors, k)
-        keys = np.array(results.keys).reshape(-1, k)
-        distances = np.array(results.distances, dtype=np.float32).reshape(-1, k)
+        keys = np.atleast_2d(results.keys)
+        distances = np.atleast_2d(results.distances)
         return list(zip(keys, distances))
 
     def insert(self, vectors: npt.NDArray) -> None:
@@ -128,9 +129,12 @@ class UsearchBackend(AbstractBackend[UsearchArgs]):
         """Delete vectors from the index (not supported by Usearch)."""
         raise NotImplementedError("Dynamic deletion is not supported in Usearch.")
 
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
-        """Threshold the backend and return filtered keys."""
-        return [
-            np.array(keys_row)[np.array(distances_row, dtype=np.float32) < threshold]
-            for keys_row, distances_row in self.query(vectors, 100)
-        ]
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
+        """Query vectors within a distance threshold and return keys and distances."""
+        out: QueryResult = []
+        for keys_row, distances_row in self.query(vectors, max_k):
+            keys_row = np.array(keys_row)
+            distances_row = np.array(distances_row, dtype=np.float32)
+            mask = distances_row < threshold
+            out.append((keys_row[mask], distances_row[mask]))
+        return out

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/backends/voyager.py

@@ -68,6 +68,7 @@ class VoyagerBackend(AbstractBackend[VoyagerArgs]):
 
     def query(self, query: npt.NDArray, k: int) -> QueryResult:
         """Query the backend for the nearest neighbors."""
+        k = min(k, len(self))
         indices, distances = self.index.query(query, k)
         return list(zip(indices, distances))
 
@@ -93,11 +94,12 @@ class VoyagerBackend(AbstractBackend[VoyagerArgs]):
         """Delete vectors from the backend."""
         raise NotImplementedError("Deletion is not supported in Voyager backend.")
 
-    def threshold(self, vectors: npt.NDArray, threshold: float) -> list[npt.NDArray]:
+    def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult:
         """Threshold the backend."""
-        out: list[npt.NDArray] = []
-        for x, y in self.query(vectors, len(self)):
-            out.append(x[y < threshold])
+        out: list[tuple[npt.NDArray, npt.NDArray]] = []
+        for x, y in self.query(vectors, max_k):
+            mask = y < threshold
+            out.append((x[mask], y[mask]))
 
         return out
 

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/version.py

@@ -1,2 +1,2 @@
-__version_triple__ = (0, 3, 3)
+__version_triple__ = (0, 4, 0)
 __version__ = ".".join(map(str, __version_triple__))

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity/vicinity.py

@@ -14,7 +14,7 @@ from numpy import typing as npt
 
 from vicinity import Metric
 from vicinity.backends import AbstractBackend, BasicBackend, BasicVectorStore, get_backend_class
-from vicinity.datatypes import Backend, PathLike
+from vicinity.datatypes import Backend, PathLike, QueryResult
 
 logger = logging.getLogger(__name__)
 
@@ -114,7 +114,7 @@ class Vicinity:
         self,
         vectors: npt.NDArray,
         k: int = 10,
-    ) -> list[list[tuple[str, float]]]:
+    ) -> list[QueryResult]:
         """
         Find the nearest neighbors to some arbitrary vector.
 
@@ -140,22 +140,26 @@ class Vicinity:
         self,
         vectors: npt.NDArray,
         threshold: float = 0.5,
-    ) -> list[list[str]]:
+        max_k: int = 100,
+    ) -> list[QueryResult]:
         """
-        Find the nearest neighbors to some arbitrary vector with some threshold.
+        Find the nearest neighbors to some arbitrary vector with some threshold. Note: the output is not sorted.
 
         :param vectors: The vectors to find the most similar vectors to.
         :param threshold: The threshold to use.
+        :param max_k: The maximum number of neighbors to consider for the threshold query.
 
-        :return: For each item in the input, all items above the threshold are returned.
+        :return: For each item in the input, the items above the threshold are returned in the form of
+                (NAME, SIMILARITY) tuples.
         """
-        vectors = np.array(vectors)
+        vectors = np.asarray(vectors)
         if np.ndim(vectors) == 1:
             vectors = vectors[None, :]
 
         out = []
-        for indexes in self.backend.threshold(vectors, threshold):
-            out.append([self.items[idx] for idx in indexes])
+        for indices, distances in self.backend.threshold(vectors, threshold, max_k=max_k):
+            distances.clip(min=0, out=distances)
+            out.append([(self.items[idx], dist) for idx, dist in zip(indices, distances)])
 
         return out
 

{vicinity-0.3.3 → vicinity-0.4.0}/vicinity.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vicinity
-Version: 0.3.3
+Version: 0.4.0
 Summary: Lightweight Nearest Neighbors with Flexible Backends
 Author-email: Stéphan Tulkens <stephantul@gmail.com>, Thomas van Dongen <thomas123@live.nl>
 License: MIT License