Class BasicVectorStore

vicinity/backends/basic.py:21–95

class BasicVectorStore:
    def __init__(self, *, vectors: npt.NDArray, **kwargs: Any) -> None:
        """
        A basic vector store that just stores vectors.

        Note that we use kwargs in order to use this class as a mixin.

        :param vectors: The vectors to store.
        :param **kwargs: Additional arguments. These are passed on to the super class.
        """
        super().__init__(**kwargs)
        self._vectors = vectors

    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, path: Path) -> None:
        """Save the vectors to a path."""
        path = path / "vectors.npy"
        with open(path, "wb") as f:
            np.save(f, self._vectors)

    @staticmethod
    def _load_vectors(path: Path) -> npt.NDArray:
        """Load the vectors from a path."""
        path = path / "vectors.npy"
        with open(path, "rb") as f:
            vectors = np.load(f)

        return vectors

    @classmethod
    def load(cls, path: Path) -> BasicVectorStore:
        """Load the vectors from a path."""
        vectors = cls._load_vectors(path)
        return cls(vectors=vectors)

    @property
    def dim(self) -> int:
        """The size of the space."""
        return self.vectors.shape[1]