From 56791015552b8c88578aa8d9d27f2ed88da2a6c9 Mon Sep 17 00:00:00 2001 From: Denise2004 <3087753261@qq.com> Date: Tue, 2 Dec 2025 02:41:44 +0000 Subject: [PATCH] cli for elastic --- .../backend/clients/elastic_cloud/cli.py | 288 ++++++++++++++++++ vectordb_bench/cli/vectordbbench.py | 10 + 2 files changed, 298 insertions(+) create mode 100644 vectordb_bench/backend/clients/elastic_cloud/cli.py diff --git a/vectordb_bench/backend/clients/elastic_cloud/cli.py b/vectordb_bench/backend/clients/elastic_cloud/cli.py new file mode 100644 index 000000000..55e521d8f --- /dev/null +++ b/vectordb_bench/backend/clients/elastic_cloud/cli.py @@ -0,0 +1,288 @@ +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.ElasticCloud + + +class ElasticCloudTypedDict(TypedDict): + cloud_id: Annotated[ + str, + click.option("--cloud-id", type=str, help="Elastic Cloud ID", required=True), + ] + password: Annotated[ + str, + click.option("--password", type=str, help="Elastic Cloud password", required=True), + ] + number_of_shards: Annotated[ + int, + click.option( + "--number-of-shards", + type=int, + help="Number of shards", + required=False, + default=1, + show_default=True, + ), + ] + number_of_replicas: Annotated[ + int, + click.option( + "--number-of-replicas", + type=int, + help="Number of replicas", + required=False, + default=0, + show_default=True, + ), + ] + refresh_interval: Annotated[ + str, + click.option( + "--refresh-interval", + type=str, + help="Index refresh interval", + required=False, + default="30s", + show_default=True, + ), + ] + merge_max_thread_count: Annotated[ + int, + click.option( + "--merge-max-thread-count", + type=int, + help="Maximum thread count for merge", + required=False, + default=8, + show_default=True, + ), + ] + use_force_merge: Annotated[ + bool, + click.option( + "--use-force-merge", + type=bool, + help="Whether to use force merge", + required=False, + default=True, + show_default=True, + ), + ] + use_routing: Annotated[ + bool, + click.option( + "--use-routing", + type=bool, + help="Whether to use routing", + required=False, + default=False, + show_default=True, + ), + ] + use_rescore: Annotated[ + bool, + click.option( + "--use-rescore", + type=bool, + help="Whether to use rescore", + required=False, + default=False, + show_default=True, + ), + ] + oversample_ratio: Annotated[ + float, + click.option( + "--oversample-ratio", + type=float, + help="Oversample ratio for rescore", + required=False, + default=2.0, + show_default=True, + ), + ] + + +class ElasticCloudHNSWTypedDict(CommonTypedDict, ElasticCloudTypedDict): + m: Annotated[ + int, + click.option( + "--m", + type=int, + help="HNSW M parameter", + required=False, + default=16, + show_default=True, + ), + ] + ef_construction: Annotated[ + int, + click.option( + "--ef-construction", + type=int, + help="HNSW efConstruction parameter", + required=False, + default=100, + show_default=True, + ), + ] + num_candidates: Annotated[ + int, + click.option( + "--num-candidates", + type=int, + help="Number of candidates for search", + required=False, + default=100, + show_default=True, + ), + ] + element_type: Annotated[ + str, + click.option( + "--element-type", + type=click.Choice(["float", "byte"], case_sensitive=False), + help="Element type for vectors (float: 4 bytes, byte: 1 byte)", + required=False, + default="float", + show_default=True, + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(ElasticCloudHNSWTypedDict) +def ElasticCloudHNSW(**parameters: Unpack[ElasticCloudHNSWTypedDict]): + from ..api import IndexType + from .config import ElasticCloudConfig, ElasticCloudIndexConfig, ESElementType + + run( + db=DBTYPE, + db_config=ElasticCloudConfig( + db_label=parameters["db_label"], + cloud_id=SecretStr(parameters["cloud_id"]), + password=SecretStr(parameters["password"]), + ), + db_case_config=ElasticCloudIndexConfig( + index=IndexType.ES_HNSW, + M=parameters["m"], + efConstruction=parameters["ef_construction"], + num_candidates=parameters["num_candidates"], + element_type=ESElementType(parameters["element_type"]), + number_of_shards=parameters["number_of_shards"], + number_of_replicas=parameters["number_of_replicas"], + refresh_interval=parameters["refresh_interval"], + merge_max_thread_count=parameters["merge_max_thread_count"], + use_force_merge=parameters["use_force_merge"], + use_routing=parameters["use_routing"], + use_rescore=parameters["use_rescore"], + oversample_ratio=parameters["oversample_ratio"], + ), + **parameters, + ) + + +@cli.command() +@click_parameter_decorators_from_typed_dict(ElasticCloudHNSWTypedDict) +def ElasticCloudHNSWInt8(**parameters: Unpack[ElasticCloudHNSWTypedDict]): + from ..api import IndexType + from .config import ElasticCloudConfig, ElasticCloudIndexConfig, ESElementType + + run( + db=DBTYPE, + db_config=ElasticCloudConfig( + db_label=parameters["db_label"], + cloud_id=SecretStr(parameters["cloud_id"]), + password=SecretStr(parameters["password"]), + ), + db_case_config=ElasticCloudIndexConfig( + index=IndexType.ES_HNSW_INT8, + M=parameters["m"], + efConstruction=parameters["ef_construction"], + num_candidates=parameters["num_candidates"], + element_type=ESElementType(parameters["element_type"]), + number_of_shards=parameters["number_of_shards"], + number_of_replicas=parameters["number_of_replicas"], + refresh_interval=parameters["refresh_interval"], + merge_max_thread_count=parameters["merge_max_thread_count"], + use_force_merge=parameters["use_force_merge"], + use_routing=parameters["use_routing"], + use_rescore=parameters["use_rescore"], + oversample_ratio=parameters["oversample_ratio"], + ), + **parameters, + ) + + +@cli.command() +@click_parameter_decorators_from_typed_dict(ElasticCloudHNSWTypedDict) +def ElasticCloudHNSWInt4(**parameters: Unpack[ElasticCloudHNSWTypedDict]): + from ..api import IndexType + from .config import ElasticCloudConfig, ElasticCloudIndexConfig, ESElementType + + run( + db=DBTYPE, + db_config=ElasticCloudConfig( + db_label=parameters["db_label"], + cloud_id=SecretStr(parameters["cloud_id"]), + password=SecretStr(parameters["password"]), + ), + db_case_config=ElasticCloudIndexConfig( + index=IndexType.ES_HNSW_INT4, + M=parameters["m"], + efConstruction=parameters["ef_construction"], + num_candidates=parameters["num_candidates"], + element_type=ESElementType(parameters["element_type"]), + number_of_shards=parameters["number_of_shards"], + number_of_replicas=parameters["number_of_replicas"], + refresh_interval=parameters["refresh_interval"], + merge_max_thread_count=parameters["merge_max_thread_count"], + use_force_merge=parameters["use_force_merge"], + use_routing=parameters["use_routing"], + use_rescore=parameters["use_rescore"], + oversample_ratio=parameters["oversample_ratio"], + ), + **parameters, + ) + + +@cli.command() +@click_parameter_decorators_from_typed_dict(ElasticCloudHNSWTypedDict) +def ElasticCloudHNSWBBQ(**parameters: Unpack[ElasticCloudHNSWTypedDict]): + from ..api import IndexType + from .config import ElasticCloudConfig, ElasticCloudIndexConfig, ESElementType + + run( + db=DBTYPE, + db_config=ElasticCloudConfig( + db_label=parameters["db_label"], + cloud_id=SecretStr(parameters["cloud_id"]), + password=SecretStr(parameters["password"]), + ), + db_case_config=ElasticCloudIndexConfig( + index=IndexType.ES_HNSW_BBQ, + M=parameters["m"], + efConstruction=parameters["ef_construction"], + num_candidates=parameters["num_candidates"], + element_type=ESElementType(parameters["element_type"]), + number_of_shards=parameters["number_of_shards"], + number_of_replicas=parameters["number_of_replicas"], + refresh_interval=parameters["refresh_interval"], + merge_max_thread_count=parameters["merge_max_thread_count"], + use_force_merge=parameters["use_force_merge"], + use_routing=parameters["use_routing"], + use_rescore=parameters["use_rescore"], + oversample_ratio=parameters["oversample_ratio"], + ), + **parameters, + ) diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index 1c8df6c19..8924f5007 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -4,6 +4,12 @@ from ..backend.clients.clickhouse.cli import Clickhouse from ..backend.clients.cockroachdb.cli import CockroachDB as CockroachDBCli from ..backend.clients.doris.cli import Doris +from ..backend.clients.elastic_cloud.cli import ( + ElasticCloudHNSW, + ElasticCloudHNSWBBQ, + ElasticCloudHNSWInt4, + ElasticCloudHNSWInt8, +) from ..backend.clients.hologres.cli import HologresHGraph from ..backend.clients.lancedb.cli import LanceDB from ..backend.clients.mariadb.cli import MariaDBHNSW @@ -53,6 +59,10 @@ cli.add_command(HologresHGraph) cli.add_command(QdrantCloud) cli.add_command(QdrantLocal) +cli.add_command(ElasticCloudHNSW) +cli.add_command(ElasticCloudHNSWInt8) +cli.add_command(ElasticCloudHNSWInt4) +cli.add_command(ElasticCloudHNSWBBQ) cli.add_command(BatchCli) cli.add_command(S3Vectors) cli.add_command(TencentElasticsearch)