diff --git a/README.md b/README.md index 9b69e95a0..d1fe6e1cb 100644 --- a/README.md +++ b/README.md @@ -42,23 +42,24 @@ All the database client supported | Optional database client | install command | |--------------------------|---------------------------------------------| | pymilvus, zilliz_cloud (*default*) | `pip install vectordb-bench` | -| all (*clients requirements might be conflict with each other*) | `pip install 'vectordb-bench[all]'` | -| qdrant | `pip install 'vectordb-bench[qdrant]'` | -| pinecone | `pip install 'vectordb-bench[pinecone]'` | -| weaviate | `pip install 'vectordb-bench[weaviate]'` | -| elastic, aliyun_elasticsearch| `pip install 'vectordb-bench[elastic]'` | -| pgvector, pgvectorscale, pgdiskann, alloydb | `pip install 'vectordb-bench[pgvector]'` | -| pgvecto.rs | `pip install 'vectordb-bench[pgvecto_rs]'` | -| redis | `pip install 'vectordb-bench[redis]'` | -| memorydb | `pip install 'vectordb-bench[memorydb]'` | -| chromadb | `pip install 'vectordb-bench[chromadb]'` | -| awsopensearch | `pip install 'vectordb-bench[opensearch]'` | -| aliyun_opensearch | `pip install 'vectordb-bench[aliyun_opensearch]'` | -| mongodb | `pip install 'vectordb-bench[mongodb]'` | -| tidb | `pip install 'vectordb-bench[tidb]'` | -| vespa | `pip install 'vectordb-bench[vespa]'` | -| oceanbase | `pip install 'vectordb-bench[oceanbase]'` | -| hologres | `pip install 'vectordb-bench[hologres]'` | +| all (*clients requirements might be conflict with each other*) | `pip install vectordb-bench[all]` | +| qdrant | `pip install vectordb-bench[qdrant]` | +| pinecone | `pip install vectordb-bench[pinecone]` | +| weaviate | `pip install vectordb-bench[weaviate]` | +| elastic, aliyun_elasticsearch| `pip install vectordb-bench[elastic]` | +| pgvector, pgvectorscale, pgdiskann, alloydb | `pip install vectordb-bench[pgvector]` | +| pgvecto.rs | `pip install vectordb-bench[pgvecto_rs]` | +| redis | `pip install vectordb-bench[redis]` | +| memorydb | `pip install vectordb-bench[memorydb]` | +| chromadb | `pip install vectordb-bench[chromadb]` | +| awsopensearch | `pip install vectordb-bench[opensearch]` | +| aliyun_opensearch | `pip install vectordb-bench[aliyun_opensearch]` | +| mongodb | `pip install vectordb-bench[mongodb]` | +| tidb | `pip install vectordb-bench[tidb]` | +| vespa | `pip install vectordb-bench[vespa]` | +| oceanbase | `pip install vectordb-bench[oceanbase]` | +| hologres | `pip install vectordb-bench[hologres]` | +| tencent_es | `pip install vectordb-bench[tencent_es]` | ### Run diff --git a/install/requirements_py3.11.txt b/install/requirements_py3.11.txt index 0ae328a6f..6b051fc47 100644 --- a/install/requirements_py3.11.txt +++ b/install/requirements_py3.11.txt @@ -3,7 +3,7 @@ grpcio-tools==1.53.0 qdrant-client pinecone-client weaviate-client -elasticsearch +elasticsearch==8.16.0 pgvector pgvecto_rs[psycopg3]>=0.2.1 sqlalchemy diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index 79a6f964a..7ddb04383 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -51,6 +51,7 @@ class DB(Enum): OceanBase = "OceanBase" S3Vectors = "S3Vectors" Hologres = "Alibaba Cloud Hologres" + TencentElasticsearch = "TencentElasticsearch" @property def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 @@ -200,6 +201,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return Hologres + if self == DB.TencentElasticsearch: + from .tencent_elasticsearch.tencent_elasticsearch import TencentElasticsearch + + return TencentElasticsearch + msg = f"Unknown DB: {self.name}" raise ValueError(msg) @@ -351,6 +357,11 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return HologresConfig + if self == DB.TencentElasticsearch: + from .tencent_elasticsearch.config import TencentElasticsearchConfig + + return TencentElasticsearchConfig + msg = f"Unknown DB: {self.name}" raise ValueError(msg) @@ -477,6 +488,11 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912 return HologresIndexConfig + if self == DB.TencentElasticsearch: + from .tencent_elasticsearch.config import TencentElasticsearchIndexConfig + + return TencentElasticsearchIndexConfig + # DB.Pinecone, DB.Chroma, DB.Redis return EmptyDBCaseConfig diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index fdc445618..605e85ac0 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -34,6 +34,7 @@ class IndexType(str, Enum): ES_HNSW_INT8 = "int8_hnsw" ES_HNSW_INT4 = "int4_hnsw" ES_HNSW_BBQ = "bbq_hnsw" + TES_VSEARCH = "vsearch" ES_IVFFlat = "ivfflat" GPU_IVF_FLAT = "GPU_IVF_FLAT" GPU_BRUTE_FORCE = "GPU_BRUTE_FORCE" diff --git a/vectordb_bench/backend/clients/elastic_cloud/config.py b/vectordb_bench/backend/clients/elastic_cloud/config.py index 4d9ec32d4..13826a61e 100644 --- a/vectordb_bench/backend/clients/elastic_cloud/config.py +++ b/vectordb_bench/backend/clients/elastic_cloud/config.py @@ -57,6 +57,7 @@ def __hash__(self) -> int: self.use_routing, self.efConstruction, self.M, + 2, ) ) diff --git a/vectordb_bench/backend/clients/tencent_elasticsearch/cli.py b/vectordb_bench/backend/clients/tencent_elasticsearch/cli.py new file mode 100644 index 000000000..40d78e02c --- /dev/null +++ b/vectordb_bench/backend/clients/tencent_elasticsearch/cli.py @@ -0,0 +1,96 @@ +import os +from typing import Annotated, 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, +) + + +class TencentElasticsearchTypedDict(CommonTypedDict): + scheme: Annotated[ + str, + click.option( + "--scheme", + type=str, + help="Protocol in use to connect to the node", + default="http", + show_default=True, + ), + ] + host: Annotated[ + str, + click.option("--host", type=str, help="shot connection string", required=True), + ] + port: Annotated[ + int, + click.option("--port", type=int, help="Port to connect to", default=9200, show_default=True), + ] + user: Annotated[ + str, + click.option("--user", type=str, help="Db username", required=True), + ] + password: Annotated[ + str, + click.option( + "--password", + type=str, + help="TencentElasticsearch password", + default=lambda: os.environ.get("TES_PASSWORD", ""), + show_default="$TES_PASSWORD", + ), + ] + m: Annotated[ + int, + click.option("--m", type=int, help="HNSW M parameter", default=16, show_default=True), + ] + ef_construction: Annotated[ + int, + click.option( + "--ef_construction", + type=int, + help="HNSW efConstruction parameter", + default=200, + show_default=True, + ), + ] + num_candidates: Annotated[ + int, + click.option( + "--num_candidates", + type=int, + help="Number of candidates to consider during searching", + default=200, + show_default=True, + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(TencentElasticsearchTypedDict) +def TencentElasticsearch(**parameters: Unpack[TencentElasticsearchTypedDict]): + from .config import TencentElasticsearchConfig, TencentElasticsearchIndexConfig + + run( + db=DB.TencentElasticsearch, + db_config=TencentElasticsearchConfig( + db_label=parameters["db_label"], + scheme=parameters["scheme"], + host=parameters["host"], + port=parameters["port"], + user=parameters["user"], + password=SecretStr(parameters["password"]), + ), + db_case_config=TencentElasticsearchIndexConfig( + M=parameters["m"], + efConstruction=parameters["ef_construction"], + num_candidates=parameters["num_candidates"], + ), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/tencent_elasticsearch/config.py b/vectordb_bench/backend/clients/tencent_elasticsearch/config.py new file mode 100644 index 000000000..53a42c058 --- /dev/null +++ b/vectordb_bench/backend/clients/tencent_elasticsearch/config.py @@ -0,0 +1,92 @@ +from enum import Enum + +from pydantic import BaseModel, SecretStr + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType + + +class TencentElasticsearchConfig(DBConfig, BaseModel): + #: Protocol in use to connect to the node + scheme: str = "http" + host: str = "" + port: int = 9200 + user: str = "elastic" + password: SecretStr + + def to_dict(self) -> dict: + return { + "hosts": [{"scheme": self.scheme, "host": self.host, "port": self.port}], + "basic_auth": (self.user, self.password.get_secret_value()), + } + + +class ESElementType(str, Enum): + float = "float" # 4 byte + byte = "byte" # 1 byte, -128 to 127 + + +class TencentElasticsearchIndexConfig(BaseModel, DBCaseConfig): + element_type: ESElementType = ESElementType.float + index: IndexType = IndexType.TES_VSEARCH + number_of_shards: int = 1 + number_of_replicas: int = 0 + refresh_interval: str = "3s" + merge_max_thread_count: int = 8 + use_rescore: bool = False + oversample_ratio: float = 2.0 + use_routing: bool = False + use_force_merge: bool = True + + metric_type: MetricType | None = None + efConstruction: int | None = None + M: int | None = None + num_candidates: int | None = None + + def __eq__(self, obj: any): + return ( + self.index == obj.index + and self.number_of_shards == obj.number_of_shards + and self.number_of_replicas == obj.number_of_replicas + and self.use_routing == obj.use_routing + and self.efConstruction == obj.efConstruction + and self.M == obj.M + ) + + def __hash__(self) -> int: + return hash( + ( + self.index, + self.number_of_shards, + self.number_of_replicas, + self.use_routing, + self.efConstruction, + self.M, + 2, + ) + ) + + def parse_metric(self) -> str: + if self.metric_type == MetricType.L2: + return "l2_norm" + if self.metric_type == MetricType.IP: + return "dot_product" + return "cosine" + + def index_param(self) -> dict: + return { + "type": "dense_vector", + "index": True, + "element_type": self.element_type.value, + "similarity": self.parse_metric(), + "index_options": { + "type": self.index.value, + "index": "hnsw", + "m": self.M, + "ef_construction": self.efConstruction, + }, + } + + def search_param(self) -> dict: + return { + "num_candidates": self.num_candidates, + } diff --git a/vectordb_bench/backend/clients/tencent_elasticsearch/tencent_elasticsearch.py b/vectordb_bench/backend/clients/tencent_elasticsearch/tencent_elasticsearch.py new file mode 100644 index 000000000..887797212 --- /dev/null +++ b/vectordb_bench/backend/clients/tencent_elasticsearch/tencent_elasticsearch.py @@ -0,0 +1,53 @@ +import logging +import time +from contextlib import contextmanager + +from vectordb_bench.backend.filter import Filter, FilterOp + +from ..elastic_cloud.elastic_cloud import ElasticCloud +from .config import TencentElasticsearchIndexConfig + +for logger in ("elasticsearch", "elastic_transport"): + logging.getLogger(logger).setLevel(logging.WARNING) + +log = logging.getLogger(__name__) + + +SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC = 30 + + +class TencentElasticsearch(ElasticCloud): + supported_filter_types: list[FilterOp] = [ + FilterOp.NonFilter, + FilterOp.NumGE, + FilterOp.StrEqual, + ] + + @contextmanager + def init(self) -> None: + """connect to elasticsearch""" + from elasticsearch import Elasticsearch + + self.client = Elasticsearch(**self.db_config, request_timeout=1800) + + yield + self.client = None + del self.client + + def optimize(self, data_size: int | None = None): + """optimize will be called between insertion and search in performance cases.""" + assert self.client is not None, "should self.init() first" + self.client.indices.refresh(index=self.indice) + time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC) + if self.case_config.use_force_merge: + force_merge_task_id = self.client.indices.forcemerge( + index=self.indice, + max_num_segments=1, + wait_for_completion=False, + )["task"] + log.info(f"Elasticsearch force merge task id: {force_merge_task_id}") + while True: + time.sleep(SECONDS_WAITING_FOR_FORCE_MERGE_API_CALL_SEC) + task_status = self.client.tasks.get(task_id=force_merge_task_id) + if task_status["completed"]: + return diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index 83dab74f6..3d8309fa6 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -16,6 +16,7 @@ from ..backend.clients.qdrant_local.cli import QdrantLocal from ..backend.clients.redis.cli import Redis from ..backend.clients.s3_vectors.cli import S3Vectors +from ..backend.clients.tencent_elasticsearch.cli import TencentElasticsearch from ..backend.clients.test.cli import Test from ..backend.clients.tidb.cli import TiDB from ..backend.clients.vespa.cli import Vespa @@ -50,6 +51,7 @@ cli.add_command(QdrantLocal) cli.add_command(BatchCli) cli.add_command(S3Vectors) +cli.add_command(TencentElasticsearch) if __name__ == "__main__": diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index a2bfbac05..5b4c0473b 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -1492,6 +1492,44 @@ class CaseConfigInput(BaseModel): }, ) +CaseConfigParamInput_IndexType_TES = CaseConfigInput( + label=CaseConfigParamType.IndexType, + inputHelp="hnsw or vsearch", + inputType=InputType.Text, + inputConfig={ + "value": "hnsw", + }, +) + +CaseConfigParamInput_EFConstruction_TES = CaseConfigInput( + label=CaseConfigParamType.EFConstruction, + inputType=InputType.Number, + inputConfig={ + "min": 8, + "max": 512, + "value": 360, + }, +) + +CaseConfigParamInput_M_TES = CaseConfigInput( + label=CaseConfigParamType.M, + inputType=InputType.Number, + inputConfig={ + "min": 4, + "max": 64, + "value": 30, + }, +) +CaseConfigParamInput_NumCandidates_TES = CaseConfigInput( + label=CaseConfigParamType.numCandidates, + inputType=InputType.Number, + inputConfig={ + "min": 1, + "max": 10000, + "value": 100, + }, +) + CaseConfigParamInput_IndexType_MariaDB = CaseConfigInput( label=CaseConfigParamType.IndexType, inputHelp="Select Index Type", @@ -1933,6 +1971,17 @@ class CaseConfigInput(BaseModel): CaseConfigParamInput_UseRouting_ES, ] +TencentElasticsearchLoadingConfig = [ + CaseConfigParamInput_EFConstruction_TES, + CaseConfigParamInput_M_TES, + CaseConfigParamInput_IndexType_TES, +] +TencentElasticsearchPerformanceConfig = [ + CaseConfigParamInput_EFConstruction_TES, + CaseConfigParamInput_M_TES, + CaseConfigParamInput_NumCandidates_TES, +] + MongoDBLoadingConfig = [ CaseConfigParamInput_MongoDBQuantizationType, ] @@ -2182,6 +2231,10 @@ class CaseConfigInput(BaseModel): CaseLabel.Load: LanceDBLoadConfig, CaseLabel.Performance: LanceDBPerformanceConfig, }, + DB.TencentElasticsearch: { + CaseLabel.Load: TencentElasticsearchLoadingConfig, + CaseLabel.Performance: TencentElasticsearchPerformanceConfig, + }, }