vectordb-bench 0.0.28__py3-none-any.whl → 0.0.30__py3-none-any.whl

vectordb_bench/backend/clients/qdrant_local/qdrant_local.py

@@ -0,0 +1,232 @@
+"""Wrapper around the Qdrant over VectorDB"""
+
+import logging
+import time
+from collections.abc import Iterable
+from contextlib import contextmanager
+
+from qdrant_client import QdrantClient
+from qdrant_client.http.models import (
+    Batch,
+    CollectionStatus,
+    FieldCondition,
+    Filter,
+    HnswConfigDiff,
+    OptimizersConfigDiff,
+    PayloadSchemaType,
+    Range,
+    SearchParams,
+    VectorParams,
+)
+
+from ..api import VectorDB
+from .config import QdrantLocalIndexConfig
+
+log = logging.getLogger(__name__)
+
+SECONDS_WAITING_FOR_INDEXING_API_CALL = 5
+QDRANT_BATCH_SIZE = 100
+
+
+def qdrant_collection_exists(client: QdrantClient, collection_name: str) -> bool:
+    collection_exists = True
+
+    try:
+        client.get_collection(collection_name)
+    except Exception:
+        collection_exists = False
+
+    return collection_exists
+
+
+class QdrantLocal(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: dict,
+        db_case_config: QdrantLocalIndexConfig,
+        collection_name: str = "QdrantLocalCollection",
+        drop_old: bool = False,
+        name: str = "QdrantLocal",
+        **kwargs,
+    ):
+        """Initialize wrapper around the qdrant."""
+        self.name = name
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.search_parameter = self.case_config.search_param()
+        self.collection_name = collection_name
+        self.client = None
+
+        self._primary_field = "pk"
+        self._vector_field = "vector"
+
+        client = QdrantClient(**self.db_config)
+
+        # Lets just print the parameters here for double check
+        log.info(f"Case config: {self.case_config.index_param()}")
+        log.info(f"Search parameter: {self.search_parameter}")
+
+        if drop_old and qdrant_collection_exists(client, self.collection_name):
+            log.info(f"{self.name} client drop_old collection: {self.collection_name}")
+            client.delete_collection(self.collection_name)
+
+        if not qdrant_collection_exists(client, self.collection_name):
+            log.info(f"{self.name} create collection: {self.collection_name}")
+            self._create_collection(dim, client)
+
+        client = None
+
+    @contextmanager
+    def init(self):
+        """
+        Examples:
+            >>> with self.init():
+            >>>     self.insert_embeddings()
+            >>>     self.search_embedding()
+        """
+        # create connection
+        self.client = QdrantClient(**self.db_config)
+        yield
+        self.client = None
+        del self.client
+
+    def _create_collection(self, dim: int, qdrant_client: QdrantClient):
+        log.info(f"Create collection: {self.collection_name}")
+        log.info(
+            f"Index parameters: m={self.case_config.index_param()['m']}, "
+            f"ef_construct={self.case_config.index_param()['ef_construct']}, "
+            f"on_disk={self.case_config.index_param()['on_disk']}"
+        )
+
+        # If the on_disk is true, we enable both on disk index and vectors.
+        try:
+            qdrant_client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=VectorParams(
+                    size=dim,
+                    distance=self.case_config.index_param()["distance"],
+                    on_disk=self.case_config.index_param()["on_disk"],
+                ),
+                hnsw_config=HnswConfigDiff(
+                    m=self.case_config.index_param()["m"],
+                    ef_construct=self.case_config.index_param()["ef_construct"],
+                    on_disk=self.case_config.index_param()["on_disk"],
+                ),
+            )
+
+            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 optimize(self, data_size: int | None = None):
+        assert self.client, "Please call self.init() before"
+        # wait for vectors to be fully indexed
+        try:
+            while True:
+                info = self.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"Finishing building index for collection: {self.collection_name}")
+                    msg = (
+                        f"Stored vectors: {info.vectors_count}, Indexed vectors: {info.indexed_vectors_count}, "
+                        f"Collection status: {info.indexed_vectors_count}"
+                    )
+                    log.info(msg)
+                    return
+
+        except Exception as e:
+            log.warning(f"QdrantCloud ready to search error: {e}")
+            raise e from None
+
+    def insert_embeddings(
+        self,
+        embeddings: Iterable[list[float]],
+        metadata: list[int],
+        **kwargs,
+    ) -> tuple[int, Exception]:
+        """Insert embeddings into the database.
+
+        Args:
+            embeddings(list[list[float]]): list of embeddings
+            metadata(list[int]): list of metadata
+            kwargs: other arguments
+
+        Returns:
+            tuple[int, Exception]: number of embeddings inserted and exception if any
+        """
+        assert self.client is not None
+        assert len(embeddings) == len(metadata)
+        insert_count = 0
+
+        # disable indexing for quick insertion
+        self.client.update_collection(
+            collection_name=self.collection_name,
+            optimizer_config=OptimizersConfigDiff(indexing_threshold=0),
+        )
+        try:
+            for offset in range(0, len(embeddings), QDRANT_BATCH_SIZE):
+                vectors = embeddings[offset : offset + QDRANT_BATCH_SIZE]
+                ids = metadata[offset : offset + QDRANT_BATCH_SIZE]
+                payloads = [{self._primary_field: v} for v in ids]
+                _ = self.client.upsert(
+                    collection_name=self.collection_name,
+                    wait=True,
+                    points=Batch(ids=ids, payloads=payloads, vectors=vectors),
+                )
+                insert_count += QDRANT_BATCH_SIZE
+            # enable indexing after insertion
+            self.client.update_collection(
+                collection_name=self.collection_name,
+                optimizer_config=OptimizersConfigDiff(indexing_threshold=100),
+            )
+
+        except Exception as e:
+            log.info(f"Failed to insert data, {e}")
+            return insert_count, e
+        else:
+            return insert_count, 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 with score.
+        Should call self.init() first.
+        """
+        assert self.client is not None
+
+        f = None
+        if filters:
+            f = Filter(
+                must=[
+                    FieldCondition(
+                        key=self._primary_field,
+                        range=Range(
+                            gt=filters.get("id"),
+                        ),
+                    ),
+                ],
+            )
+        res = self.client.query_points(
+            collection_name=self.collection_name,
+            query=query,
+            limit=k,
+            query_filter=f,
+            search_params=SearchParams(**self.search_parameter),
+        ).points
+
+        return [result.id for result in res]

vectordb_bench/backend/clients/weaviate_cloud/cli.py

@@ -15,12 +15,33 @@ from .. import DB
 class WeaviateTypedDict(CommonTypedDict):
     api_key: Annotated[
         str,
-        click.option("--api-key", type=str, help="Weaviate api key", required=True),
+        click.option("--api-key", type=str, help="Weaviate api key", required=False, default=""),
     ]
     url: Annotated[
         str,
         click.option("--url", type=str, help="Weaviate url", required=True),
     ]
+    no_auth: Annotated[
+        bool,
+        click.option(
+            "--no-auth",
+            is_flag=True,
+            help="Do not use api-key, set it to true if you are using a local setup. Default is False.",
+            default=False,
+        ),
+    ]
+    m: Annotated[
+        int,
+        click.option("--m", type=int, default=16, help="HNSW index parameter m."),
+    ]
+    ef_construct: Annotated[
+        int,
+        click.option("--ef-construction", type=int, default=256, help="HNSW index parameter ef_construction"),
+    ]
+    ef: Annotated[
+        int,
+        click.option("--ef", type=int, default=256, help="HNSW index parameter ef for search"),
+    ]
 
 
 @cli.command()
@@ -32,9 +53,14 @@ def Weaviate(**parameters: Unpack[WeaviateTypedDict]):
         db=DB.WeaviateCloud,
         db_config=WeaviateConfig(
             db_label=parameters["db_label"],
-            api_key=SecretStr(parameters["api_key"]),
+            api_key=SecretStr(parameters["api_key"]) if parameters["api_key"] != "" else SecretStr("-"),
             url=SecretStr(parameters["url"]),
+            no_auth=parameters["no_auth"],
+        ),
+        db_case_config=WeaviateIndexConfig(
+            efConstruction=parameters["ef_construction"],
+            maxConnections=parameters["m"],
+            ef=parameters["ef"],
         ),
-        db_case_config=WeaviateIndexConfig(ef=256, efConstruction=256, maxConnections=16),
         **parameters,
     )

vectordb_bench/backend/clients/weaviate_cloud/config.py

@@ -6,11 +6,13 @@ from ..api import DBCaseConfig, DBConfig, MetricType
 class WeaviateConfig(DBConfig):
     url: SecretStr
     api_key: SecretStr
+    no_auth: bool | None = False
 
     def to_dict(self) -> dict:
         return {
             "url": self.url.get_secret_value(),
             "auth_client_secret": self.api_key.get_secret_value(),
+            "no_auth": self.no_auth,
         }
 
 

vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py

@@ -38,6 +38,11 @@ class WeaviateCloud(VectorDB):
         self._vector_field = "vector"
         self._index_name = "vector_idx"
 
+        # If local setup is used, we
+        if db_config["no_auth"]:
+            del db_config["auth_client_secret"]
+        del db_config["no_auth"]
+
         from weaviate import Client
 
         client = Client(**db_config)

vectordb_bench/backend/runner/mp_runner.py

@@ -5,10 +5,12 @@ import random
 import time
 import traceback
 from collections.abc import Iterable
+from multiprocessing.queues import Queue
 
 import numpy as np
 
 from ... import config
+from ...models import ConcurrencySlotTimeoutError
 from ..clients import api
 
 NUM_PER_BATCH = config.NUM_PER_BATCH
@@ -28,16 +30,18 @@ class MultiProcessingSearchRunner:
         self,
         db: api.VectorDB,
         test_data: list[list[float]],
-        k: int = 100,
+        k: int = config.K_DEFAULT,
         filters: dict | None = None,
         concurrencies: Iterable[int] = config.NUM_CONCURRENCY,
-        duration: int = 30,
+        duration: int = config.CONCURRENCY_DURATION,
+        concurrency_timeout: int = config.CONCURRENCY_TIMEOUT,
     ):
         self.db = db
         self.k = k
         self.filters = filters
         self.concurrencies = concurrencies
         self.duration = duration
+        self.concurrency_timeout = concurrency_timeout
 
         self.test_data = test_data
         log.debug(f"test dataset columns: {len(test_data)}")
@@ -114,9 +118,7 @@ class MultiProcessingSearchRunner:
                         log.info(f"Start search {self.duration}s in concurrency {conc}, filters: {self.filters}")
                         future_iter = [executor.submit(self.search, self.test_data, q, cond) for i in range(conc)]
                         # Sync all processes
-                        while q.qsize() < conc:
-                            sleep_t = conc if conc < 10 else 10
-                            time.sleep(sleep_t)
+                        self._wait_for_queue_fill(q, size=conc)
 
                         with cond:
                             cond.notify_all()
@@ -160,6 +162,15 @@ class MultiProcessingSearchRunner:
             conc_latency_avg_list,
         )
 
+    def _wait_for_queue_fill(self, q: Queue, size: int):
+        wait_t = 0
+        while q.qsize() < size:
+            sleep_t = size if size < 10 else 10
+            wait_t += sleep_t
+            if wait_t > self.concurrency_timeout > 0:
+                raise ConcurrencySlotTimeoutError
+            time.sleep(sleep_t)
+
     def run(self) -> float:
         """
         Returns:

vectordb_bench/backend/task_runner.py

@@ -275,6 +275,7 @@ class CaseRunner(BaseModel):
                 filters=self.ca.filters,
                 concurrencies=self.config.case_config.concurrency_search_config.num_concurrency,
                 duration=self.config.case_config.concurrency_search_config.concurrency_duration,
+                concurrency_timeout=self.config.case_config.concurrency_search_config.concurrency_timeout,
                 k=self.config.case_config.k,
             )
 

vectordb_bench/cli/batch_cli.py

@@ -0,0 +1,121 @@
+import logging
+import time
+from collections.abc import MutableMapping
+from concurrent.futures import wait
+from pathlib import Path
+from typing import Annotated, Any, TypedDict
+
+import click
+from click.testing import CliRunner
+from yaml import Loader, load
+
+from .. import config
+from ..cli.cli import (
+    cli,
+    click_parameter_decorators_from_typed_dict,
+)
+
+log = logging.getLogger(__name__)
+
+
+def click_get_defaults_from_file(ctx, param, value):  # noqa: ANN001, ARG001
+    if not value:
+        raise click.MissingParameter
+    path = Path(value)
+    input_file = path if path.exists() else Path(config.CONFIG_LOCAL_DIR, path)
+    try:
+        with input_file.open() as f:
+            _config: dict[str, list[dict[str, Any]]] = load(f.read(), Loader=Loader)  # noqa: S506
+            ctx.default_map = _config
+    except Exception as e:
+        msg = f"Failed to load batch config file: {e}"
+        raise click.BadParameter(msg) from e
+    return value
+
+
+class BatchCliTypedDict(TypedDict):
+    batch_config_file: Annotated[
+        bool,
+        click.option(
+            "--batch-config-file",
+            type=click.Path(),
+            callback=click_get_defaults_from_file,
+            is_eager=True,
+            expose_value=False,
+            help="Read batch configuration from yaml file",
+        ),
+    ]
+
+
+def build_sub_cmd_args(batch_config: MutableMapping[str, Any] | None):
+    bool_options = {
+        "drop_old": True,
+        "load": True,
+        "search_serial": True,
+        "search_concurrent": True,
+        "dry_run": False,
+        "custom_dataset_use_shuffled": True,
+        "custom_dataset_with_gt": True,
+    }
+
+    def format_option(key: str, value: Any):
+        opt_name = key.replace("_", "-")
+
+        if key in bool_options:
+            return format_bool_option(opt_name, value, skip=False)
+
+        if key.startswith("skip_"):
+            raw_key = key[5:]
+            raw_opt = raw_key.replace("_", "-")
+            return format_bool_option(raw_opt, value, skip=True, raw_key=raw_key)
+
+        return [f"--{opt_name}", str(value)]
+
+    def format_bool_option(opt_name: str, value: Any, skip: bool = False, raw_key: str | None = None):
+        if isinstance(value, bool):
+            if skip:
+                if bool_options.get(raw_key, False):
+                    return [f"--skip-{opt_name}"] if value else [f"--{opt_name}"]
+                return [f"--{opt_name}", str(value)]
+            if value:
+                return [f"--{opt_name}"]
+            if bool_options.get(opt_name.replace("-", "_"), False):
+                return [f"--skip-{opt_name}"]
+            return []
+        return [f"--{opt_name}", str(value)]
+
+    args_arr = []
+    for sub_cmd_key, sub_cmd_config_list in batch_config.items():
+        for sub_cmd_args in sub_cmd_config_list:
+            args = [sub_cmd_key]
+            for k, v in sub_cmd_args.items():
+                args.extend(format_option(k, v))
+            args_arr.append(args)
+
+    return args_arr
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(BatchCliTypedDict)
+def BatchCli():
+    ctx = click.get_current_context()
+    batch_config = ctx.default_map
+
+    runner = CliRunner()
+
+    args_arr = build_sub_cmd_args(batch_config)
+
+    for args in args_arr:
+        log.info(f"got batch config: {' '.join(args)}")
+
+    for args in args_arr:
+        result = runner.invoke(cli, args)
+        time.sleep(5)
+
+        from ..interface import global_result_future
+
+        if global_result_future:
+            wait([global_result_future])
+
+        if result.exception:
+            log.exception(f"failed to run sub command: {args[0]}", exc_info=result.exception)

vectordb_bench/cli/cli.py

@@ -17,10 +17,9 @@ from typing import (
 import click
 from yaml import load
 
-from vectordb_bench.backend.clients.api import MetricType
-
 from .. import config
 from ..backend.clients import DB
+from ..backend.clients.api import MetricType
 from ..interface import benchmark_runner, global_result_future
 from ..models import (
     CaseConfig,
@@ -303,6 +302,17 @@ class CommonTypedDict(TypedDict):
             callback=lambda *args: list(map(int, click_arg_split(*args))),
         ),
     ]
+    concurrency_timeout: Annotated[
+        int,
+        click.option(
+            "--concurrency-timeout",
+            type=int,
+            default=config.CONCURRENCY_TIMEOUT,
+            show_default=True,
+            help="Timeout (in seconds) to wait for a concurrency slot before failing. "
+            "Set to a negative value to wait indefinitely.",
+        ),
+    ]
     custom_case_name: Annotated[
         str,
         click.option(
@@ -490,6 +500,7 @@ def run(
             concurrency_search_config=ConcurrencySearchConfig(
                 concurrency_duration=parameters["concurrency_duration"],
                 num_concurrency=[int(s) for s in parameters["num_concurrency"]],
+                concurrency_timeout=parameters["concurrency_timeout"],
             ),
             custom_case=get_custom_case_config(parameters),
         ),

vectordb_bench/cli/vectordbbench.py

@@ -9,12 +9,15 @@ from ..backend.clients.pgdiskann.cli import PgDiskAnn
 from ..backend.clients.pgvecto_rs.cli import PgVectoRSHNSW, PgVectoRSIVFFlat
 from ..backend.clients.pgvector.cli import PgVectorHNSW
 from ..backend.clients.pgvectorscale.cli import PgVectorScaleDiskAnn
+from ..backend.clients.qdrant_cloud.cli import QdrantCloud
+from ..backend.clients.qdrant_local.cli import QdrantLocal
 from ..backend.clients.redis.cli import Redis
 from ..backend.clients.test.cli import Test
 from ..backend.clients.tidb.cli import TiDB
 from ..backend.clients.vespa.cli import Vespa
 from ..backend.clients.weaviate_cloud.cli import Weaviate
 from ..backend.clients.zilliz_cloud.cli import ZillizAutoIndex
+from .batch_cli import BatchCli
 from .cli import cli
 
 cli.add_command(PgVectorHNSW)
@@ -35,6 +38,9 @@ cli.add_command(TiDB)
 cli.add_command(Clickhouse)
 cli.add_command(Vespa)
 cli.add_command(LanceDB)
+cli.add_command(QdrantCloud)
+cli.add_command(QdrantLocal)
+cli.add_command(BatchCli)
 
 
 if __name__ == "__main__":

vectordb_bench/config-files/batch_sample_config.yml

@@ -0,0 +1,17 @@
+pgvectorhnsw:
+  - db_label: pgConfigTest
+    user_name: vectordbbench
+    db_name:  vectordbbench
+    host: localhost
+    m: 16
+    ef_construction: 128
+    ef_search: 128
+milvushnsw:
+  - skip_search_serial: True
+    case_type: Performance1536D50K
+    uri: http://localhost:19530
+    m: 16
+    ef_construction: 128
+    ef_search: 128
+    drop_old: False
+    load: False

vectordb_bench/frontend/components/run_test/dbConfigSetting.py

@@ -36,21 +36,27 @@ def dbConfigSettingItem(st, activeDb: DB):
     columns = st.columns(DB_CONFIG_SETTING_COLUMNS)
 
     dbConfigClass = activeDb.config_cls
-    properties = dbConfigClass.schema().get("properties")
+    schema = dbConfigClass.schema()
+    property_items = schema.get("properties").items()
+    required_fields = set(schema.get("required", []))
     dbConfig = {}
     idx = 0
 
     # db config (unique)
-    for key, property in properties.items():
+    for key, property in property_items:
         if key not in dbConfigClass.common_short_configs() and key not in dbConfigClass.common_long_configs():
             column = columns[idx % DB_CONFIG_SETTING_COLUMNS]
             idx += 1
-            dbConfig[key] = column.text_input(
+            input_value = column.text_input(
                 key,
-                key="%s-%s" % (activeDb.name, key),
+                key=f"{activeDb.name}-{key}",
                 value=property.get("default", ""),
                 type="password" if inputIsPassword(key) else "default",
+                placeholder="optional" if key not in required_fields else None,
             )
+            if key in required_fields or input_value:
+                dbConfig[key] = input_value
+
     # db config (common short labels)
     for key in dbConfigClass.common_short_configs():
         column = columns[idx % DB_CONFIG_SETTING_COLUMNS]