fix: Fixed wrong distances and neighbours across backends

tests/test_vicinity.py

@@ -7,7 +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 Metric, normalize
 
 BackendType = tuple[Backend, str]
 
@@ -41,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:
@@ -144,6 +147,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:
@@ -364,3 +368,89 @@ 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,atol",
+    [
+        (Backend.BASIC, {}, 1e-4),
+        (Backend.HNSW, {}, 1e-4),
+        (Backend.PYNNDESCENT, {}, 1e-4),
+        (Backend.VOYAGER, {}, 1e-4),
+        (Backend.FAISS, {"index_type": "flat"}, 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),
+        (Backend.FAISS, {"index_type": "ivf_scalar", "nlist": 50}, 0.05),
+    ],
+)
+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
+    )
+
+    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, 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":
+        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(
+    "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}),
+        (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(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)
+    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()

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

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
@@ -55,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
@@ -79,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)
 
@@ -88,11 +93,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.")
@@ -103,9 +108,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:
@@ -128,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."""
@@ -147,16 +154,44 @@ 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":
+        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)
-        if self.arguments.metric == "cosine":
-            distances = 1 - distances
-        return list(zip(indices, distances))
+        out: QueryResult = []
+        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], idx[found], zero_query)))
+        return out
+
+    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
+        raw = np.maximum(raw, 0)
+        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."""
+        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":
+        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)
 
@@ -167,38 +202,36 @@ 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":
+        zero_queries = np.linalg.norm(vectors, axis=1) == 0
+        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), 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]))
 
         return out
 
     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
     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)

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."""

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
@@ -110,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)

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."""

vicinity/version.py

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