diff --git a/tests/test_vicinity.py b/tests/test_vicinity.py index 1f87175..1a24ac2 100644 --- a/tests/test_vicinity.py +++ b/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() 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..24433c2 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 @@ -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,33 +202,29 @@ 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 @@ -201,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/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..3062882 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 @@ -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) 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.""" 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__))