vectordb-bench 0.0.10__py3-none-any.whl → 0.0.12__py3-none-any.whl

vectordb_bench/backend/clients/test/cli.py

@@ -0,0 +1,25 @@
+from typing import Unpack
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+from ..test.config import TestConfig, TestIndexConfig
+
+
+class TestTypedDict(CommonTypedDict):
+    ...
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(TestTypedDict)
+def Test(**parameters: Unpack[TestTypedDict]):
+    run(
+        db=DB.NewClient,
+        db_config=TestConfig(db_label=parameters["db_label"]),
+        db_case_config=TestIndexConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/test/config.py

@@ -0,0 +1,18 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class TestConfig(DBConfig):
+    def to_dict(self) -> dict:
+        return {"db_label": self.db_label}
+
+
+class TestIndexConfig(BaseModel, DBCaseConfig):
+    metric_type: MetricType | None = None
+
+    def index_param(self) -> dict:
+        return {}
+
+    def search_param(self) -> dict:
+        return {}

vectordb_bench/backend/clients/test/test.py

@@ -0,0 +1,62 @@
+import logging
+from contextlib import contextmanager
+from typing import Any, Generator, Optional, Tuple
+
+from ..api import DBCaseConfig, VectorDB
+
+log = logging.getLogger(__name__)
+
+
+class Test(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: DBCaseConfig,
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.db_config = db_config
+        self.case_config = db_case_config
+
+        log.info("Starting Test DB")
+
+    @contextmanager
+    def init(self) -> Generator[None, None, None]:
+        """create and destroy connections to database.
+
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+        """
+
+        yield
+
+    def ready_to_load(self) -> bool:
+        return True
+
+    def optimize(self) -> None:
+        pass
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+        **kwargs: Any,
+    ) -> Tuple[int, Optional[Exception]]:
+        """Insert embeddings into the database.
+        Should call self.init() first.
+        """
+        raise RuntimeError("Not implemented")
+        return len(metadata), None
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+        timeout: int | None = None,
+        **kwargs: Any,
+    ) -> list[int]:
+        raise NotImplementedError
+        return [i for i in range(k)]

vectordb_bench/backend/clients/weaviate_cloud/cli.py

@@ -0,0 +1,41 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+
+
+class WeaviateTypedDict(CommonTypedDict):
+    api_key: Annotated[
+        str, click.option("--api-key", type=str, help="Weaviate api key", required=True)
+    ]
+    url: Annotated[
+        str,
+        click.option("--url", type=str, help="Weaviate url", required=True),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(WeaviateTypedDict)
+def Weaviate(**parameters: Unpack[WeaviateTypedDict]):
+    from .config import WeaviateConfig, WeaviateIndexConfig
+
+    run(
+        db=DB.WeaviateCloud,
+        db_config=WeaviateConfig(
+            db_label=parameters["db_label"],
+            api_key=SecretStr(parameters["api_key"]),
+            url=SecretStr(parameters["url"]),
+        ),
+        db_case_config=WeaviateIndexConfig(
+            ef=256, efConstruction=256, maxConnections=16
+        ),
+        **parameters,
+    )

vectordb_bench/backend/clients/zilliz_cloud/cli.py

@@ -0,0 +1,55 @@
+from typing import Annotated, Unpack
+
+import click
+import os
+from pydantic import SecretStr
+
+from vectordb_bench.cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from vectordb_bench.backend.clients import DB
+
+
+class ZillizTypedDict(CommonTypedDict):
+    uri: Annotated[
+        str, click.option("--uri", type=str, help="uri connection string", required=True)
+    ]
+    user_name: Annotated[
+        str, click.option("--user-name", type=str, help="Db username", required=True)
+    ]
+    password: Annotated[
+        str,
+        click.option("--password",
+                     type=str,
+                     help="Zilliz password",
+                     default=lambda: os.environ.get("ZILLIZ_PASSWORD", ""),
+                     show_default="$ZILLIZ_PASSWORD",
+                     ),
+    ]
+    level: Annotated[
+        str,
+        click.option("--level", type=str, help="Zilliz index level", required=False),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(ZillizTypedDict)
+def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
+    from .config import ZillizCloudConfig, AutoIndexConfig
+
+    run(
+        db=DB.ZillizCloud,
+        db_config=ZillizCloudConfig(
+            db_label=parameters["db_label"],
+            uri=SecretStr(parameters["uri"]),
+            user=parameters["user_name"],
+            password=SecretStr(parameters["password"]),
+        ),
+        db_case_config=AutoIndexConfig(
+            params={parameters["level"]},
+        ),
+        **parameters,
+    )

vectordb_bench/backend/dataset.py

@@ -33,6 +33,7 @@ class BaseDataset(BaseModel):
     use_shuffled: bool
     with_gt: bool = False
     _size_label: dict[int, SizeLabel] = PrivateAttr()
+    isCustom: bool = False
 
     @validator("size")
     def verify_size(cls, v):
@@ -52,7 +53,27 @@ class BaseDataset(BaseModel):
     def file_count(self) -> int:
         return self._size_label.get(self.size).file_count
 
+class CustomDataset(BaseDataset):
+    dir: str
+    file_num: int
+    isCustom: bool = True
+    
+    @validator("size")
+    def verify_size(cls, v):
+        return v
+    
+    @property
+    def label(self) -> str:
+        return "Custom"
 
+    @property
+    def dir_name(self) -> str:
+        return self.dir
+
+    @property
+    def file_count(self) -> int:
+        return self.file_num
+    
 class LAION(BaseDataset):
     name: str = "LAION"
     dim: int = 768
@@ -186,11 +207,12 @@ class DatasetManager(BaseModel):
             gt_file, test_file = utils.compose_gt_file(filters), "test.parquet"
             all_files.extend([gt_file, test_file])
 
-        source.reader().read(
-            dataset=self.data.dir_name.lower(),
-            files=all_files,
-            local_ds_root=self.data_dir,
-        )
+        if not self.data.isCustom:
+            source.reader().read(
+                dataset=self.data.dir_name.lower(),
+                files=all_files,
+                local_ds_root=self.data_dir,
+            )
 
         if gt_file is not None and test_file is not None:
             self.test_data = self._read_file(test_file)

vectordb_bench/backend/runner/mp_runner.py

@@ -4,6 +4,7 @@ import concurrent
 import multiprocessing as mp
 import logging
 from typing import Iterable
+import numpy as np
 from ..clients import api
 from ... import config
 
@@ -49,6 +50,7 @@ class MultiProcessingSearchRunner:
 
             start_time = time.perf_counter()
             count = 0
+            latencies = []
             while time.perf_counter() < start_time + self.duration:
                 s = time.perf_counter()
                 try:
@@ -61,7 +63,8 @@ class MultiProcessingSearchRunner:
                     log.warning(f"VectorDB search_embedding error: {e}")
                     traceback.print_exc(chain=True)
                     raise e from None
-
+                
+                latencies.append(time.perf_counter() - s)
                 count += 1
                 # loop through the test data
                 idx = idx + 1 if idx < num - 1 else 0
@@ -75,7 +78,7 @@ class MultiProcessingSearchRunner:
             f"actual_dur={total_dur}s, count={count}, qps in this process: {round(count / total_dur, 4):3}"
          )
 
-        return (count, total_dur)
+        return (count, total_dur, latencies)
 
     @staticmethod
     def get_mp_context():
@@ -85,6 +88,9 @@ class MultiProcessingSearchRunner:
 
     def _run_all_concurrencies_mem_efficient(self) -> float:
         max_qps = 0
+        conc_num_list = []
+        conc_qps_list = []
+        conc_latency_p99_list = []
         try:
             for conc in self.concurrencies:
                 with mp.Manager() as m:
@@ -103,9 +109,14 @@ class MultiProcessingSearchRunner:
 
                         start = time.perf_counter()
                         all_count = sum([r.result()[0] for r in future_iter])
+                        latencies = sum([r.result()[2] for r in future_iter], start=[])
+                        latency_p99 = np.percentile(latencies, 0.99)
                         cost = time.perf_counter() - start
 
                         qps = round(all_count / cost, 4)
+                        conc_num_list.append(conc)
+                        conc_qps_list.append(qps)
+                        conc_latency_p99_list.append(latency_p99)
                         log.info(f"End search in concurrency {conc}: dur={cost}s, total_count={all_count}, qps={qps}")
 
                 if qps > max_qps:
@@ -122,7 +133,7 @@ class MultiProcessingSearchRunner:
         finally:
             self.stop()
 
-        return max_qps
+        return max_qps, conc_num_list, conc_qps_list, conc_latency_p99_list
 
     def run(self) -> float:
         """

vectordb_bench/backend/runner/serial_runner.py

@@ -10,7 +10,7 @@ import numpy as np
 import pandas as pd
 
 from ..clients import api
-from ...metric import calc_recall
+from ...metric import calc_ndcg, calc_recall, get_ideal_dcg
 from ...models import LoadTimeoutError, PerformanceTimeoutError
 from .. import utils
 from ... import config
@@ -171,11 +171,12 @@ class SerialSearchRunner:
         log.info(f"{mp.current_process().name:14} start search the entire test_data to get recall and latency")
         with self.db.init():
             test_data, ground_truth = args
+            ideal_dcg = get_ideal_dcg(self.k)
 
             log.debug(f"test dataset size: {len(test_data)}")
             log.debug(f"ground truth size: {ground_truth.columns}, shape: {ground_truth.shape}")
 
-            latencies, recalls = [], []
+            latencies, recalls, ndcgs = [], [], []
             for idx, emb in enumerate(test_data):
                 s = time.perf_counter()
                 try:
@@ -194,6 +195,7 @@ class SerialSearchRunner:
 
                 gt = ground_truth['neighbors_id'][idx]
                 recalls.append(calc_recall(self.k, gt[:self.k], results))
+                ndcgs.append(calc_ndcg(gt[:self.k], results, ideal_dcg))
 
 
                 if len(latencies) % 100 == 0:
@@ -201,6 +203,7 @@ class SerialSearchRunner:
 
         avg_latency = round(np.mean(latencies), 4)
         avg_recall = round(np.mean(recalls), 4)
+        avg_ndcg = round(np.mean(ndcgs), 4)
         cost = round(np.sum(latencies), 4)
         p99 = round(np.percentile(latencies, 99), 4)
         log.info(
@@ -208,10 +211,11 @@ class SerialSearchRunner:
             f"cost={cost}s, "
             f"queries={len(latencies)}, "
             f"avg_recall={avg_recall}, "
+            f"avg_ndcg={avg_ndcg},"
             f"avg_latency={avg_latency}, "
             f"p99={p99}"
          )
-        return (avg_recall, p99)
+        return (avg_recall, avg_ndcg, p99)
 
 
     def _run_in_subprocess(self) -> tuple[float, float]:

vectordb_bench/backend/task_runner.py

@@ -8,7 +8,7 @@ from enum import Enum, auto
 from . import utils
 from .cases import Case, CaseLabel
 from ..base import BaseModel
-from ..models import TaskConfig, PerformanceTimeoutError
+from ..models import TaskConfig, PerformanceTimeoutError, TaskStage
 
 from .clients import (
     api,
@@ -29,7 +29,7 @@ class RunningStatus(Enum):
 
 
 class CaseRunner(BaseModel):
-    """ DataSet, filter_rate, db_class with db config
+    """DataSet, filter_rate, db_class with db config
 
     Fields:
         run_id(str): run_id of this case runner,
@@ -49,8 +49,9 @@ class CaseRunner(BaseModel):
 
     db: api.VectorDB | None = None
     test_emb: list[list[float]] | None = None
-    search_runner: MultiProcessingSearchRunner | None = None
     serial_search_runner: SerialSearchRunner | None = None
+    search_runner: MultiProcessingSearchRunner | None = None
+    final_search_runner: MultiProcessingSearchRunner | None = None
 
     def __eq__(self, obj):
         if isinstance(obj, CaseRunner):
@@ -58,7 +59,7 @@ class CaseRunner(BaseModel):
                 self.config.db == obj.config.db and \
                 self.config.db_case_config == obj.config.db_case_config and \
                 self.ca.dataset == obj.ca.dataset
-            return False
+        return False
 
     def display(self) -> dict:
         c_dict = self.ca.dict(include={'label':True, 'filters': True,'dataset':{'data': {'name': True, 'size': True, 'dim': True, 'metric_type': True, 'label': True}} })
@@ -79,20 +80,25 @@ class CaseRunner(BaseModel):
             db_config=self.config.db_config.to_dict(),
             db_case_config=self.config.db_case_config,
             drop_old=drop_old,
-        )
+        )  # type:ignore
+
 
     def _pre_run(self, drop_old: bool = True):
         try:
             self.init_db(drop_old)
             self.ca.dataset.prepare(self.dataset_source, filters=self.ca.filter_rate)
         except ModuleNotFoundError as e:
-            log.warning(f"pre run case error: please install client for db: {self.config.db}, error={e}")
+            log.warning(
+                f"pre run case error: please install client for db: {self.config.db}, error={e}"
+            )
             raise e from None
         except Exception as e:
             log.warning(f"pre run case error: {e}")
             raise e from None
 
     def run(self, drop_old: bool = True) -> Metric:
+        log.info("Starting run")
+
         self._pre_run(drop_old)
 
         if self.ca.label == CaseLabel.Load:
@@ -105,31 +111,35 @@ class CaseRunner(BaseModel):
             raise ValueError(msg)
 
     def _run_capacity_case(self) -> Metric:
-        """ run capacity cases
+        """run capacity cases
 
         Returns:
             Metric: the max load count
         """
+        assert self.db is not None
         log.info("Start capacity case")
         try:
-            runner = SerialInsertRunner(self.db, self.ca.dataset, self.normalize, self.ca.load_timeout)
+            runner = SerialInsertRunner(
+                self.db, self.ca.dataset, self.normalize, self.ca.load_timeout
+            )
             count = runner.run_endlessness()
         except Exception as e:
             log.warning(f"Failed to run capacity case, reason = {e}")
             raise e from None
         else:
-            log.info(f"Capacity case loading dataset reaches VectorDB's limit: max capacity = {count}")
+            log.info(
+                f"Capacity case loading dataset reaches VectorDB's limit: max capacity = {count}"
+            )
             return Metric(max_load_count=count)
 
     def _run_perf_case(self, drop_old: bool = True) -> Metric:
-        """ run performance cases
+        """run performance cases
 
         Returns:
             Metric: load_duration, recall, serial_latency_p99, and, qps
         """
-        try:
-            m = Metric()
-            if drop_old:
+        '''
+                    if drop_old:
                 _, load_dur = self._load_train_data()
                 build_dur = self._optimize()
                 m.load_duration = round(load_dur+build_dur, 4)
@@ -140,8 +150,43 @@ class CaseRunner(BaseModel):
                 )
 
             self._init_search_runner()
-            m.qps = self._conc_search()
+    
+            m.qps, m.conc_num_list, m.conc_qps_list, m.conc_latency_p99_list = self._conc_search()
             m.recall, m.serial_latency_p99 = self._serial_search()
+        '''
+
+        log.info("Start performance case")
+        try:
+            m = Metric()
+            if drop_old:
+                if TaskStage.LOAD in self.config.stages:
+                    # self._load_train_data()
+                    _, load_dur = self._load_train_data()
+                    build_dur = self._optimize()
+                    m.load_duration = round(load_dur + build_dur, 4)
+                    log.info(
+                        f"Finish loading the entire dataset into VectorDB,"
+                        f" insert_duration={load_dur}, optimize_duration={build_dur}"
+                        f" load_duration(insert + optimize) = {m.load_duration}"
+                    )
+                else:
+                    log.info("Data loading skipped")
+            if (
+                TaskStage.SEARCH_SERIAL in self.config.stages
+                or TaskStage.SEARCH_CONCURRENT in self.config.stages
+            ):
+                self._init_search_runner()
+                if TaskStage.SEARCH_SERIAL in self.config.stages:
+                    search_results = self._serial_search()
+                    '''
+                    m.recall = search_results.recall
+                    m.serial_latencies = search_results.serial_latencies
+                    '''
+                    m.recall, m.ndcg, m.serial_latency_p99 = search_results
+                if TaskStage.SEARCH_CONCURRENT in self.config.stages:
+                    search_results = self._conc_search()
+                    m.qps, m.conc_num_list, m.conc_qps_list, m.conc_latency_p99_list = search_results
+            
         except Exception as e:
             log.warning(f"Failed to run performance case, reason = {e}")
             traceback.print_exc()
@@ -217,18 +262,23 @@ class CaseRunner(BaseModel):
 
         gt_df = self.ca.dataset.gt_data
 
-        self.serial_search_runner = SerialSearchRunner(
-            db=self.db,
-            test_data=self.test_emb,
-            ground_truth=gt_df,
-            filters=self.ca.filters,
-        )
-
-        self.search_runner =  MultiProcessingSearchRunner(
-            db=self.db,
-            test_data=self.test_emb,
-            filters=self.ca.filters,
-        )
+        if TaskStage.SEARCH_SERIAL in self.config.stages:
+            self.serial_search_runner = SerialSearchRunner(
+                db=self.db,
+                test_data=self.test_emb,
+                ground_truth=gt_df,
+                filters=self.ca.filters,
+                k=self.config.case_config.k,
+            )
+        if TaskStage.SEARCH_CONCURRENT in self.config.stages:
+            self.search_runner = MultiProcessingSearchRunner(
+                db=self.db,
+                test_data=self.test_emb,
+                filters=self.ca.filters,
+                concurrencies=self.config.case_config.concurrency_search_config.num_concurrency,
+                duration=self.config.case_config.concurrency_search_config.concurrency_duration,
+                k=self.config.case_config.k,
+            )
 
     def stop(self):
         if self.search_runner: