vectordb-bench 0.0.1__1-py3-none-any.whl

vectordb_bench/backend/clients/milvus/config.py

@@ -0,0 +1,123 @@
+from pydantic import BaseModel, SecretStr
+from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+
+
+class MilvusConfig(DBConfig, BaseModel):
+    uri: SecretStr | None = "http://localhost:19530"
+
+    def to_dict(self) -> dict:
+        return {"uri": self.uri.get_secret_value()}
+
+
+
+class MilvusIndexConfig(BaseModel):
+    """Base config for milvus"""
+
+    index: IndexType
+    metric_type: MetricType | None = None
+
+    def parse_metric(self) -> str:
+        if not self.metric_type:
+            return ""
+
+        if self.metric_type == MetricType.COSINE:
+            return MetricType.L2.value
+        return self.metric_type.value
+
+
+class AutoIndexConfig(MilvusIndexConfig, DBCaseConfig):
+    index: IndexType = IndexType.AUTOINDEX
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+        }
+
+class HNSWConfig(MilvusIndexConfig, DBCaseConfig):
+    M: int
+    efConstruction: int
+    ef: int | None = None
+    index: IndexType = IndexType.HNSW
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"M": self.M, "efConstruction": self.efConstruction},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"ef": self.ef},
+        }
+
+
+class DISKANNConfig(MilvusIndexConfig, DBCaseConfig):
+    search_list: int | None = None
+    index: IndexType = IndexType.DISKANN
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"search_list": self.search_list},
+        }
+
+
+class IVFFlatConfig(MilvusIndexConfig, DBCaseConfig):
+    nlist: int
+    nprobe: int | None = None
+    index: IndexType = IndexType.IVFFlat
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {"nlist": self.nlist},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {"nprobe": self.nprobe},
+        }
+
+
+class FLATConfig(MilvusIndexConfig, DBCaseConfig):
+    index: IndexType = IndexType.Flat
+
+    def index_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "index_type": self.index.value,
+            "params": {},
+        }
+
+    def search_param(self) -> dict:
+        return {
+            "metric_type": self.parse_metric(),
+            "params": {},
+        }
+
+_milvus_case_config = {
+    IndexType.AUTOINDEX: AutoIndexConfig,
+    IndexType.HNSW: HNSWConfig,
+    IndexType.DISKANN: DISKANNConfig,
+    IndexType.IVFFlat: IVFFlatConfig,
+    IndexType.Flat: FLATConfig,
+}
+

vectordb_bench/backend/clients/milvus/milvus.py

@@ -0,0 +1,182 @@
+"""Wrapper around the Milvus vector database over VectorDB"""
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Iterable, Type
+
+from pymilvus import Collection, utility
+from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusException
+
+from ..api import VectorDB, DBCaseConfig, DBConfig, IndexType
+from .config import MilvusConfig, _milvus_case_config
+
+
+log = logging.getLogger(__name__)
+
+
+class Milvus(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        collection_name: str = "VectorDBBenchCollection",
+        drop_old: bool = False,
+        name: str = "Milvus",
+    ):
+        """Initialize wrapper around the milvus vector database."""
+        self.name = name
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.collection_name = collection_name
+
+        self._primary_field = "pk"
+        self._scalar_field = "id"
+        self._vector_field = "vector"
+        self._index_name = "vector_idx"
+
+        from pymilvus import connections
+        connections.connect(**self.db_config, timeout=30)
+        if drop_old and utility.has_collection(self.collection_name):
+            log.info(f"{self.name} client drop_old collection: {self.collection_name}")
+            utility.drop_collection(self.collection_name)
+
+        if not utility.has_collection(self.collection_name):
+            fields = [
+                FieldSchema(self._primary_field, DataType.INT64, is_primary=True),
+                FieldSchema(self._scalar_field, DataType.INT64),
+                FieldSchema(self._vector_field, DataType.FLOAT_VECTOR, dim=dim)
+            ]
+
+            log.info(f"{self.name} create collection: {self.collection_name}")
+
+            # Create the collection
+            coll = Collection(
+                name=self.collection_name,
+                schema=CollectionSchema(fields),
+                consistency_level="Session",
+            )
+
+            #  self._pre_load(coll)
+
+        connections.disconnect("default")
+
+    @classmethod
+    def config_cls(cls) -> Type[DBConfig]:
+        return MilvusConfig
+
+    @classmethod
+    def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        return _milvus_case_config.get(index_type)
+
+
+    @contextmanager
+    def init(self) -> None:
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+        from pymilvus import connections
+        self.col: Collection | None = None
+
+        connections.connect(**self.db_config, timeout=60)
+        # Grab the existing colection with connections
+        self.col = Collection(self.collection_name)
+
+        yield
+        connections.disconnect("default")
+
+    def _pre_load(self, coll: Collection):
+        if not coll.has_index(index_name=self._index_name):
+            log.info(f"{self.name} create index and load")
+            try:
+                coll.create_index(
+                    self._vector_field,
+                    self.case_config.index_param(),
+                    index_name=self._index_name,
+                )
+
+                coll.load()
+            except Exception as e:
+                log.warning(f"{self.name} pre load error: {e}")
+                raise e from None
+
+    def _optimize(self):
+        log.info(f"{self.name} optimizing before search")
+        try:
+            self.col.flush()
+            self.col.compact()
+            self.col.wait_for_compaction_completed()
+
+            # wait for index done and load refresh
+            self.col.create_index(
+                self._vector_field,
+                self.case_config.index_param(),
+                index_name=self._index_name,
+            )
+            utility.wait_for_index_building_complete(self.collection_name)
+            self.col.load()
+            #  self.col.load(_refresh=True)
+            #  utility.wait_for_loading_complete(self.collection_name)
+            #  import time; time.sleep(10)
+        except Exception as e:
+            log.warning(f"{self.name} optimize error: {e}")
+            raise e from None
+
+    def ready_to_load(self):
+        assert self.col, "Please call self.init() before"
+        self._pre_load(self.col)
+        pass
+
+    def ready_to_search(self):
+        assert self.col, "Please call self.init() before"
+        self._optimize()
+
+    def insert_embeddings(
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> int:
+        """Insert embeddings into Milvus. should call self.init() first"""
+        # use the first insert_embeddings to init collection
+        assert self.col is not None
+        insert_data = [
+                metadata,
+                metadata,
+                embeddings,
+        ]
+
+        try:
+            res = self.col.insert(insert_data, **kwargs)
+            return len(res.primary_keys)
+        except MilvusException as e:
+            log.warning("Failed to insert data")
+            raise e from None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+    ) -> list[int]:
+        """Perform a search on a query embedding and return results."""
+        assert self.col is not None
+
+        expr = f"{self._scalar_field} {filters.get('metadata')}" if filters else ""
+
+        # Perform the search.
+        res = self.col.search(
+            data=[query],
+            anns_field=self._vector_field,
+            param=self.case_config.search_param(),
+            limit=k,
+            expr=expr,
+        )
+
+        # Organize results.
+        ret = [result.id for result in res[0]]
+        return ret

vectordb_bench/backend/clients/pinecone/config.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel, SecretStr
+from ..api import DBConfig
+
+
+class PineconeConfig(DBConfig, BaseModel):
+    api_key: SecretStr | None = None
+    environment: SecretStr | None = None
+    index_name: str
+
+    def to_dict(self) -> dict:
+        return {
+            "api_key": self.api_key.get_secret_value(),
+            "environment": self.environment.get_secret_value(),
+            "index_name": self.index_name,
+        }

vectordb_bench/backend/clients/pinecone/pinecone.py

@@ -0,0 +1,113 @@
+"""Wrapper around the Pinecone vector database over VectorDB"""
+
+import logging
+from contextlib import contextmanager
+from typing import Any, Type
+
+from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
+from .config import PineconeConfig
+
+
+log = logging.getLogger(__name__)
+
+PINECONE_MAX_NUM_PER_BATCH = 1000
+PINECONE_MAX_SIZE_PER_BATCH = 2 * 1024 * 1024 # 2MB
+
+class Pinecone(VectorDB):
+    def __init__(
+        self,
+        dim,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        drop_old: bool = False,
+    ):
+        """Initialize wrapper around the milvus vector database."""
+        self.index_name = db_config["index_name"]
+        self.api_key = db_config["api_key"]
+        self.environment = db_config["environment"]
+        self.batch_size = int(min(PINECONE_MAX_SIZE_PER_BATCH / (dim * 5), PINECONE_MAX_NUM_PER_BATCH))
+        # Pincone will make connections with server while import
+        # so place the import here.
+        import pinecone
+        pinecone.init(
+            api_key=self.api_key, environment=self.environment)
+        if drop_old:
+            list_indexes = pinecone.list_indexes()
+            if self.index_name in list_indexes:
+                index = pinecone.Index(self.index_name)
+                index_dim = index.describe_index_stats()["dimension"]
+                if (index_dim != dim):
+                    raise ValueError(
+                        f"Pinecone index {self.index_name} dimension mismatch, expected {index_dim} got {dim}")
+                log.info(
+                    f"Pinecone client delete old index: {self.index_name}")
+                index.delete(delete_all=True)
+                index.close()
+            else:
+                raise ValueError(
+                    f"Pinecone index {self.index_name} does not exist")
+
+        self._metadata_key = "meta"
+
+    @classmethod
+    def config_cls(cls) -> Type[DBConfig]:
+        return PineconeConfig
+
+    @classmethod
+    def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        return EmptyDBCaseConfig
+
+    @contextmanager
+    def init(self) -> None:
+        import pinecone
+        pinecone.init(
+            api_key=self.api_key, environment=self.environment)
+        self.index = pinecone.Index(self.index_name)
+        yield
+        self.index.close()
+
+    def ready_to_load(self):
+        pass
+
+    def ready_to_search(self):
+        pass
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+    ) -> list[str]:
+        assert len(embeddings) == len(metadata)
+        for batch_start_offset in range(0, len(embeddings), self.batch_size):
+            batch_end_offset = min(batch_start_offset + self.batch_size, len(embeddings))
+            insert_datas = []
+            for i in range(batch_start_offset, batch_end_offset):
+                insert_data = (str(metadata[i]), embeddings[i], {
+                            self._metadata_key: metadata[i]})
+                insert_datas.append(insert_data)
+            self.index.upsert(insert_datas)
+        return len(embeddings)
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> list[tuple[int, float]]:
+        if filters is None:
+            pinecone_filters = {}
+        else:
+            pinecone_filters = {self._metadata_key: {"$gte": filters["id"]}}
+        try:
+            res = self.index.query(
+                top_k=k,
+                vector=query,
+                filter=pinecone_filters,
+            )['matches']
+        except Exception as e:
+            print(f"Error querying index: {e}")
+            raise e
+        id_res = [int(one_res['id']) for one_res in res]
+        return id_res

vectordb_bench/backend/clients/qdrant_cloud/config.py

@@ -0,0 +1,16 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBConfig
+
+
+class QdrantConfig(DBConfig, BaseModel):
+    url: SecretStr | None = None
+    api_key: SecretStr | None = None
+    prefer_grpc: bool = True
+
+    def to_dict(self) -> dict:
+        return {
+            "url": self.url.get_secret_value(),
+            "api_key": self.api_key.get_secret_value(),
+            "prefer_grpc": self.prefer_grpc,
+        }

vectordb_bench/backend/clients/qdrant_cloud/qdrant_cloud.py

@@ -0,0 +1,169 @@
+"""Wrapper around the QdrantCloud vector database over VectorDB"""
+
+import logging
+import time
+from contextlib import contextmanager
+from typing import Any, Type
+
+from ..api import VectorDB, DBConfig, DBCaseConfig, EmptyDBCaseConfig, IndexType
+from .config import QdrantConfig
+from qdrant_client.http.models import (
+    CollectionStatus,
+    Distance,
+    VectorParams,
+    PayloadSchemaType,
+    Batch,
+    Filter,
+    FieldCondition,
+    Range,
+)
+
+from qdrant_client import QdrantClient
+
+
+log = logging.getLogger(__name__)
+
+
+class QdrantCloud(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        collection_name: str = "QdrantCloudCollection",
+        drop_old: bool = False,
+    ):
+        """Initialize wrapper around the QdrantCloud vector database."""
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.collection_name = collection_name
+
+        self._primary_field = "pk"
+        self._vector_field = "vector"
+
+        tmp_client = QdrantClient(**self.db_config)
+        if drop_old:
+            log.info(f"QdrantCloud client drop_old collection: {self.collection_name}")
+            tmp_client.delete_collection(self.collection_name)
+
+        self._create_collection(dim, tmp_client)
+        tmp_client = None
+
+    @classmethod
+    def config_cls(cls) -> Type[DBConfig]:
+        return QdrantConfig
+
+    @classmethod
+    def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        return EmptyDBCaseConfig
+
+    @contextmanager
+    def init(self) -> None:
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+        self.qdrant_client = QdrantClient(**self.db_config)
+        yield
+        self.qdrant_client = None
+        del(self.qdrant_client)
+
+    def ready_to_load(self):
+        pass
+
+
+    def ready_to_search(self):
+        assert self.qdrant_client, "Please call self.init() before"
+        # wait for vectors to be fully indexed
+        SECONDS_WAITING_FOR_INDEXING_API_CALL = 5
+        try:
+            while True:
+                info = self.qdrant_client.get_collection(self.collection_name)
+                time.sleep(SECONDS_WAITING_FOR_INDEXING_API_CALL)
+                if info.status != CollectionStatus.GREEN:
+                    continue
+                if info.status == CollectionStatus.GREEN:
+                    log.info(f"Stored vectors: {info.vectors_count}, Indexed vectors: {info.indexed_vectors_count}, Collection status: {info.indexed_vectors_count}")
+                    return
+        except Exception as e:
+            log.warning(f"QdrantCloud ready to search error: {e}")
+            raise e from None
+
+    def _create_collection(self, dim, qdrant_client: int):
+        log.info(f"Create collection: {self.collection_name}")
+
+        try:
+            qdrant_client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=VectorParams(size=dim, distance=Distance.EUCLID)
+            )
+
+            qdrant_client.create_payload_index(
+                collection_name=self.collection_name,
+                field_name=self._primary_field,
+                field_schema=PayloadSchemaType.INTEGER,
+            )
+
+        except Exception as e:
+            if "already exists!" in str(e):
+                return
+            log.warning(f"Failed to create collection: {self.collection_name} error: {e}")
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> list[str]:
+        """Insert embeddings into Milvus. should call self.init() first"""
+        assert self.qdrant_client is not None
+        try:
+            # TODO: counts
+            _ = self.qdrant_client.upsert(
+                collection_name=self.collection_name,
+                wait=True,
+                points=Batch(ids=metadata, payloads=[{self._primary_field: v} for v in metadata], vectors=embeddings)
+            )
+
+            return len(metadata)
+        except Exception as e:
+            log.info(f"Failed to insert data, {e}")
+            raise e from None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> list[int]:
+        """Perform a search on a query embedding and return results with score.
+        Should call self.init() first.
+        """
+        assert self.qdrant_client is not None
+
+        f = None
+        if filters:
+            f = Filter(
+                must=[FieldCondition(
+                    key = self._primary_field,
+                    range = Range(
+                        gt=filters.get('id'),
+                    ),
+                )]
+            )
+
+        res = self.qdrant_client.search(
+            collection_name=self.collection_name,
+            query_vector=query,
+            limit=k,
+            query_filter=f,
+            #  with_payload=True,
+        ),
+
+        ret = [result.id for result in res[0]]
+        return ret

vectordb_bench/backend/clients/weaviate_cloud/config.py

@@ -0,0 +1,45 @@
+from pydantic import BaseModel, SecretStr
+import weaviate
+
+from ..api import DBConfig, DBCaseConfig, MetricType
+
+
+class WeaviateConfig(DBConfig, BaseModel):
+    url: SecretStr | None = None
+    api_key: SecretStr | None = None
+
+    def to_dict(self) -> dict:
+        return {
+            "url": self.url.get_secret_value(),
+            "auth_client_secret": weaviate.AuthApiKey(api_key=self.api_key.get_secret_value()),
+        }
+
+
+class WeaviateIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+    ef: int | None = -1
+    efConstruction: int | None = None
+    maxConnections: int | None = None
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "l2-squared"
+        elif self.metric_type == MetricType.IP:
+            return "dot"
+        return "cosine"
+
+    def index_param(self) -> dict:
+        if self.maxConnections is not None and self.efConstruction is not None:
+            params = {
+                "distance": self.parse_metric(),
+                "maxConnections": self.maxConnections,
+                "efConstruction": self.efConstruction,
+            }
+        else:
+            params = {"distance": self.parse_metric()}
+        return params
+
+    def search_param(self) -> dict:
+        return {
+            "ef": self.ef,
+        }