vicinity 0.3.3__tar.gz → 0.3.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vicinity
-Version: 0.3.3
+Version: 0.3.4
 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.3.4}/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.3.4}/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)

{vicinity-0.3.3 → vicinity-0.3.4}/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:

{vicinity-0.3.3 → vicinity-0.3.4}/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:

{vicinity-0.3.3 → vicinity-0.3.4}/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))
 

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

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

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

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