diff --git a/README.md b/README.md index 91b3e2c..b8777bd 100644 --- a/README.md +++ b/README.md @@ -135,8 +135,9 @@ The following backends are supported: - `ivfpq`: Inverted file search with product quantizer. - `ivfpqr`: Inverted file search with product quantizer and refinement. - [VOYAGER](https://github.com/spotify/voyager): Voyager is a library for performing fast approximate nearest-neighbor searches on an in-memory collection of vectors. +- [TURBOVEC](https://github.com/RyanCodrai/turbovec): Quantized flat index using Google's TurboQuant, with 2-4 bits per dimension. -NOTE: the ANN backends do not support dynamic deletion. To delete items, you need to recreate the index. Insertion is supported in the following backends: `FAISS`, `HNSW`, and `Usearch`. The `BASIC` backend supports both insertion and deletion. +NOTE: most ANN backends do not support dynamic deletion. To delete items, you need to recreate the index. Insertion is supported in the following backends: `FAISS`, `HNSW`, `Usearch`, and `TurboVec`. The `BASIC` and `TurboVec` backends support both insertion and deletion. ### Backend Parameters @@ -165,6 +166,9 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee | **VOYAGER** | `metric` | Similarity space to use (`cosine`, `l2`). | `"cosine"` | | | `ef_construction` | The number of vectors that this index searches through when inserting a new vector into the index. | `200` | | | `m` | The number of connections between nodes in the tree’s internal data structure. | `16` | +| **TURBOVEC** | `metric` | Similarity metric to use (`cosine`). | `"cosine"` | +| | `bit_width` | Bits per dimension (`2`, `3`, or `4`). | `4` | +| | `calibrate` | Fit TQ+ calibration on a random sample of the vectors before indexing, which improves recall. | `True` | ## Installation @@ -193,6 +197,7 @@ pip install vicinity[hnsw] pip install vicinity[pynndescent] pip install vicinity[usearch] pip install vicinity[voyager] +pip install vicinity[turbovec] ``` ## License diff --git a/pyproject.toml b/pyproject.toml index 7a6998e..ac68aaf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,7 @@ annoy = ["annoy"] faiss = ["faiss-cpu"] usearch = ["usearch"] voyager = ["voyager"] +turbovec = ["turbovec"] backends = [ "hnswlib", "pynndescent>=0.5.10", @@ -69,7 +70,8 @@ backends = [ "annoy", "faiss-cpu", "usearch", - "voyager" + "voyager", + "turbovec" ] all = [ @@ -82,7 +84,8 @@ all = [ "annoy", "faiss-cpu", "usearch", - "voyager" + "voyager", + "turbovec" ] [project.urls] diff --git a/tests/conftest.py b/tests/conftest.py index 63a8199..1c3b40a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -22,13 +22,13 @@ @pytest.fixture(scope="session") -def items() -> list[str]: +def items() -> list[str | dict[str, object]]: """Fixture providing a list of item names.""" return [f"item{i}" if i % 2 == 0 else {"name": f"item{i}", "id": i} for i in range(1, 10001)] @pytest.fixture(scope="session") -def non_serializable_items() -> list[str]: +def non_serializable_items() -> list[object]: """Fixture providing a list of non-serializable items.""" class NonSerializable: @@ -58,6 +58,7 @@ def query_vector() -> np.ndarray: (Backend.PYNNDESCENT, None), (Backend.USEARCH, None), (Backend.VOYAGER, None), + (Backend.TURBOVEC, None), ] diff --git a/tests/test_vicinity.py b/tests/test_vicinity.py index 1a24ac2..fa3f17b 100644 --- a/tests/test_vicinity.py +++ b/tests/test_vicinity.py @@ -80,7 +80,10 @@ def test_vicinity_query_threshold(vicinity_instance: Vicinity, query_vector: np. results = vicinity_instance.query_threshold(np.stack([query_vector, query_vector]), threshold=0.7) - assert results[0] == results[1] + # PyNNDescent seeds each query in a batch from a shared random state, so identical queries can find different + # approximate neighbours. + if vicinity_instance.backend.backend_type != Backend.PYNNDESCENT: + assert results[0] == results[1] def test_vicinity_insert(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None: @@ -112,8 +115,8 @@ def test_vicinity_delete(vicinity_instance: Vicinity, items: list[str], vectors: :param items: List of item names. :param vectors: Array of vectors corresponding to items. """ - if vicinity_instance.backend.backend_type != Backend.BASIC: - # Skip delete for non-basic backends + if vicinity_instance.backend.backend_type not in {Backend.BASIC, Backend.TURBOVEC}: + # Only the basic and TurboVec backends support deletion. return # Get the vector corresponding to "item2" @@ -223,8 +226,8 @@ def test_vicinity_delete_nonexistent(vicinity_instance: Vicinity) -> None: :param vicinity_instance: A Vicinity instance. :raises ValueError: If deleting items that do not exist. """ - if vicinity_instance.backend.backend_type != Backend.BASIC: - # Skip delete for non-basic backends + if vicinity_instance.backend.backend_type not in {Backend.BASIC, Backend.TURBOVEC}: + # Only the basic and TurboVec backends support deletion. return with pytest.raises(ValueError): vicinity_instance.delete(["item10002"]) @@ -295,8 +298,8 @@ def test_vicinity_delete_and_query(vicinity_instance: Vicinity, items: list[str] :param items: List of item names. :param vectors: Array of vectors corresponding to items. """ - if vicinity_instance.backend.backend_type != Backend.BASIC: - # Skip delete for non-basic backends + if vicinity_instance.backend.backend_type not in {Backend.BASIC, Backend.TURBOVEC}: + # Only the basic and TurboVec backends support deletion. return # Delete some items from the Vicinity instance @@ -385,12 +388,16 @@ def test_vicinity_usearch_binary_metrics(tmp_path: Path, metric: str) -> None: # Scalar quantization makes distances approximate. (Backend.FAISS, {"index_type": "scalar"}, 0.05), (Backend.FAISS, {"index_type": "ivf_scalar", "nlist": 50}, 0.05), + # 4-bit quantization of 8-dimensional vectors is coarse, with errors up to 0.055 depending on the platform. + (Backend.TURBOVEC, {}, 0.1), ], ) 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.""" + if backend_type == Backend.TURBOVEC and metric == "euclidean": + pytest.skip("TurboVec only supports cosine.") # 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( @@ -428,6 +435,7 @@ def distance(a: np.ndarray, b: np.ndarray) -> np.ndarray: (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}), + (Backend.TURBOVEC, {}), ], ) def test_cosine_distance_to_zero_vector(tmp_path: Path, backend_type: Backend, kwargs: dict) -> None: @@ -454,3 +462,18 @@ def test_faiss_lsh_returns_hamming_distances(vectors: np.ndarray, query_vector: 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() + + +def test_turbovec_delete_keeps_items_aligned(tmp_path: Path) -> None: + """Deleting from TurboVec moves vectors between slots, but every item still maps to its own vector.""" + vectors = np.eye(16, dtype=np.float32)[:10] + vicinity = Vicinity.from_vectors_and_items(vectors, list(range(10)), backend_type=Backend.TURBOVEC) + vicinity.delete([1, 4, 9]) + vicinity.insert([10], np.eye(16, dtype=np.float32)[10:11]) + vicinity.save(tmp_path / "turbovec") + vicinity = Vicinity.load(tmp_path / "turbovec") + + remaining = [0, 2, 3, 5, 6, 7, 8, 10] + assert vicinity.items == remaining + results = vicinity.query(np.eye(16, dtype=np.float32)[remaining], k=1) + assert [result[0][0] for result in results] == remaining diff --git a/uv.lock b/uv.lock index 47fa4fc..765c469 100644 --- a/uv.lock +++ b/uv.lock @@ -8,7 +8,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-04-26T18:37:56.643847Z" +exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P1W" [[package]] @@ -1620,6 +1620,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/00/c0/8f5d070730d7836adc9c9b6408dec68c6ced86b304a9b26a14df072a6e8c/traitlets-5.14.3-py3-none-any.whl", hash = "sha256:b74e89e397b1ed28cc831db7aea759ba6640cb3de13090ca145426688ff1ac4f", size = 85359, upload-time = "2024-04-19T11:11:46.763Z" }, ] +[[package]] +name = "turbovec" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/45/d2/00c981bd1a6be5b6a875435bf678e627f191e3daf47ea70e1b5ee3109c36/turbovec-1.0.0.tar.gz", hash = "sha256:4ceb3c4ad08e48c6b3483b88aca63d81a06449f90d5fbc1b162dd11df10b8fb2", size = 702952, upload-time = "2026-08-18T20:43:31.249Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/e0/d759918244236356199b0f4d868ee331d0440a9824f147780aaefa8fe1af/turbovec-1.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:a69e8e2e26ba943f31fba546ab595b60e9b4c25087ed905c3e6ceb0d297ee234", size = 668359, upload-time = "2026-08-18T20:43:25.45Z" }, + { url = "https://files.pythonhosted.org/packages/9c/e9/8de1e7d2636a8f5c73d39c64ec6fb4bd20180d8c966962d919356d06e175/turbovec-1.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:09aa495e802701f139820e21700dc947ea21734446112afefb4e480a8f65fa90", size = 700020, upload-time = "2026-08-18T20:43:27.028Z" }, + { url = "https://files.pythonhosted.org/packages/e5/6d/c1acf31a1ef381b588de86ebebb32c3ca7d310cbb84279ac1f65c3a3e854/turbovec-1.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:c810f50b967c2163b78d1e27d29e153ab5e8820de23921ad959227acca8bb06a", size = 719275, upload-time = "2026-08-18T20:43:28.402Z" }, + { url = "https://files.pythonhosted.org/packages/4e/7d/e0c14b09f6dfb7ff156d1f884119daebc5c8bf9f5b8fc72ada26c2674256/turbovec-1.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:cd855e0b318a57dc57c733f9a62ae98de5192f4f6c2c760e305523e8ceb1b090", size = 611788, upload-time = "2026-08-18T20:43:29.854Z" }, +] + [[package]] name = "typing-extensions" version = "4.12.2" @@ -1705,6 +1720,7 @@ all = [ { name = "numba" }, { name = "numpy" }, { name = "pynndescent" }, + { name = "turbovec" }, { name = "usearch" }, { name = "voyager" }, ] @@ -1719,6 +1735,7 @@ backends = [ { name = "numba" }, { name = "numpy" }, { name = "pynndescent" }, + { name = "turbovec" }, { name = "usearch" }, { name = "voyager" }, ] @@ -1750,6 +1767,9 @@ pynndescent = [ { name = "numpy" }, { name = "pynndescent" }, ] +turbovec = [ + { name = "turbovec" }, +] usearch = [ { name = "usearch" }, ] @@ -1794,6 +1814,9 @@ requires-dist = [ { name = "ruff", marker = "extra == 'dev'" }, { name = "setuptools", marker = "extra == 'dev'" }, { name = "tqdm" }, + { name = "turbovec", marker = "extra == 'all'" }, + { name = "turbovec", marker = "extra == 'backends'" }, + { name = "turbovec", marker = "extra == 'turbovec'" }, { name = "usearch", marker = "extra == 'all'" }, { name = "usearch", marker = "extra == 'backends'" }, { name = "usearch", marker = "extra == 'usearch'" }, @@ -1801,7 +1824,7 @@ requires-dist = [ { name = "voyager", marker = "extra == 'backends'" }, { name = "voyager", marker = "extra == 'voyager'" }, ] -provides-extras = ["dev", "huggingface", "integrations", "hnsw", "pynndescent", "annoy", "faiss", "usearch", "voyager", "backends", "all"] +provides-extras = ["dev", "huggingface", "integrations", "hnsw", "pynndescent", "annoy", "faiss", "usearch", "voyager", "turbovec", "backends", "all"] [[package]] name = "virtualenv" diff --git a/vicinity/backends/__init__.py b/vicinity/backends/__init__.py index 6c5720c..c469099 100644 --- a/vicinity/backends/__init__.py +++ b/vicinity/backends/__init__.py @@ -62,5 +62,11 @@ def get_backend_class(backend: Backend | str) -> type[AbstractBackend]: return VoyagerBackend + elif backend == Backend.TURBOVEC: + _require("turbovec", backend, "turbovec") + from vicinity.backends.turbovec import TurboVecBackend + + return TurboVecBackend + __all__ = ["get_backend_class", "AbstractBackend", "BasicVectorStore"] diff --git a/vicinity/backends/turbovec.py b/vicinity/backends/turbovec.py new file mode 100644 index 0000000..4cbce68 --- /dev/null +++ b/vicinity/backends/turbovec.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +from numpy import typing as npt +from turbovec import TurboQuantIndex + +from vicinity.backends.base import AbstractBackend, BaseArgs +from vicinity.datatypes import Backend, QueryResult +from vicinity.utils import Metric, normalize + +# Rows used to fit TQ+ calibration; turbovec recommends around 1024. +_CALIBRATION_SAMPLE_SIZE = 1024 + + +@dataclass +class TurboVecArgs(BaseArgs): + dim: int = 0 + metric: Metric = Metric.COSINE + bit_width: int = 4 + calibrate: bool = True + + +class TurboVecBackend(AbstractBackend[TurboVecArgs]): + argument_class = TurboVecArgs + supported_metrics = {Metric.COSINE} + + def __init__( + self, + index: TurboQuantIndex, + arguments: TurboVecArgs, + positions: npt.NDArray | None = None, + ) -> None: + """ + Initialize the backend using TurboVec. + + :param index: The TurboVec index. + :param arguments: The arguments of the backend. + :param positions: The position of the vector in each index slot. Defaults to the slot itself. + """ + super().__init__(arguments) + self.index = index + # Deletion moves the last vector into the freed slot, so slots are mapped back to their positions. + self.positions = np.arange(len(index)) if positions is None else positions + + @classmethod + def from_vectors( + cls: type[TurboVecBackend], + vectors: npt.NDArray, + metric: str | Metric = Metric.COSINE, + bit_width: int = 4, + calibrate: bool = True, + **kwargs: Any, + ) -> TurboVecBackend: + """Create a new instance from vectors, optionally fitting TQ+ calibration on a random sample first.""" + metric_enum = Metric.from_string(metric) + + if metric_enum not in cls.supported_metrics: + raise ValueError(f"Metric '{metric_enum.value}' is not supported by TurboVecBackend.") + + if bit_width not in (2, 3, 4): + raise ValueError(f"bit_width must be 2, 3, or 4, got {bit_width}.") + + arguments = TurboVecArgs(dim=vectors.shape[1], metric=metric_enum, bit_width=bit_width, calibrate=calibrate) + backend = cls(TurboQuantIndex(dim=_padded_dim(arguments.dim), bit_width=bit_width), arguments) + prepared = backend._prepare(vectors) + # turbovec needs at least two rows to fit a calibration. + if calibrate and len(prepared) > 1: + rng = np.random.default_rng(42) + sample = rng.choice(len(prepared), min(len(prepared), _CALIBRATION_SAMPLE_SIZE), replace=False) + backend.index.calibrate(prepared[sample]) + backend.index.add(prepared) + backend.positions = np.arange(len(prepared)) + return backend + + @property + def backend_type(self) -> Backend: + """The type of the backend.""" + return Backend.TURBOVEC + + @property + def dim(self) -> int: + """Get the dimension of the space.""" + return self.arguments.dim + + def __len__(self) -> int: + """Get the number of vectors.""" + return len(self.index) + + @classmethod + def load(cls: type[TurboVecBackend], path: Path) -> TurboVecBackend: + """Load the index from a path.""" + index_path = path / "index.tv" + arguments = TurboVecArgs.load(path / "arguments.json") + index = TurboQuantIndex.load(str(index_path)) + return cls(index, arguments=arguments, positions=np.load(path / "positions.npy")) + + def save(self, path: Path) -> None: + """Save the index to a path.""" + self.index.write(str(path / "index.tv")) + np.save(path / "positions.npy", self.positions) + self.arguments.dump(path / "arguments.json") + + def query(self, vectors: npt.NDArray, k: int) -> QueryResult: + """Query the backend and return results as tuples of keys and distances.""" + scores, indices = self.index.search(self._prepare(vectors), k=k) + # Inner products of unit vectors are cosine similarities. + return list(zip(self.positions[indices], 1.0 - scores)) + + def insert(self, vectors: npt.NDArray) -> None: + """Insert vectors into the backend.""" + self.positions = np.concatenate([self.positions, np.arange(len(self), len(self) + len(vectors))]) + self.index.add(self._prepare(vectors)) + + def _prepare(self, vectors: npt.NDArray) -> npt.NDArray: + """Normalize, zero-pad to the index dim and convert to contiguous float32, as turbovec requires.""" + vectors = normalize(np.asarray(vectors, dtype=np.float32)) + padding = _padded_dim(self.dim) - self.dim + return np.ascontiguousarray(np.pad(vectors, ((0, 0), (0, padding))), dtype=np.float32) + + def delete(self, indices: list[int]) -> None: + """Delete vectors at the given positions, shifting later positions down to stay aligned with the items.""" + deleted = np.sort(np.asarray(indices, dtype=np.int64)) + slots = np.flatnonzero(np.isin(self.positions, deleted)) + # Remove the highest slots first, so the last vector moved into a freed slot is never one being deleted. + for slot in slots[::-1]: + self.index.swap_remove(int(slot)) + self.positions[slot] = self.positions[-1] + self.positions = self.positions[:-1] + self.positions = self.positions - np.searchsorted(deleted, self.positions) + + 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 + + +def _padded_dim(dim: int) -> int: + """Round up to a multiple of 8, which turbovec requires; zero padding leaves inner products unchanged.""" + return -(-dim // 8) * 8 diff --git a/vicinity/datatypes.py b/vicinity/datatypes.py index f09db5e..c332b30 100644 --- a/vicinity/datatypes.py +++ b/vicinity/datatypes.py @@ -25,3 +25,4 @@ class Backend(str, Enum): FAISS = "faiss" USEARCH = "usearch" VOYAGER = "voyager" + TURBOVEC = "turbovec"