feat: add vector store and integer indexing (#46)

* feat: add vector store and integer indexing

* add saving and loading of vector store

* fix: bug in loading

* Add tests

* tests: add additional tests

Author: stephantul

tests/conftest.py

@@ -50,10 +50,7 @@ def query_vector() -> np.ndarray:
 
 
 # Create human-readable ids for each backend type
-BACKEND_IDS = [
-    f"{backend.name}-{index_type}" if index_type else backend.name
-    for backend, index_type in BACKEND_PARAMS
-]
+BACKEND_IDS = [f"{backend.name}-{index_type}" if index_type else backend.name for backend, index_type in BACKEND_PARAMS]
 
 
 @pytest.fixture(params=BACKEND_PARAMS)
@@ -63,9 +60,7 @@ def backend_type(request: pytest.FixtureRequest) -> Backend:
 
 
 @pytest.fixture(params=BACKEND_PARAMS, ids=BACKEND_IDS)
-def vicinity_instance(
-    request: pytest.FixtureRequest, items: list[str], vectors: np.ndarray
-) -> Vicinity:
+def vicinity_instance(request: pytest.FixtureRequest, items: list[str], vectors: np.ndarray) -> Vicinity:
     """Fixture providing a Vicinity instance for each backend type."""
     backend_type, index_type = request.param
     # Handle FAISS backend with specific FAISS index types
@@ -91,3 +86,30 @@ def vicinity_instance(
             )
 
     return Vicinity.from_vectors_and_items(vectors, items, backend_type=backend_type)
+
+
+@pytest.fixture(params=BACKEND_PARAMS, ids=BACKEND_IDS)
+def vicinity_instance_with_stored_vectors(
+    request: pytest.FixtureRequest, items: list[str], vectors: np.ndarray
+) -> Vicinity:
+    """Fixture providing a Vicinity instance for each backend type."""
+    backend_type, index_type = request.param
+    # Handle FAISS backend with specific FAISS index types
+    if backend_type == Backend.FAISS:
+        if index_type in ("pq", "ivfpq", "ivfpqr"):
+            # Use smaller values for pq indexes since the dataset is small
+            return Vicinity.from_vectors_and_items(
+                vectors, items, backend_type=backend_type, index_type=index_type, m=2, nbits=4, store_vectors=True
+            )
+        else:
+            return Vicinity.from_vectors_and_items(
+                vectors, items, backend_type=backend_type, index_type=index_type, nlist=2, nbits=32, store_vectors=True
+            )
+
+    return Vicinity.from_vectors_and_items(vectors, items, backend_type=backend_type, store_vectors=True)
+
+
+@pytest.fixture()
+def vicinity_with_basic_backend(vectors: np.ndarray, items: list[str]) -> Vicinity:
+    """Fixture providing a BasicBackend instance."""
+    return Vicinity.from_vectors_and_items(vectors, items, backend_type=Backend.BASIC, store_vectors=True)

tests/test_vicinity.py

@@ -131,8 +131,48 @@ def test_vicinity_save_and_load(tmp_path: Path, vicinity_instance: Vicinity) ->
     """
     save_path = tmp_path / "vicinity_data"
     vicinity_instance.save(save_path)
+    assert vicinity_instance.vector_store is None
 
-    Vicinity.load(save_path)
+    v = Vicinity.load(save_path)
+    assert v.vector_store is None
+
+
+def test_vicinity_save_and_load_vector_store(tmp_path: Path, vicinity_instance_with_stored_vectors: Vicinity) -> None:
+    """
+    Test Vicinity.save and Vicinity.load.
+
+    :param tmp_path: Temporary directory provided by pytest.
+    :param vicinity_instance: A Vicinity instance.
+    """
+    save_path = tmp_path / "vicinity_data"
+    vicinity_instance_with_stored_vectors.save(save_path)
+
+    assert (save_path / "store").exists()
+    assert (save_path / "store" / "vectors.npy").exists()
+
+    v = Vicinity.load(save_path)
+    assert v.vector_store is not None
+
+
+def test_index_vector_store(vicinity_with_basic_backend: Vicinity, vectors: np.ndarray) -> None:
+    """
+    Index vectors in the Vicinity instance.
+
+    :param vicinity_instance: A Vicinity instance.
+    :param vectors: Array of vectors to index.
+    """
+    v = vicinity_with_basic_backend.get_vector_by_index(0)
+    assert np.allclose(v, vectors[0])
+
+    idx = [0, 1, 2, 3, 4, 10]
+    v = vicinity_with_basic_backend.get_vector_by_index(idx)
+    assert np.allclose(v, vectors[idx])
+
+    with pytest.raises(ValueError):
+        vicinity_with_basic_backend.get_vector_by_index([10_000])
+
+    with pytest.raises(ValueError):
+        vicinity_with_basic_backend.get_vector_by_index([-1])
 
 
 def test_vicinity_insert_duplicate(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:

vicinity/backends/__init__.py

@@ -1,7 +1,7 @@
 from typing import Union
 
 from vicinity.backends.base import AbstractBackend
-from vicinity.backends.basic import BasicBackend
+from vicinity.backends.basic import BasicBackend, BasicVectorStore
 from vicinity.datatypes import Backend
 
 
@@ -39,4 +39,4 @@ def get_backend_class(backend: Union[Backend, str]) -> type[AbstractBackend]:
         return VoyagerBackend
 
 
-__all__ = ["get_backend_class", "AbstractBackend"]
+__all__ = ["get_backend_class", "AbstractBackend", "BasicVectorStore"]

vicinity/backends/basic.py

@@ -18,23 +18,58 @@ class BasicArgs(BaseArgs):
     metric: Metric = Metric.COSINE
 
 
-class BasicBackend(AbstractBackend[BasicArgs], ABC):
-    argument_class = BasicArgs
-    _vectors: npt.NDArray
-    supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN}
+class BasicVectorStore:
+    def __init__(self, *, vectors: npt.NDArray, **kwargs: Any) -> None:
+        """
+        A basic vector store that just stores vectors.
 
-    def __init__(self, arguments: BasicArgs) -> None:
-        """Initialize the backend."""
-        super().__init__(arguments)
+        Note that we use kwargs in order to use this class as a mixin.
 
-    def __len__(self) -> int:
-        """Get the number of vectors."""
-        return self.vectors.shape[0]
+        :param vectors: The vectors to store.
+        :param **kwargs: Additional arguments. These are passed on to the super class.
+        """
+        super().__init__(**kwargs)
+        self._vectors = vectors
 
-    @property
-    def backend_type(self) -> Backend:
-        """The type of the backend."""
-        return Backend.BASIC
+    def _update_precomputed_data(self) -> None:
+        """Update precomputed data based on the metric."""
+        # NOTE: this is a no-op in the base implementation.
+        return
+
+    def get_by_index(self, indices: list[int]) -> npt.NDArray:
+        """Get vectors by index."""
+        return self._vectors[indices]
+
+    def insert(self, vectors: npt.NDArray) -> None:
+        """Insert vectors into the vector space."""
+        self._vectors = np.vstack([self._vectors, vectors])
+        self._update_precomputed_data()
+
+    def delete(self, indices: list[int]) -> None:
+        """Deletes specific indices from the vector space."""
+        self._vectors = np.delete(self._vectors, indices, axis=0)
+        self._update_precomputed_data()
+
+    def save(self, folder: Path) -> None:
+        """Save the vectors to a path."""
+        path = folder / "vectors.npy"
+        with open(path, "wb") as f:
+            np.save(f, self._vectors)
+
+    @staticmethod
+    def _load_vectors(folder: Path) -> npt.NDArray:
+        """Load the vectors from a path."""
+        path = folder / "vectors.npy"
+        with open(path, "rb") as f:
+            vectors = np.load(f)
+
+        return vectors
+
+    @classmethod
+    def load(cls, folder: Path) -> BasicVectorStore:
+        """Load the vectors from a path."""
+        vectors = cls._load_vectors(folder)
+        return cls(vectors=vectors)
 
     @property
     def dim(self) -> int:
@@ -55,10 +90,24 @@ def vectors(self, x: Matrix) -> None:
         self._vectors = matrix
         self._update_precomputed_data()
 
-    @abstractmethod
-    def _update_precomputed_data(self) -> None:
-        """Update precomputed data based on the metric."""
-        raise NotImplementedError()
+
+class BasicBackend(BasicVectorStore, AbstractBackend[BasicArgs], ABC):
+    argument_class = BasicArgs
+    _vectors: npt.NDArray
+    supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN}
+
+    def __init__(self, vectors: npt.NDArray, arguments: BasicArgs) -> None:
+        """Initialize the backend."""
+        super().__init__(vectors=vectors, arguments=arguments)
+
+    def __len__(self) -> int:
+        """Get the number of vectors."""
+        return self.vectors.shape[0]
+
+    @property
+    def backend_type(self) -> Backend:
+        """The type of the backend."""
+        return Backend.BASIC
 
     @abstractmethod
     def _dist(self, x: npt.NDArray) -> npt.NDArray:
@@ -83,10 +132,8 @@ def from_vectors(cls, vectors: npt.NDArray, metric: Union[str, Metric] = "cosine
     @classmethod
     def load(cls, folder: Path) -> BasicBackend:
         """Load the vectors from a path."""
-        path = folder / "vectors.npy"
         arguments = BasicArgs.load(folder / "arguments.json")
-        with open(path, "rb") as f:
-            vectors = np.load(f)
+        vectors = cls._load_vectors(folder)
         if arguments.metric == Metric.COSINE:
             return CosineBasicBackend(vectors, arguments)
         elif arguments.metric == Metric.EUCLIDEAN:
@@ -96,10 +143,8 @@ def load(cls, folder: Path) -> BasicBackend:
 
     def save(self, folder: Path) -> None:
         """Save the vectors to a path."""
-        path = folder / "vectors.npy"
+        super().save(folder)
         self.arguments.dump(folder / "arguments.json")
-        with open(path, "wb") as f:
-            np.save(f, self._vectors)
 
     def threshold(
         self,
@@ -160,26 +205,12 @@ def query(
 
         return out
 
-    def insert(self, vectors: npt.NDArray) -> None:
-        """Insert vectors into the vector space."""
-        self._vectors = np.vstack([self._vectors, vectors])
-        self._update_precomputed_data()
-
-    def delete(self, indices: list[int]) -> None:
-        """Deletes specific indices from the vector space."""
-        self._vectors = np.delete(self._vectors, indices, axis=0)
-        self._update_precomputed_data()
-
 
 class CosineBasicBackend(BasicBackend):
     def __init__(self, vectors: npt.NDArray, arguments: BasicArgs) -> None:
         """Initialize the cosine basic backend."""
-        super().__init__(arguments)
-        self._vectors = normalize_or_copy(vectors)
-
-    def _update_precomputed_data(self) -> None:
-        """Update precomputed data for cosine similarity."""
-        pass
+        super().__init__(vectors=vectors, arguments=arguments)
+        self._vectors = normalize_or_copy(self._vectors)
 
     def _dist(self, x: npt.NDArray) -> npt.NDArray:
         """Compute cosine distance."""
@@ -197,21 +228,12 @@ def insert(self, vectors: npt.NDArray) -> None:
 class EuclideanBasicBackend(BasicBackend):
     def __init__(self, vectors: npt.NDArray, arguments: BasicArgs) -> None:
         """Initialize the Euclidean basic backend."""
-        super().__init__(arguments)
-        self._vectors = vectors
-        self._squared_norm_vectors: npt.NDArray | None = None
-        self._update_precomputed_data()
+        super().__init__(vectors=vectors, arguments=arguments)
+        self.squared_norm_vectors = (self._vectors**2).sum(1)
 
     def _update_precomputed_data(self) -> None:
         """Update precomputed data for Euclidean distance."""
-        self._squared_norm_vectors = (self._vectors**2).sum(1)
-
-    @property
-    def squared_norm_vectors(self) -> npt.NDArray:
-        """Return squared norms of vectors."""
-        if self._squared_norm_vectors is None:
-            self._squared_norm_vectors = (self._vectors**2).sum(1)
-        return self._squared_norm_vectors
+        self.squared_norm_vectors = (self._vectors**2).sum(1)
 
     def _dist(self, x: npt.NDArray) -> npt.NDArray:
         """Compute Euclidean distance."""