From 2b739710cff426be9de184e2992a21c23343e372 Mon Sep 17 00:00:00 2001 From: Pringled Date: Fri, 25 Sep 2026 17:33:09 +0200 Subject: [PATCH 1/5] fix: return true cosine and Euclidean distances from all backends - faiss: compare against Metric.COSINE (the string comparison never matched), convert inner-product and squared-L2 scores to distances, use metric-aware range search radii, and drop results padded with index -1 - annoy: store the metric enum so cosine queries are normalized - voyager: add items with explicit ids (ids could be assigned out of order) and return non-squared Euclidean distances - hnsw: return non-squared Euclidean distances - pynndescent: only normalize query vectors for the cosine metric --- tests/test_vicinity.py | 31 ++++++++++++++++++ vicinity/backends/annoy.py | 4 ++- vicinity/backends/faiss.py | 54 +++++++++++++++++++------------- vicinity/backends/hnsw.py | 7 ++++- vicinity/backends/pynndescent.py | 10 +++--- vicinity/backends/voyager.py | 9 ++++-- 6 files changed, 86 insertions(+), 29 deletions(-) diff --git a/tests/test_vicinity.py b/tests/test_vicinity.py index 1f87175..9b5d24e 100644 --- a/tests/test_vicinity.py +++ b/tests/test_vicinity.py @@ -8,6 +8,7 @@ from vicinity import Vicinity from vicinity.datatypes import Backend +from vicinity.utils import normalize BackendType = tuple[Backend, str] @@ -364,3 +365,33 @@ def test_vicinity_usearch_binary_metrics(tmp_path: Path, metric: str) -> None: with pytest.raises(ValueError, match="bit-packed"): Vicinity.from_vectors_and_items(bits.astype(np.float32), items, backend_type=Backend.USEARCH, metric=metric) + + +@pytest.mark.parametrize("metric", ["cosine", "euclidean"]) +@pytest.mark.parametrize( + "backend_type,kwargs", + [ + (Backend.BASIC, {}), + (Backend.HNSW, {}), + (Backend.PYNNDESCENT, {}), + (Backend.VOYAGER, {}), + (Backend.FAISS, {"index_type": "flat"}), + (Backend.FAISS, {"index_type": "ivf", "nlist": 50}), + (Backend.FAISS, {"index_type": "hnsw"}), + ], +) +def test_backend_distances_match_metric( + backend_type: Backend, kwargs: dict, metric: str, vectors: np.ndarray, query_vector: np.ndarray +) -> None: + """Backends return true cosine or Euclidean distances for the returned items, without padding.""" + vicinity = Vicinity.from_vectors_and_items( + vectors, list(range(len(vectors))), backend_type=backend_type, metric=metric, **kwargs + ) + if metric == "cosine": + expected = 1 - normalize(vectors) @ normalize(query_vector) + else: + expected = np.linalg.norm(vectors - query_vector, axis=1) + threshold = float(np.sort(expected)[10]) + for result in (vicinity.query(query_vector, k=100)[0], vicinity.query_threshold(query_vector, threshold)[0]): + items, distances = zip(*result) + assert np.allclose(distances, expected[list(items)], atol=1e-4) diff --git a/vicinity/backends/annoy.py b/vicinity/backends/annoy.py index 82b738a..34d412a 100644 --- a/vicinity/backends/annoy.py +++ b/vicinity/backends/annoy.py @@ -67,7 +67,9 @@ def from_vectors( index.add_item(i, vector) index.build(trees) - arguments = AnnoyArgs(dim=dim, metric=metric, trees=trees, length=len(vectors), internal_metric=internal_metric) # type: ignore + arguments = AnnoyArgs( + dim=dim, metric=metric_enum, trees=trees, length=len(vectors), internal_metric=internal_metric + ) return AnnoyBackend(index, arguments=arguments) @property diff --git a/vicinity/backends/faiss.py b/vicinity/backends/faiss.py index 2b3c6da..6e741d4 100644 --- a/vicinity/backends/faiss.py +++ b/vicinity/backends/faiss.py @@ -6,6 +6,7 @@ from typing import Any import faiss +import numpy as np from numpy import typing as npt from vicinity.backends.base import AbstractBackend, BaseArgs @@ -147,16 +148,33 @@ def dim(self) -> int: 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": + if self.arguments.metric == Metric.COSINE: vectors = normalize(vectors) distances, indices = self.index.search(vectors, k) - if self.arguments.metric == "cosine": - distances = 1 - distances - return list(zip(indices, distances)) + out: QueryResult = [] + for idx, raw in zip(indices, distances): + # FAISS pads missing results with index -1. + found = idx >= 0 + out.append((idx[found], self._to_distances(raw[found]))) + return out + + def _to_distances(self, raw: npt.NDArray) -> npt.NDArray: + """Convert raw FAISS scores to distances for the configured metric.""" + if self.index.metric_type == faiss.METRIC_INNER_PRODUCT: + return 1 - raw + # L2 indexes return squared distances, which are 2 - 2 * cosine for normalized vectors. + raw = np.maximum(raw, 0) + return raw / 2 if self.arguments.metric == Metric.COSINE else np.sqrt(raw) + + def _radius(self, threshold: float) -> float: + """Convert a distance threshold to a FAISS range search radius.""" + if self.index.metric_type == faiss.METRIC_INNER_PRODUCT: + return 1 - threshold + return 2 * threshold if self.arguments.metric == Metric.COSINE else threshold**2 def insert(self, vectors: npt.NDArray) -> None: """Insert vectors into the backend.""" - if self.arguments.metric == "cosine": + if self.arguments.metric == Metric.COSINE: vectors = normalize(vectors) self.index.add(vectors) @@ -167,27 +185,21 @@ def delete(self, indices: list[int]) -> None: def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult: """Query vectors within a distance threshold, using range_search if supported.""" out: QueryResult = [] - if self.arguments.metric == "cosine": + if self.arguments.metric == Metric.COSINE: vectors = normalize(vectors) if isinstance(self.index, RANGE_SEARCH_INDEXES): - radius = threshold - lims, D, I = self.index.range_search(vectors, radius) - for i in range(vectors.shape[0]): - start, end = lims[i], lims[i + 1] - idx = I[start:end] - dist = D[start:end] - if self.arguments.metric == "cosine": - dist = 1 - dist - mask = dist < threshold - out.append((idx[mask], dist[mask])) + lims, D, I = self.index.range_search(vectors, self._radius(threshold)) + results = [(I[lims[i] : lims[i + 1]], D[lims[i] : lims[i + 1]]) for i in range(vectors.shape[0])] else: distances, indices = self.index.search(vectors, max_k) - for dist, idx in zip(distances, indices): - if self.arguments.metric == "cosine": - dist = 1 - dist - mask = dist < threshold - out.append((idx[mask], dist[mask])) + results = list(zip(indices, distances)) + + for idx, raw in results: + dist = self._to_distances(raw) + # FAISS pads missing results with index -1. + mask = (idx >= 0) & (dist < threshold) + out.append((idx[mask], dist[mask])) return out diff --git a/vicinity/backends/hnsw.py b/vicinity/backends/hnsw.py index 7f65006..c12d00a 100644 --- a/vicinity/backends/hnsw.py +++ b/vicinity/backends/hnsw.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any +import numpy as np from hnswlib import Index as HnswIndex from numpy import typing as npt @@ -94,7 +95,11 @@ def save(self, path: Path) -> None: 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))) + indices, distances = self.index.knn_query(vectors, k) + if self.arguments.metric == Metric.EUCLIDEAN: + # hnswlib returns squared Euclidean distances. + distances = np.sqrt(distances) + return list(zip(indices, distances)) def insert(self, vectors: npt.NDArray) -> None: """Insert vectors into the backend.""" diff --git a/vicinity/backends/pynndescent.py b/vicinity/backends/pynndescent.py index 0b398e9..56a23b9 100644 --- a/vicinity/backends/pynndescent.py +++ b/vicinity/backends/pynndescent.py @@ -68,8 +68,9 @@ def dim(self) -> int: def query(self, vectors: npt.NDArray, k: int) -> QueryResult: """Batched approximate nearest neighbors search.""" - normalized_vectors = normalize_or_copy(vectors) - indices, distances = self.index.query(normalized_vectors, k=k) + if self.arguments.metric == Metric.COSINE: + vectors = normalize_or_copy(vectors) + indices, distances = self.index.query(vectors, k=k) return list(zip(indices, distances)) def insert(self, vectors: npt.NDArray) -> None: @@ -82,8 +83,9 @@ def delete(self, indices: list[int]) -> None: 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=max_k) + if self.arguments.metric == Metric.COSINE: + vectors = normalize_or_copy(vectors) + indices, distances = self.index.query(vectors, k=max_k) out: QueryResult = [] for idx, dist in zip(indices, distances): mask = dist < threshold diff --git a/vicinity/backends/voyager.py b/vicinity/backends/voyager.py index 86f18be..55b3490 100644 --- a/vicinity/backends/voyager.py +++ b/vicinity/backends/voyager.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any +import numpy as np from numpy import typing as npt from voyager import Index, Space @@ -60,7 +61,8 @@ def from_vectors( M=m, ef_construction=ef_construction, ) - index.add_items(vectors) + # Explicit ids, since Voyager does not guarantee input order when assigning them. + index.add_items(vectors, ids=np.arange(len(vectors))) return cls( index, VoyagerArgs(dim=dim, metric=metric_enum, ef_construction=ef_construction, m=m), @@ -70,6 +72,9 @@ def query(self, vectors: npt.NDArray, k: int) -> QueryResult: """Query the backend for the nearest neighbors.""" k = min(k, len(self)) indices, distances = self.index.query(vectors, k) + if self.arguments.metric == Metric.EUCLIDEAN: + # Voyager returns squared Euclidean distances. + distances = np.sqrt(distances) return list(zip(indices, distances)) @classmethod @@ -89,7 +94,7 @@ def save(self, path: Path) -> None: def insert(self, vectors: npt.NDArray) -> None: """Insert vectors into the backend.""" - self.index.add_items(vectors) + self.index.add_items(vectors, ids=np.arange(len(self), len(self) + len(vectors))) def delete(self, indices: list[int]) -> None: """Delete vectors from the backend.""" From a93309be04ee82515c678bb1b4da148efb4b7130 Mon Sep 17 00:00:00 2001 From: Pringled Date: Sat, 26 Sep 2026 08:04:08 +0200 Subject: [PATCH 2/5] fix: build cosine FAISS indexes with inner product and leave LSH scores unconverted Halving squared L2 only gives cosine distance for unit vectors, so zero vectors came out at distance 0.5 instead of 1 on hnsw/scalar/ivf_scalar/ivfpq indexes. These are now built with inner product where FAISS supports it. LSH returns Hamming distances, which cannot be converted, so they are returned as is. The contract test now covers scalar indexes, zero vectors and threshold membership. --- tests/test_vicinity.py | 53 +++++++++++++++++++++++++++++--------- vicinity/backends/faiss.py | 15 +++++++---- 2 files changed, 51 insertions(+), 17 deletions(-) diff --git a/tests/test_vicinity.py b/tests/test_vicinity.py index 9b5d24e..7098e70 100644 --- a/tests/test_vicinity.py +++ b/tests/test_vicinity.py @@ -369,21 +369,24 @@ def test_vicinity_usearch_binary_metrics(tmp_path: Path, metric: str) -> None: @pytest.mark.parametrize("metric", ["cosine", "euclidean"]) @pytest.mark.parametrize( - "backend_type,kwargs", + "backend_type,kwargs,atol", [ - (Backend.BASIC, {}), - (Backend.HNSW, {}), - (Backend.PYNNDESCENT, {}), - (Backend.VOYAGER, {}), - (Backend.FAISS, {"index_type": "flat"}), - (Backend.FAISS, {"index_type": "ivf", "nlist": 50}), - (Backend.FAISS, {"index_type": "hnsw"}), + (Backend.BASIC, {}, 1e-4), + (Backend.HNSW, {}, 1e-4), + (Backend.PYNNDESCENT, {}, 1e-4), + (Backend.VOYAGER, {}, 1e-4), + (Backend.FAISS, {"index_type": "flat"}, 1e-4), + (Backend.FAISS, {"index_type": "ivf", "nlist": 50}, 1e-4), + (Backend.FAISS, {"index_type": "hnsw"}, 1e-4), + # Scalar quantization makes distances approximate. + (Backend.FAISS, {"index_type": "scalar"}, 0.05), + (Backend.FAISS, {"index_type": "ivf_scalar", "nlist": 50}, 0.05), ], ) def test_backend_distances_match_metric( - backend_type: Backend, kwargs: dict, metric: str, vectors: np.ndarray, query_vector: np.ndarray + backend_type: Backend, kwargs: dict, atol: float, metric: str, vectors: np.ndarray, query_vector: np.ndarray ) -> None: - """Backends return true cosine or Euclidean distances for the returned items, without padding.""" + """Backends return true cosine or Euclidean distances without padding; exact backends return every close item.""" vicinity = Vicinity.from_vectors_and_items( vectors, list(range(len(vectors))), backend_type=backend_type, metric=metric, **kwargs ) @@ -391,7 +394,33 @@ def test_backend_distances_match_metric( expected = 1 - normalize(vectors) @ normalize(query_vector) else: expected = np.linalg.norm(vectors - query_vector, axis=1) - threshold = float(np.sort(expected)[10]) + # Halfway between the 11th and 12th closest items, so no item sits on the boundary. + threshold = float(np.sort(expected)[10:12].mean()) for result in (vicinity.query(query_vector, k=100)[0], vicinity.query_threshold(query_vector, threshold)[0]): items, distances = zip(*result) - assert np.allclose(distances, expected[list(items)], atol=1e-4) + assert np.allclose(distances, expected[list(items)], atol=atol) + if backend_type == Backend.BASIC or kwargs.get("index_type") == "flat": + returned = {item for item, _ in vicinity.query_threshold(query_vector, threshold)[0]} + assert returned == set(np.flatnonzero(expected < threshold).tolist()) + + +@pytest.mark.parametrize( + "backend_type,kwargs", + [ + (Backend.BASIC, {}), + (Backend.FAISS, {"index_type": "flat"}), + (Backend.FAISS, {"index_type": "hnsw"}), + (Backend.FAISS, {"index_type": "scalar"}), + (Backend.FAISS, {"index_type": "ivf_scalar", "nlist": 1}), + ], +) +def test_cosine_distance_to_zero_vector(backend_type: Backend, kwargs: dict) -> None: + """A zero vector has cosine distance 1 to everything, so it never falls within a threshold below 1.""" + vectors = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float32) + vicinity = Vicinity.from_vectors_and_items( + vectors, ["zero", "same", "orthogonal"], backend_type=backend_type, **kwargs + ) + query = np.array([1.0, 0.0], dtype=np.float32) + distances = dict(vicinity.query(query, k=3)[0]) + assert np.allclose([distances["zero"], distances["same"], distances["orthogonal"]], [1.0, 0.0, 1.0], atol=0.01) + assert [item for item, _ in vicinity.query_threshold(query, threshold=0.75)[0]] == ["same"] diff --git a/vicinity/backends/faiss.py b/vicinity/backends/faiss.py index 6e741d4..af72145 100644 --- a/vicinity/backends/faiss.py +++ b/vicinity/backends/faiss.py @@ -89,11 +89,11 @@ def from_vectors( # noqa: C901 if index_type == "flat": index = faiss.IndexFlat(dim, faiss_metric) elif index_type == "hnsw": - index = faiss.IndexHNSWFlat(dim, m) + index = faiss.IndexHNSWFlat(dim, m, faiss_metric) elif index_type == "lsh": index = faiss.IndexLSH(dim, nbits) elif index_type == "scalar": - index = faiss.IndexScalarQuantizer(dim, faiss.ScalarQuantizer.QT_8bit) + index = faiss.IndexScalarQuantizer(dim, faiss.ScalarQuantizer.QT_8bit, faiss_metric) elif index_type == "pq": if not (1 <= nbits <= 16): logger.warning(f"Invalid nbits={nbits} for IndexPQ. Setting nbits to 16.") @@ -104,9 +104,11 @@ def from_vectors( # noqa: C901 if index_type == "ivf": index = faiss.IndexIVFFlat(quantizer, dim, nlist, faiss_metric) elif index_type == "ivf_scalar": - index = faiss.IndexIVFScalarQuantizer(quantizer, dim, nlist, faiss.ScalarQuantizer.QT_8bit) + index = faiss.IndexIVFScalarQuantizer( + quantizer, dim, nlist, faiss.ScalarQuantizer.QT_8bit, faiss_metric + ) elif index_type == "ivfpq": - index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, nbits) + index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, nbits, faiss_metric) elif index_type == "ivfpqr": index = faiss.IndexIVFPQR(quantizer, dim, nlist, m, nbits, m, refine_nbits) else: @@ -160,9 +162,12 @@ def query(self, vectors: npt.NDArray, k: int) -> QueryResult: def _to_distances(self, raw: npt.NDArray) -> npt.NDArray: """Convert raw FAISS scores to distances for the configured metric.""" + if isinstance(self.index, faiss.IndexLSH): + # LSH returns Hamming distances between binary codes, which cannot be converted. + return raw if self.index.metric_type == faiss.METRIC_INNER_PRODUCT: return 1 - raw - # L2 indexes return squared distances, which are 2 - 2 * cosine for normalized vectors. + # L2 indexes (pq and ivfpqr) return squared distances, which are 2 - 2 * cosine for unit vectors. raw = np.maximum(raw, 0) return raw / 2 if self.arguments.metric == Metric.COSINE else np.sqrt(raw) From 0db0d284a1602d005a380e834e0c5cf8440c43aa Mon Sep 17 00:00:00 2001 From: Pringled Date: Sat, 26 Sep 2026 08:10:10 +0200 Subject: [PATCH 3/5] fix: handle zero vectors on L2-built FAISS indexes and restore PyNNDescent graphs on load - faiss: pq and ivfpqr stay on L2, where halving squared L2 only gives cosine distance for unit vectors. Zero vectors (stored or queried) now get cosine distance 1; stored zero vectors are tracked through inserts and save/load. - pynndescent: the saved (indices, distances) neighbour graph came back as one float array, so querying a loaded index failed. The original dtypes are restored on load, which also fixes indexes saved by earlier versions. - The save/load test now queries the loaded index. --- tests/test_vicinity.py | 20 +++++++++++-------- vicinity/backends/faiss.py | 34 +++++++++++++++++++++++--------- vicinity/backends/pynndescent.py | 4 +++- 3 files changed, 40 insertions(+), 18 deletions(-) diff --git a/tests/test_vicinity.py b/tests/test_vicinity.py index 7098e70..885aa05 100644 --- a/tests/test_vicinity.py +++ b/tests/test_vicinity.py @@ -145,6 +145,7 @@ def test_vicinity_save_and_load(tmp_path: Path, vicinity_instance: Vicinity) -> v = Vicinity.load(save_path) assert v.vector_store is None + assert v.query(np.ones(v.dim), k=5)[0] def test_vicinity_save_and_load_vector_store(tmp_path: Path, vicinity_instance_with_stored_vectors: Vicinity) -> None: @@ -412,15 +413,18 @@ def test_backend_distances_match_metric( (Backend.FAISS, {"index_type": "hnsw"}), (Backend.FAISS, {"index_type": "scalar"}), (Backend.FAISS, {"index_type": "ivf_scalar", "nlist": 1}), + (Backend.FAISS, {"index_type": "pq", "m": 1, "nbits": 3}), + (Backend.FAISS, {"index_type": "ivfpq", "nlist": 1, "m": 1, "nbits": 3}), + (Backend.FAISS, {"index_type": "ivfpqr", "nlist": 1, "m": 1, "nbits": 3, "refine_nbits": 3}), ], ) def test_cosine_distance_to_zero_vector(backend_type: Backend, kwargs: dict) -> None: - """A zero vector has cosine distance 1 to everything, so it never falls within a threshold below 1.""" - vectors = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]], dtype=np.float32) - vicinity = Vicinity.from_vectors_and_items( - vectors, ["zero", "same", "orthogonal"], backend_type=backend_type, **kwargs - ) + """Zero vectors have cosine distance 1 to everything, both as stored items and as queries.""" + # Apart from the zero vector, every vector points away from the query, so nothing falls within the threshold. + vectors = np.array([[0.0, 0.0]] + [[-1.0, 0.1 * i] for i in range(15)], dtype=np.float32) + vicinity = Vicinity.from_vectors_and_items(vectors, list(range(len(vectors))), backend_type=backend_type, **kwargs) query = np.array([1.0, 0.0], dtype=np.float32) - distances = dict(vicinity.query(query, k=3)[0]) - assert np.allclose([distances["zero"], distances["same"], distances["orthogonal"]], [1.0, 0.0, 1.0], atol=0.01) - assert [item for item, _ in vicinity.query_threshold(query, threshold=0.75)[0]] == ["same"] + assert dict(vicinity.query(query, k=len(vectors))[0])[0] == pytest.approx(1.0, abs=0.01) + assert vicinity.query_threshold(query, threshold=0.75)[0] == [] + zero_query_distances = [distance for _, distance in vicinity.query(np.zeros(2, dtype=np.float32), k=5)[0]] + assert np.allclose(zero_query_distances, 1.0, atol=0.01) diff --git a/vicinity/backends/faiss.py b/vicinity/backends/faiss.py index af72145..24433c2 100644 --- a/vicinity/backends/faiss.py +++ b/vicinity/backends/faiss.py @@ -56,10 +56,13 @@ def __init__( self, index: faiss.Index, arguments: FaissArgs, + zero_indices: npt.NDArray | None = None, ) -> None: """Initialize the backend using a FAISS index.""" super().__init__(arguments) self.index = index + # Indices of zero vectors, whose cosine distance cannot be derived from squared L2. + self.zero_indices = np.zeros(0, dtype=np.int64) if zero_indices is None else zero_indices @classmethod def from_vectors( # noqa: C901 @@ -80,6 +83,7 @@ def from_vectors( # noqa: C901 raise ValueError(f"Metric '{metric_enum.value}' is not supported by FaissBackend.") faiss_metric = cls._map_metric_to_string(metric_enum) + zero_indices = np.flatnonzero(np.linalg.norm(vectors, axis=1) == 0) if faiss_metric == faiss.METRIC_INNER_PRODUCT: vectors = normalize(vectors) @@ -131,7 +135,7 @@ def from_vectors( # noqa: C901 nbits=nbits, refine_nbits=refine_nbits, ) - return cls(index=index, arguments=arguments) + return cls(index=index, arguments=arguments, zero_indices=zero_indices) def __len__(self) -> int: """Return the number of vectors in the index.""" @@ -150,26 +154,32 @@ def dim(self) -> int: def query(self, vectors: npt.NDArray, k: int) -> QueryResult: """Perform a k-NN search in the FAISS index.""" k = min(len(self), k) + zero_queries = np.linalg.norm(vectors, axis=1) == 0 if self.arguments.metric == Metric.COSINE: vectors = normalize(vectors) distances, indices = self.index.search(vectors, k) out: QueryResult = [] - for idx, raw in zip(indices, distances): + for idx, raw, zero_query in zip(indices, distances, zero_queries): # FAISS pads missing results with index -1. found = idx >= 0 - out.append((idx[found], self._to_distances(raw[found]))) + out.append((idx[found], self._to_distances(raw[found], idx[found], zero_query))) return out - def _to_distances(self, raw: npt.NDArray) -> npt.NDArray: + def _to_distances(self, raw: npt.NDArray, indices: npt.NDArray, zero_query: bool) -> npt.NDArray: """Convert raw FAISS scores to distances for the configured metric.""" if isinstance(self.index, faiss.IndexLSH): # LSH returns Hamming distances between binary codes, which cannot be converted. return raw if self.index.metric_type == faiss.METRIC_INNER_PRODUCT: return 1 - raw - # L2 indexes (pq and ivfpqr) return squared distances, which are 2 - 2 * cosine for unit vectors. raw = np.maximum(raw, 0) - return raw / 2 if self.arguments.metric == Metric.COSINE else np.sqrt(raw) + if self.arguments.metric != Metric.COSINE: + return np.sqrt(raw) + # L2 indexes (pq and ivfpqr) return squared distances, which are 2 - 2 * cosine for unit vectors. + # Zero vectors are not unit vectors, and have cosine distance 1 to everything. + distances = raw / 2 + distances[np.isin(indices, self.zero_indices)] = 1.0 + return np.ones_like(distances) if zero_query else distances def _radius(self, threshold: float) -> float: """Convert a distance threshold to a FAISS range search radius.""" @@ -179,6 +189,8 @@ def _radius(self, threshold: float) -> float: def insert(self, vectors: npt.NDArray) -> None: """Insert vectors into the backend.""" + new_zero_indices = np.flatnonzero(np.linalg.norm(vectors, axis=1) == 0) + len(self) + self.zero_indices = np.concatenate([self.zero_indices, new_zero_indices]) if self.arguments.metric == Metric.COSINE: vectors = normalize(vectors) self.index.add(vectors) @@ -190,6 +202,7 @@ def delete(self, indices: list[int]) -> None: def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> QueryResult: """Query vectors within a distance threshold, using range_search if supported.""" out: QueryResult = [] + zero_queries = np.linalg.norm(vectors, axis=1) == 0 if self.arguments.metric == Metric.COSINE: vectors = normalize(vectors) @@ -200,8 +213,8 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query distances, indices = self.index.search(vectors, max_k) results = list(zip(indices, distances)) - for idx, raw in results: - dist = self._to_distances(raw) + for (idx, raw), zero_query in zip(results, zero_queries): + dist = self._to_distances(raw, idx, zero_query) # FAISS pads missing results with index -1. mask = (idx >= 0) & (dist < threshold) out.append((idx[mask], dist[mask])) @@ -211,6 +224,7 @@ def threshold(self, vectors: npt.NDArray, threshold: float, max_k: int) -> Query def save(self, path: Path) -> None: """Save the FAISS index and arguments.""" faiss.write_index(self.index, str(path / "index.faiss")) + np.save(path / "zero_indices.npy", self.zero_indices) self.arguments.dump(path / "arguments.json") @classmethod @@ -218,4 +232,6 @@ def load(cls: type[FaissBackend], path: Path) -> FaissBackend: """Load a FAISS index and arguments.""" arguments = FaissArgs.load(path / "arguments.json") index = faiss.read_index(str(path / "index.faiss")) - return cls(index=index, arguments=arguments) + zero_indices_path = path / "zero_indices.npy" + zero_indices = np.load(zero_indices_path) if zero_indices_path.exists() else None + return cls(index=index, arguments=arguments, zero_indices=zero_indices) diff --git a/vicinity/backends/pynndescent.py b/vicinity/backends/pynndescent.py index 56a23b9..3062882 100644 --- a/vicinity/backends/pynndescent.py +++ b/vicinity/backends/pynndescent.py @@ -112,6 +112,8 @@ def load(cls: type[PyNNDescentBackend], path: Path) -> PyNNDescentBackend: # Load the neighbor graph if it was saved neighbor_graph_path = path / "neighbor_graph.npy" if neighbor_graph_path.exists(): - index._neighbor_graph = np.load(str(neighbor_graph_path), allow_pickle=True) + # The (indices, distances) tuple is saved as one float array, so restore the original dtypes. + indices, distances = np.load(str(neighbor_graph_path), allow_pickle=True) + index._neighbor_graph = (indices.astype(np.int32), distances.astype(np.float32)) return cls(index=index, arguments=arguments) From 42d6a339eff2af4d4a2d192e28fe8dbf33daa6a1 Mon Sep 17 00:00:00 2001 From: Pringled Date: Sat, 26 Sep 2026 08:43:31 +0200 Subject: [PATCH 4/5] test: cover every backend distance fix Every reachable fix now fails a test when reverted: FAISS padding (IVF with small clusters), range-search radii (thresholds beyond 0.5 cosine / 1 Euclidean), zero vectors through insert and save/load, LSH Hamming distances, item-to-vector mapping (every stored vector queried for itself), and string metrics stored as enums. --- tests/test_vicinity.py | 54 +++++++++++++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 14 deletions(-) diff --git a/tests/test_vicinity.py b/tests/test_vicinity.py index 885aa05..1a24ac2 100644 --- a/tests/test_vicinity.py +++ b/tests/test_vicinity.py @@ -7,8 +7,9 @@ from orjson import JSONEncodeError from vicinity import Vicinity +from vicinity.backends.faiss import FaissBackend from vicinity.datatypes import Backend -from vicinity.utils import normalize +from vicinity.utils import Metric, normalize BackendType = tuple[Backend, str] @@ -42,11 +43,12 @@ def test_vicinity_from_vectors_and_items(backend_type: BackendType, items: list[ :param vectors: An array of vectors. """ backend = backend_type[0] - vicinity = Vicinity.from_vectors_and_items(vectors, items, backend_type=backend) + vicinity = Vicinity.from_vectors_and_items(vectors, items, backend_type=backend, metric="cosine") assert len(vicinity) == len(items) assert vicinity.items == items assert vicinity.dim == vectors.shape[1] + assert vicinity.metric is Metric.COSINE def test_vicinity_query(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None: @@ -377,7 +379,8 @@ def test_vicinity_usearch_binary_metrics(tmp_path: Path, metric: str) -> None: (Backend.PYNNDESCENT, {}, 1e-4), (Backend.VOYAGER, {}, 1e-4), (Backend.FAISS, {"index_type": "flat"}, 1e-4), - (Backend.FAISS, {"index_type": "ivf", "nlist": 50}, 1e-4), + # Clusters of about 10 vectors, so a query for 100 neighbours gets padded results. + (Backend.FAISS, {"index_type": "ivf", "nlist": 1000}, 1e-4), (Backend.FAISS, {"index_type": "hnsw"}, 1e-4), # Scalar quantization makes distances approximate. (Backend.FAISS, {"index_type": "scalar"}, 0.05), @@ -388,21 +391,30 @@ def test_backend_distances_match_metric( backend_type: Backend, kwargs: dict, atol: float, metric: str, vectors: np.ndarray, query_vector: np.ndarray ) -> None: """Backends return true cosine or Euclidean distances without padding; exact backends return every close item.""" + # Centred and scaled, so distances go beyond 0.5 (cosine) and 1 (Euclidean), where FAISS range radii differ. + vectors, query = 2 * (vectors - 0.5), 2 * (query_vector - 0.5) vicinity = Vicinity.from_vectors_and_items( vectors, list(range(len(vectors))), backend_type=backend_type, metric=metric, **kwargs ) - if metric == "cosine": - expected = 1 - normalize(vectors) @ normalize(query_vector) - else: - expected = np.linalg.norm(vectors - query_vector, axis=1) + + def distance(a: np.ndarray, b: np.ndarray) -> np.ndarray: + if metric == "cosine": + return 1 - np.sum(normalize(a) * normalize(b), axis=-1) + return np.linalg.norm(a - b, axis=-1) + + expected = distance(vectors, query) # Halfway between the 11th and 12th closest items, so no item sits on the boundary. threshold = float(np.sort(expected)[10:12].mean()) - for result in (vicinity.query(query_vector, k=100)[0], vicinity.query_threshold(query_vector, threshold)[0]): + for result in (vicinity.query(query, k=100)[0], vicinity.query_threshold(query, threshold)[0]): items, distances = zip(*result) assert np.allclose(distances, expected[list(items)], atol=atol) + # Every stored vector must come back with its true distance, which catches items mapped to the wrong vector. + items, distances = zip(*(result[0] for result in vicinity.query(vectors, k=1))) + assert np.allclose(distances, distance(vectors, vectors[list(items)]), atol=atol) if backend_type == Backend.BASIC or kwargs.get("index_type") == "flat": - returned = {item for item, _ in vicinity.query_threshold(query_vector, threshold)[0]} - assert returned == set(np.flatnonzero(expected < threshold).tolist()) + for limit in (threshold, float(np.median(expected))): + returned = {item for item, _ in vicinity.query_threshold(query, limit)[0]} + assert returned == set(np.flatnonzero(expected < limit).tolist()) @pytest.mark.parametrize( @@ -418,13 +430,27 @@ def test_backend_distances_match_metric( (Backend.FAISS, {"index_type": "ivfpqr", "nlist": 1, "m": 1, "nbits": 3, "refine_nbits": 3}), ], ) -def test_cosine_distance_to_zero_vector(backend_type: Backend, kwargs: dict) -> None: - """Zero vectors have cosine distance 1 to everything, both as stored items and as queries.""" - # Apart from the zero vector, every vector points away from the query, so nothing falls within the threshold. +def test_cosine_distance_to_zero_vector(tmp_path: Path, backend_type: Backend, kwargs: dict) -> None: + """Zero vectors have cosine distance 1 to everything: built, inserted, reloaded and as queries.""" + # Apart from the zero vectors, every vector points away from the query, so nothing falls within the threshold. vectors = np.array([[0.0, 0.0]] + [[-1.0, 0.1 * i] for i in range(15)], dtype=np.float32) vicinity = Vicinity.from_vectors_and_items(vectors, list(range(len(vectors))), backend_type=backend_type, **kwargs) + vicinity.insert([16], np.zeros((1, 2), dtype=np.float32)) + vicinity.save(tmp_path / "vicinity") + vicinity = Vicinity.load(tmp_path / "vicinity") query = np.array([1.0, 0.0], dtype=np.float32) - assert dict(vicinity.query(query, k=len(vectors))[0])[0] == pytest.approx(1.0, abs=0.01) + distances = dict(vicinity.query(query, k=len(vectors) + 1)[0]) + assert [distances[0], distances[16]] == pytest.approx([1.0, 1.0], abs=0.01) assert vicinity.query_threshold(query, threshold=0.75)[0] == [] zero_query_distances = [distance for _, distance in vicinity.query(np.zeros(2, dtype=np.float32), k=5)[0]] assert np.allclose(zero_query_distances, 1.0, atol=0.01) + + +def test_faiss_lsh_returns_hamming_distances(vectors: np.ndarray, query_vector: np.ndarray) -> None: + """LSH distances are FAISS's Hamming distances, which cannot be converted to cosine distances.""" + vicinity = Vicinity.from_vectors_and_items( + vectors, list(range(len(vectors))), backend_type=Backend.FAISS, index_type="lsh", nbits=32 + ) + assert isinstance(vicinity.backend, FaissBackend) + hamming, _ = vicinity.backend.index.search(normalize(query_vector)[None], 10) + assert [distance for _, distance in vicinity.query(query_vector, k=10)[0]] == hamming[0].tolist() From 52faf317fe086d4414e0a2eddf18385fee4cab91 Mon Sep 17 00:00:00 2001 From: Pringled Date: Sat, 26 Sep 2026 09:18:52 +0200 Subject: [PATCH 5/5] chore: bump version to 0.4.6 --- vicinity/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vicinity/version.py b/vicinity/version.py index fc18fd8..cdc6356 100644 --- a/vicinity/version.py +++ b/vicinity/version.py @@ -1,2 +1,2 @@ -__version_triple__ = (0, 4, 5) +__version_triple__ = (0, 4, 6) __version__ = ".".join(map(str, __version_triple__))