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

vectordb_bench/__init__.py

@@ -0,0 +1,30 @@
+import environs
+import inspect
+import pathlib
+from . import log_util
+
+
+env = environs.Env()
+env.read_env(".env")
+
+class config:
+    LOG_LEVEL = env.str("LOG_LEVEL", "INFO")
+
+    DEFAULT_DATASET_URL = env.str("DEFAULT_DATASET_URL", "assets.zilliz.com/benchmark/")
+    DATASET_LOCAL_DIR = env.path("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset")
+    NUM_PER_BATCH = env.int("NUM_PER_BATCH", 5000)
+
+    DROP_OLD = env.bool("DROP_OLD", True)
+    USE_SHUFFLED_DATA = env.bool("USE_SHUFFLED_DATA", True)
+
+    RESULTS_LOCAL_DIR = pathlib.Path(__file__).parent.joinpath("results")
+    CASE_TIMEOUT_IN_SECOND = 24 * 60 * 60
+
+
+    def display(self) -> str:
+        tmp = [i for i in inspect.getmembers(self)
+            if not inspect.ismethod(i[1]) and not i[0].startswith('_') \
+        ]
+        return tmp
+
+log_util.init(config.LOG_LEVEL)

vectordb_bench/__main__.py

@@ -0,0 +1,39 @@
+import traceback
+import logging
+import subprocess
+import os
+from . import config
+
+log = logging.getLogger("vectordb_bench")
+
+
+def main():
+    log.info(f"all configs: {config().display()}")
+    run_streamlit()
+
+
+def run_streamlit():
+    cmd = [
+        "streamlit",
+        "run",
+        f"{os.path.dirname(__file__)}/frontend/vdb_benchmark.py",
+        "--logger.level",
+        "info",
+        "--theme.base",
+        "light",
+        "--theme.primaryColor",
+        "#3670F2",
+        "--theme.secondaryBackgroundColor",
+        "#F0F2F6",
+    ]
+    log.debug(f"cmd: {cmd}")
+    try:
+        subprocess.run(cmd, check=True)
+    except KeyboardInterrupt:
+        log.info("exit streamlit...")
+    except Exception as e:
+        log.warning(f"exit, err={e}\nstack trace={traceback.format_exc(chain=True)}")
+
+
+if __name__ == "__main__":
+    main()

vectordb_bench/backend/assembler.py

@@ -0,0 +1,57 @@
+from .cases import type2case, CaseLabel
+from .task_runner import CaseRunner, RunningStatus, TaskRunner
+from ..models import TaskConfig
+from ..backend.clients import EmptyDBCaseConfig
+import logging
+
+
+log = logging.getLogger(__name__)
+
+
+class Assembler:
+    @classmethod
+    def assemble(cls, run_id , task: TaskConfig) -> CaseRunner:
+        c_cls = type2case.get(task.case_config.case_id)
+
+        c = c_cls()
+        if type(task.db_case_config) != EmptyDBCaseConfig:
+            task.db_case_config.metric_type = c.dataset.data.metric_type
+
+        runner = CaseRunner(
+            run_id=run_id,
+            config=task,
+            ca=c,
+            status=RunningStatus.PENDING,
+        )
+
+        return runner
+
+    @classmethod
+    def assemble_all(cls, run_id: str, task_label: str, tasks: list[TaskConfig]) -> TaskRunner:
+        """group by case type, db, and case dataset"""
+        runners = [cls.assemble(run_id, task) for task in tasks]
+        load_runners = [r for r in runners if r.ca.label == CaseLabel.Load]
+        perf_runners = [r for r in runners if r.ca.label == CaseLabel.Performance]
+
+        # group by db
+        db2runner = {}
+        for r in perf_runners:
+            db = r.config.db
+            if db not in db2runner:
+                db2runner[db] = []
+            db2runner[db].append(r)
+
+        # sort by dataset size
+        for k in db2runner.keys():
+            db2runner[k].sort(key=lambda x:x.ca.dataset.data.size)
+
+        all_runners = []
+        all_runners.extend(load_runners)
+        for v in db2runner.values():
+            all_runners.extend(v)
+
+        return TaskRunner(
+            run_id=run_id,
+            task_label=task_label,
+            case_runners=all_runners,
+        )

vectordb_bench/backend/cases.py

@@ -0,0 +1,124 @@
+import logging
+from enum import Enum, auto
+
+from . import dataset as ds
+from ..base import BaseModel
+from ..models import CaseType
+
+
+log = logging.getLogger(__name__)
+
+
+class CaseLabel(Enum):
+    Load = auto()
+    Performance = auto()
+
+
+class Case(BaseModel):
+    """ Undifined case
+
+    Fields:
+        case_id(CaseType): default 11 case type plus one custom cases.
+        label(CaseLabel): performance or load.
+        dataset(DataSet): dataset for this case runner.
+        filter_rate(float | None): one of 99% | 1% | None
+        filters(dict | None): filters for search
+    """
+
+    case_id: CaseType
+    label: CaseLabel
+    dataset: ds.DataSet
+
+    filter_rate: float | None
+
+    @property
+    def filters(self) -> dict | None:
+        if self.filter_rate is not None:
+            ID = round(self.filter_rate * self.dataset.data.size)
+            return {
+                "metadata": f">={ID}",
+                "id": ID,
+            }
+
+        return None
+
+
+class CapacityCase(Case, BaseModel):
+    label: CaseLabel = CaseLabel.Load
+    filter_rate: float | int | None = None
+
+class PerformanceCase(Case, BaseModel):
+    label: CaseLabel = CaseLabel.Performance
+    filter_rate: float | int | None = None
+
+class CapacityLDimCase(CapacityCase):
+    case_id: CaseType = CaseType.CapacityLDim
+    dataset: ds.DataSet = ds.get(ds.Name.GIST, ds.Label.SMALL)
+
+class CapacitySDimCase(CapacityCase):
+    case_id: CaseType = CaseType.CapacitySDim
+    dataset: ds.DataSet = ds.get(ds.Name.SIFT, ds.Label.SMALL)
+
+class PerformanceLZero(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceLZero
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.LARGE)
+
+class PerformanceMZero(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceMZero
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.MEDIUM)
+
+class PerformanceSZero(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceSZero
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.SMALL)
+
+class PerformanceLLow(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceLLow
+    filter_rate: float | int | None = 0.01
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.LARGE)
+
+class PerformanceMLow(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceMLow
+    filter_rate: float | int | None = 0.01
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.MEDIUM)
+
+class PerformanceSLow(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceSLow
+    filter_rate: float | int | None = 0.01
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.SMALL)
+
+class PerformanceLHigh(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceLHigh
+    filter_rate: float | int | None = 0.99
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.LARGE)
+
+class PerformanceMHigh(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceMHigh
+    filter_rate: float | int | None = 0.99
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.MEDIUM)
+
+class PerformanceSHigh(PerformanceCase):
+    case_id: CaseType = CaseType.PerformanceSLow
+    filter_rate: float | int | None = 0.99
+    dataset: ds.DataSet = ds.get(ds.Name.Cohere, ds.Label.SMALL)
+
+class Performance100M(PerformanceCase):
+    case_id: CaseType = CaseType.Performance100M
+    filter_rate: float | int | None = None
+    dataset: ds.DataSet = ds.get(ds.Name.LAION, ds.Label.LARGE)
+
+type2case = {
+    CaseType.CapacityLDim: CapacityLDimCase,
+    CaseType.CapacitySDim: CapacitySDimCase,
+
+    CaseType.PerformanceLZero: PerformanceLZero,
+    CaseType.PerformanceMZero: PerformanceMZero,
+    CaseType.PerformanceSZero: PerformanceSZero,
+
+    CaseType.PerformanceLLow: PerformanceLLow,
+    CaseType.PerformanceMLow: PerformanceMLow,
+    CaseType.PerformanceSLow: PerformanceSLow,
+    CaseType.PerformanceLHigh: PerformanceLHigh,
+    CaseType.PerformanceMHigh: PerformanceMHigh,
+    CaseType.PerformanceSHigh: PerformanceSHigh,
+    CaseType.Performance100M: Performance100M,
+}

vectordb_bench/backend/clients/__init__.py

@@ -0,0 +1,57 @@
+from enum import Enum
+from typing import Type
+from .api import (
+    VectorDB,
+    DBConfig,
+    DBCaseConfig,
+    EmptyDBCaseConfig,
+    IndexType,
+    MetricType,
+)
+
+from .milvus.milvus import Milvus
+from .elastic_cloud.elastic_cloud import ElasticCloud
+from .pinecone.pinecone import Pinecone
+from .weaviate_cloud.weaviate_cloud import WeaviateCloud
+from .qdrant_cloud.qdrant_cloud import QdrantCloud
+from .zilliz_cloud.zilliz_cloud import ZillizCloud
+
+
+class DB(Enum):
+    """Database types
+
+    Examples:
+        >>> DB.Milvus
+        <DB.Milvus: 'Milvus'>
+        >>> DB.Milvus.value
+        "Milvus"
+        >>> DB.Milvus.name
+        "Milvus"
+    """
+
+    Milvus = "Milvus"
+    ZillizCloud = "ZillizCloud"
+    Pinecone = "Pinecone"
+    ElasticCloud = "ElasticCloud"
+    QdrantCloud = "QdrantCloud"
+    WeaviateCloud = "WeaviateCloud"
+
+
+    @property
+    def init_cls(self) -> Type[VectorDB]:
+        return db2client.get(self)
+
+
+db2client = {
+    DB.Milvus: Milvus,
+    DB.ZillizCloud: ZillizCloud,
+    DB.WeaviateCloud: WeaviateCloud,
+    DB.ElasticCloud: ElasticCloud,
+    DB.QdrantCloud: QdrantCloud,
+    DB.Pinecone: Pinecone,
+}
+
+
+__all__ = [
+    "DB", "VectorDB", "DBConfig", "DBCaseConfig", "IndexType", "MetricType", "EmptyDBCaseConfig",
+]

vectordb_bench/backend/clients/api.py

@@ -0,0 +1,179 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any, Type
+from contextlib import contextmanager
+
+from pydantic import BaseModel
+
+
+class MetricType(str, Enum):
+    L2 = "L2"
+    COSINE = "COSINE"
+    IP = "IP"
+
+
+class IndexType(str, Enum):
+    HNSW = "HNSW"
+    DISKANN = "DISKANN"
+    IVFFlat = "IVF_FLAT"
+    Flat = "FLAT"
+    AUTOINDEX = "AUTOINDEX"
+    ES_HNSW = "hnsw"
+
+
+class DBConfig(ABC, BaseModel):
+    """DBConfig contains the connection info of vector database
+
+    Args:
+        db_label(str): label to distinguish different types of DB of the same database.
+
+            MilvusConfig.db_label = 2c8g
+            MilvusConfig.db_label = 16c64g
+            ZillizCloudConfig.db_label = 1cu-perf
+    """
+
+    db_label: str | None = None
+
+    @abstractmethod
+    def to_dict(self) -> dict:
+        raise NotImplementedError
+
+
+class DBCaseConfig(ABC):
+    """Case specific vector database configs, usually uesed for index params like HNSW"""
+    @abstractmethod
+    def index_param(self) -> dict:
+        raise NotImplementedError
+
+    @abstractmethod
+    def search_param(self) -> dict:
+        raise NotImplementedError
+
+
+class EmptyDBCaseConfig(BaseModel, DBCaseConfig):
+    """EmptyDBCaseConfig will be used if the vector database has no case specific configs"""
+    null: str | None = None
+    def index_param(self) -> dict:
+        return {}
+
+    def search_param(self) -> dict:
+        return {}
+
+
+class VectorDB(ABC):
+    """Each VectorDB will be __init__ once for one case, the object will be copied into multiple processes.
+
+    In each process, the benchmark cases ensure VectorDB.init() calls before any other methods operations
+
+    insert_embeddings, search_embedding, and, ready_to_search will be timed for each call.
+
+    Examples:
+        >>> milvus = Milvus()
+        >>> with milvus.init():
+        >>>     milvus.insert_embeddings()
+        >>>     milvus.search_embedding()
+    """
+
+    @abstractmethod
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig | None,
+        collection_name: str,
+        drop_old: bool = False,
+        **kwargs
+    ) -> None:
+        """Initialize wrapper around the vector database client
+
+        Args:
+            dim(int): the dimension of the dataset
+            db_config(dict): configs to establish connections with the vector database
+            db_case_config(DBCaseConfig | None): case specific configs for indexing and searching
+            drop_old(bool): whether to drop the existing collection of the dataset.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def config_cls(self) -> Type[DBConfig]:
+        raise NotImplementedError
+
+
+    @classmethod
+    @abstractmethod
+    def case_config_cls(self, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        raise NotImplementedError
+
+
+    @abstractmethod
+    @contextmanager
+    def init(self) -> None:
+        """ create and destory connections to database.
+
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        kwargs: Any,
+    ) -> int:
+        """Insert the embeddings to the vector database. The default number of embeddings for
+        each insert_embeddings is 5000.
+
+        Args:
+            embeddings(list[list[float]]): list of embedding to add to the vector database.
+            metadatas(list[int]): metadata associated with the embeddings, for filtering.
+            kwargs(Any): vector database specific parameters.
+
+        Returns:
+            int: inserted data count
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+    ) -> list[int]:
+        """Get k most similar embeddings to query vector.
+
+        Args:
+            query(list[float]): query embedding to look up documents similar to.
+            k(int): Number of most similar embeddings to return. Defaults to 100.
+            filters(dict, optional): filtering expression to filter the data while searching.
+
+        Returns:
+            list[int]: list of k most similar embeddings IDs to the query embedding.
+        """
+        raise NotImplementedError
+
+    # TODO: remove
+    @abstractmethod
+    def ready_to_search(self):
+        """ready_to_search will be called between insertion and search in performance cases.
+
+        Should be blocked until the vectorDB is ready to be tested on
+        heavy performance cases.
+
+        Time(insert the dataset) + Time(ready_to_search) will be recorded as "load_duration" metric
+        """
+        raise NotImplementedError
+
+    # TODO: remove
+    @abstractmethod
+    def ready_to_load(self):
+        """ready_to_load will be called before load in load cases.
+
+        Should be blocked until the vectorDB is ready to be tested on
+        heavy load cases.
+        """
+        raise NotImplementedError

vectordb_bench/backend/clients/elastic_cloud/config.py

@@ -0,0 +1,56 @@
+from enum import Enum
+from pydantic import SecretStr, BaseModel
+
+from ..api import DBConfig, DBCaseConfig, MetricType, IndexType
+
+
+class ElasticsearchConfig(DBConfig, BaseModel):
+    cloud_id: SecretStr
+    password: SecretStr | None = None
+
+    def to_dict(self) -> dict:
+        return {
+            "cloud_id": self.cloud_id.get_secret_value(),
+            "basic_auth": ("elastic", self.password.get_secret_value()),
+        }
+
+
+class ESElementType(str, Enum):
+    float = "float"  # 4 byte
+    byte = "byte"  # 1 byte, -128 to 127
+
+
+class ElasticsearchIndexConfig(BaseModel, DBCaseConfig):
+    element_type: ESElementType = ESElementType.float
+    index: IndexType = IndexType.ES_HNSW  # ES only support 'hnsw'
+
+    metric_type: MetricType | None = None
+    efConstruction: int | None = None
+    M: int | None = None
+    num_candidates: int | None = None
+
+    def parse_metric(self) -> str:
+        if self.metric_type == MetricType.L2:
+            return "l2_norm"
+        elif self.metric_type == MetricType.IP:
+            return "dot_product"
+        return "cosine"
+
+    def index_param(self) -> dict:
+        params = {
+            "type": "dense_vector",
+            "index": True,
+            "element_type": self.element_type.value,
+            "similarity": self.parse_metric(),
+            "index_options": {
+                "type": self.index.value,
+                "m": self.M,
+                "ef_construction": self.efConstruction
+            }
+        }
+        return params
+
+    def search_param(self) -> dict:
+        return {
+            "num_candidates": self.num_candidates,
+        }

vectordb_bench/backend/clients/elastic_cloud/elastic_cloud.py

@@ -0,0 +1,152 @@
+import logging
+from contextlib import contextmanager
+from typing import Iterable, Type
+from ..api import VectorDB, DBCaseConfig, DBConfig, IndexType
+from .config import ElasticsearchIndexConfig, ElasticsearchConfig
+from elasticsearch.helpers import bulk
+
+
+for logger in ("elasticsearch", "elastic_transport"):
+    logging.getLogger(logger).setLevel(logging.WARNING)
+
+log = logging.getLogger(__name__)
+
+class ElasticCloud(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: ElasticsearchIndexConfig,
+        indice: str = "vdb_bench_indice",  # must be lowercase
+        id_col_name: str = "id",
+        vector_col_name: str = "vector",
+        drop_old: bool = False,
+    ):
+        self.dim = dim
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.indice = indice
+        self.id_col_name = id_col_name
+        self.vector_col_name = vector_col_name
+
+        from elasticsearch import Elasticsearch
+
+        client = Elasticsearch(**self.db_config)
+
+        if drop_old:
+            log.info(f"Elasticsearch client drop_old indices: {self.indice}")
+            is_existed_res = client.indices.exists(index=self.indice)
+            if is_existed_res.raw:
+                client.indices.delete(index=self.indice)
+            self._create_indice(client)
+
+
+    @classmethod
+    def config_cls(cls) -> Type[DBConfig]:
+        return ElasticsearchConfig
+
+
+    @classmethod
+    def case_config_cls(cls, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
+        return ElasticsearchIndexConfig
+
+
+    @contextmanager
+    def init(self) -> None:
+        """connect to elasticsearch"""
+        from elasticsearch import Elasticsearch
+        self.client = Elasticsearch(**self.db_config, request_timeout=30)
+
+        yield
+        # self.client.transport.close()
+        self.client = None
+        del(self.client)
+
+    def _create_indice(self, client) -> None:
+        mappings = {
+            "properties": {
+                self.id_col_name: {"type": "integer"},
+                self.vector_col_name: {
+                    "dims": self.dim,
+                    **self.case_config.index_param(),
+                },
+            }
+        }
+
+        try:
+            client.indices.create(index=self.indice, mappings=mappings)
+        except Exception as e:
+            log.warning(f"Failed to create indice: {self.indice} error: {str(e)}")
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+    ) -> int:
+        """Insert the embeddings to the elasticsearch."""
+        assert self.client is not None, "should self.init() first"
+
+        insert_data = [
+            {
+                "_index": self.indice,
+                "_source": {
+                    self.id_col_name: metadata[i],
+                    self.vector_col_name: embeddings[i],
+                },
+            }
+            for i in range(len(embeddings))
+        ]
+        try:
+            bulk_insert_res = bulk(self.client, insert_data)
+            return bulk_insert_res[0]
+        except Exception as e:
+            log.warning(f"Failed to insert data: {self.indice} error: {str(e)}")
+            raise e from None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+    ) -> list[int]:
+        """Get k most similar embeddings to query vector.
+
+        Args:
+            query(list[float]): query embedding to look up documents similar to.
+            k(int): Number of most similar embeddings to return. Defaults to 100.
+            filters(dict, optional): filtering expression to filter the data while searching.
+
+        Returns:
+            list[tuple[int, float]]: list of k most similar embeddings in (id, score) tuple to the query embedding.
+        """
+        assert self.client is not None, "should self.init() first"
+        # is_existed_res = self.client.indices.exists(index=self.indice)
+        # assert is_existed_res.raw == True, "should self.init() first"
+
+        knn = {
+            "field": self.vector_col_name,
+            "k": k,
+            "num_candidates": self.case_config.num_candidates,
+            "filter": [{"range": {self.id_col_name: {"gt": filters["id"]}}}]
+            if filters
+            else [],
+            "query_vector": query,
+        }
+        size = k
+        try:
+            search_res = self.client.search(index=self.indice, knn=knn, size=size)
+            res = [d["_source"][self.id_col_name] for d in search_res["hits"]["hits"]]
+
+            return res
+        except Exception as e:
+            log.warning(f"Failed to search: {self.indice} error: {str(e)}")
+            raise e from None
+
+    def ready_to_search(self):
+        """ready_to_search will be called between insertion and search in performance cases."""
+        pass
+
+    def ready_to_load(self):
+        """ready_to_load will be called before load in load cases."""
+        pass