From 295062afa6cb5b8d1d4ae9ccfbdd56b1a2114940 Mon Sep 17 00:00:00 2001 From: Zebin Ren Date: Wed, 4 Jun 2025 00:47:54 +0200 Subject: [PATCH 1/4] Add Qdrant local --- vectordb_bench/backend/clients/__init__.py | 15 ++ .../backend/clients/qdrant_local/cli.py | 59 +++++ .../backend/clients/qdrant_local/config.py | 39 +++ .../clients/qdrant_local/qdrant_local.py | 223 ++++++++++++++++++ vectordb_bench/cli/vectordbbench.py | 2 + 5 files changed, 338 insertions(+) create mode 100644 vectordb_bench/backend/clients/qdrant_local/cli.py create mode 100644 vectordb_bench/backend/clients/qdrant_local/config.py create mode 100644 vectordb_bench/backend/clients/qdrant_local/qdrant_local.py diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index f05913a06..21903594e 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -27,6 +27,7 @@ class DB(Enum): Pinecone = "Pinecone" ElasticCloud = "ElasticCloud" QdrantCloud = "QdrantCloud" + QdrantLocal = "QdrantLocal" WeaviateCloud = "WeaviateCloud" PgVector = "PgVector" PgVectoRS = "PgVectoRS" @@ -46,6 +47,7 @@ class DB(Enum): Clickhouse = "Clickhouse" Vespa = "Vespa" LanceDB = "LanceDB" + @property def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 @@ -74,6 +76,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 from .qdrant_cloud.qdrant_cloud import QdrantCloud return QdrantCloud + + if self == DB.QdrantLocal: + from .qdrant_local.qdrant_local import QdrantLocal + + return QdrantLocal if self == DB.WeaviateCloud: from .weaviate_cloud.weaviate_cloud import WeaviateCloud @@ -200,6 +207,9 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 from .qdrant_cloud.config import QdrantConfig return QdrantConfig + + if self == DB.QdrantLocal: + from .qdrant_local.config import QdrantLocalConfig if self == DB.WeaviateCloud: from .weaviate_cloud.config import WeaviateConfig @@ -322,6 +332,11 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912 from .qdrant_cloud.config import QdrantIndexConfig return QdrantIndexConfig + + if self == DB.QdrantLocal: + from .qdrant_local.config import QdrantLocalIndexConfig + + return QdrantLocalIndexConfig if self == DB.WeaviateCloud: from .weaviate_cloud.config import WeaviateIndexConfig diff --git a/vectordb_bench/backend/clients/qdrant_local/cli.py b/vectordb_bench/backend/clients/qdrant_local/cli.py new file mode 100644 index 000000000..4a2083e12 --- /dev/null +++ b/vectordb_bench/backend/clients/qdrant_local/cli.py @@ -0,0 +1,59 @@ +from typing import Annotated, TypedDict, Unpack + +import click +from pydantic import SecretStr + +from vectordb_bench.backend.clients import DB +from vectordb_bench.cli.cli import ( + CommonTypedDict, + cli, + click_parameter_decorators_from_typed_dict, + run, +) + + +DBTYPE = DB.QdrantLocal + + +class QdrantLocalTypedDict(CommonTypedDict): + url: Annotated[ + str, + click.option("--url", type=str, help="Qdrant url", required=True), + ] + on_disk: Annotated[ + bool, + click.option( + "--on-disk", type=bool, default=False, help="Store the vectors and the HNSW index on disk" + ), + ] + m: Annotated[ + int, + click.option( + "--m", type=int, default=16, help="HNSW index parameter m, set 0 to disable the index" + ), + ] + ef_construct: Annotated[ + int, + click.option( + "--ef-construct", type=int, default=200, help="HNSW index parameter ef_construct" + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(QdrantLocalTypedDict) +def QdrantLocal(**parameters: Unpack[QdrantLocalTypedDict]): + from .config import QdrantLocalConfig, QdrantLocalIndexConfig + + run( + db=DBTYPE, + db_config=QdrantLocalConfig( + url=SecretStr(parameters["url"]) + ), + db_case_config=QdrantLocalIndexConfig( + on_disk=parameters["on_disk"], + m=parameters["m"], + ef_construct=parameters["ef_construct"], + ), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/qdrant_local/config.py b/vectordb_bench/backend/clients/qdrant_local/config.py new file mode 100644 index 000000000..b4eb5da47 --- /dev/null +++ b/vectordb_bench/backend/clients/qdrant_local/config.py @@ -0,0 +1,39 @@ +from pydantic import BaseModel, SecretStr + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType + +class QdrantLocalConfig(DBConfig): + url: SecretStr + + def to_dict(self) -> dict: + return { + "url": self.url.get_secret_value(), + } + + +class QdrantLocalIndexConfig(BaseModel, DBCaseConfig): + metric_type: MetricType | None = None + m: int + ef_construct: int + on_disk: bool | None = False + + def parse_metric(self) -> str: + if self.metric_type == MetricType.L2: + return "Euclid" + + if self.metric_type == MetricType.IP: + return "Dot" + + return "Cosine" + + def index_param(self) -> dict: + + return { + "distance": self.parse_metric(), + "m": self.m, + "ef_construct": self.ef_construct, + "on_disk": self.on_disk, + } + + def search_param(self) -> dict: + return {} \ No newline at end of file diff --git a/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py b/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py new file mode 100644 index 000000000..636d7aa22 --- /dev/null +++ b/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py @@ -0,0 +1,223 @@ +"""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, + 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, collection_name: str) -> bool: + collection_exists = True + + try: + client.get_collection(collection_name) + except Exception as e: + collection_exists = False + + return collection_exists + +class QdrantLocal(VectorDB): + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: dict, + 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.collection_name = collection_name + self.client = None + + self._primary_field = "pk" + self._vector_field = "vector" + + client = QdrantClient(**self.db_config) + + 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']}, ef_construct={self.case_config.index_param()['ef_construct']}, 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, + ).points + ) + + return [result.id for result in res] + diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index d4153bc1e..05ee7fd8d 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -10,6 +10,7 @@ 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 @@ -37,6 +38,7 @@ cli.add_command(Vespa) cli.add_command(LanceDB) cli.add_command(QdrantCloud) +cli.add_command(QdrantLocal) if __name__ == "__main__": From 5cef7d9366171ad94fdca8a2ec549b340787d5d6 Mon Sep 17 00:00:00 2001 From: Zebin Ren Date: Fri, 6 Jun 2025 08:58:02 +0200 Subject: [PATCH 2/4] Add support for local setup without authentication api-key for weaviate --- vectordb_bench/backend/clients/weaviate_cloud/cli.py | 9 +++++++-- vectordb_bench/backend/clients/weaviate_cloud/config.py | 2 ++ .../backend/clients/weaviate_cloud/weaviate_cloud.py | 5 +++++ 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/vectordb_bench/backend/clients/weaviate_cloud/cli.py b/vectordb_bench/backend/clients/weaviate_cloud/cli.py index 181898c74..6880411a7 100644 --- a/vectordb_bench/backend/clients/weaviate_cloud/cli.py +++ b/vectordb_bench/backend/clients/weaviate_cloud/cli.py @@ -15,12 +15,16 @@ 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), + ] @cli.command() @@ -32,8 +36,9 @@ 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(ef=256, efConstruction=256, maxConnections=16), **parameters, diff --git a/vectordb_bench/backend/clients/weaviate_cloud/config.py b/vectordb_bench/backend/clients/weaviate_cloud/config.py index 4c58167d4..f29a307a3 100644 --- a/vectordb_bench/backend/clients/weaviate_cloud/config.py +++ b/vectordb_bench/backend/clients/weaviate_cloud/config.py @@ -6,11 +6,13 @@ 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, } diff --git a/vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py b/vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py index c31104d8b..18a17a661 100644 --- a/vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py +++ b/vectordb_bench/backend/clients/weaviate_cloud/weaviate_cloud.py @@ -37,6 +37,11 @@ def __init__( self._scalar_field = "key" 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 From 29205e7112283e916c4a526e84d250f80b6d2799 Mon Sep 17 00:00:00 2001 From: Zebin Ren Date: Fri, 6 Jun 2025 10:23:49 +0200 Subject: [PATCH 3/4] Expose HSNW index parameters to cli for weaviate --- .../backend/clients/weaviate_cloud/cli.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/vectordb_bench/backend/clients/weaviate_cloud/cli.py b/vectordb_bench/backend/clients/weaviate_cloud/cli.py index 6880411a7..9faf768a6 100644 --- a/vectordb_bench/backend/clients/weaviate_cloud/cli.py +++ b/vectordb_bench/backend/clients/weaviate_cloud/cli.py @@ -25,6 +25,24 @@ class WeaviateTypedDict(CommonTypedDict): 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() @@ -40,6 +58,10 @@ def Weaviate(**parameters: Unpack[WeaviateTypedDict]): url=SecretStr(parameters["url"]), no_auth=parameters["no_auth"], ), - db_case_config=WeaviateIndexConfig(ef=256, efConstruction=256, maxConnections=16), + db_case_config=WeaviateIndexConfig( + efConstruction=parameters["ef_construction"], + maxConnections=parameters["m"], + ef=parameters["ef"], + ), **parameters, ) From 4497dcec86a07e5ffc93b444062c4bee195f2333 Mon Sep 17 00:00:00 2001 From: Zebin Ren Date: Sat, 7 Jun 2025 03:20:24 +0200 Subject: [PATCH 4/4] Add hnsw-ef parameter to Qdrant local --- vectordb_bench/backend/clients/qdrant_local/cli.py | 8 +++++++- vectordb_bench/backend/clients/qdrant_local/config.py | 11 +++++++++-- .../backend/clients/qdrant_local/qdrant_local.py | 8 ++++++++ 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/vectordb_bench/backend/clients/qdrant_local/cli.py b/vectordb_bench/backend/clients/qdrant_local/cli.py index 4a2083e12..c01f0afb7 100644 --- a/vectordb_bench/backend/clients/qdrant_local/cli.py +++ b/vectordb_bench/backend/clients/qdrant_local/cli.py @@ -38,7 +38,12 @@ class QdrantLocalTypedDict(CommonTypedDict): "--ef-construct", type=int, default=200, help="HNSW index parameter ef_construct" ), ] - + hnsw_ef: Annotated[ + int, + click.option( + "--hnsw-ef", type=int, default=0, help="HNSW index parameter hnsw_ef, set 0 to use ef_construct for search", + ), + ] @cli.command() @click_parameter_decorators_from_typed_dict(QdrantLocalTypedDict) @@ -54,6 +59,7 @@ def QdrantLocal(**parameters: Unpack[QdrantLocalTypedDict]): on_disk=parameters["on_disk"], m=parameters["m"], ef_construct=parameters["ef_construct"], + hnsw_ef=parameters["hnsw_ef"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/qdrant_local/config.py b/vectordb_bench/backend/clients/qdrant_local/config.py index b4eb5da47..b2949313f 100644 --- a/vectordb_bench/backend/clients/qdrant_local/config.py +++ b/vectordb_bench/backend/clients/qdrant_local/config.py @@ -15,6 +15,7 @@ class QdrantLocalIndexConfig(BaseModel, DBCaseConfig): metric_type: MetricType | None = None m: int ef_construct: int + hnsw_ef: int | None = 0 on_disk: bool | None = False def parse_metric(self) -> str: @@ -27,7 +28,6 @@ def parse_metric(self) -> str: return "Cosine" def index_param(self) -> dict: - return { "distance": self.parse_metric(), "m": self.m, @@ -36,4 +36,11 @@ def index_param(self) -> dict: } def search_param(self) -> dict: - return {} \ No newline at end of file + search_params = { + "exact": False, # Force to use ANNs + } + + if self.hnsw_ef != 0: + search_params["hnsw_ef"] = self.hnsw_ef + + return search_params \ No newline at end of file diff --git a/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py b/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py index 636d7aa22..723808e8f 100644 --- a/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py +++ b/vectordb_bench/backend/clients/qdrant_local/qdrant_local.py @@ -15,6 +15,7 @@ OptimizersConfigDiff, PayloadSchemaType, Range, + SearchParams, VectorParams, ) @@ -52,6 +53,7 @@ def __init__( 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 @@ -60,6 +62,10 @@ def __init__( 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) @@ -216,6 +222,8 @@ def search_embedding( query=query, limit=k, query_filter=f, + search_params=SearchParams(**self.search_parameter), + ).points )