Skip to content

Commit

Permalink
support alibaba cloud elasticsearch (#418)
Browse files Browse the repository at this point in the history
* add aliyun elasticsearch

* code reuse
  • Loading branch information
xingshaomin authored Dec 11, 2024
1 parent 6a63416 commit 854278a
Show file tree
Hide file tree
Showing 4 changed files with 98 additions and 0 deletions.
13 changes: 13 additions & 0 deletions vectordb_bench/backend/clients/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ class DB(Enum):
MemoryDB = "MemoryDB"
Chroma = "Chroma"
AWSOpenSearch = "OpenSearch"
AliyunElasticsearch = "AliyunElasticsearch"
Test = "test"


Expand Down Expand Up @@ -103,6 +104,10 @@ def init_cls(self) -> Type[VectorDB]:
from .alloydb.alloydb import AlloyDB
return AlloyDB

if self == DB.AliyunElasticsearch:
from .aliyun_elasticsearch.aliyun_elasticsearch import AliyunElasticsearch
return AliyunElasticsearch

@property
def config_cls(self) -> Type[DBConfig]:
"""Import while in use"""
Expand Down Expand Up @@ -166,6 +171,10 @@ def config_cls(self) -> Type[DBConfig]:
from .alloydb.config import AlloyDBConfig
return AlloyDBConfig

if self == DB.AliyunElasticsearch:
from .aliyun_elasticsearch.config import AliyunElasticsearchConfig
return AliyunElasticsearchConfig

def case_config_cls(self, index_type: IndexType | None = None) -> Type[DBCaseConfig]:
if self == DB.Milvus:
from .milvus.config import _milvus_case_config
Expand Down Expand Up @@ -211,6 +220,10 @@ def case_config_cls(self, index_type: IndexType | None = None) -> Type[DBCaseCon
from .alloydb.config import _alloydb_case_config
return _alloydb_case_config.get(index_type)

if self == DB.AliyunElasticsearch:
from .elastic_cloud.config import ElasticCloudIndexConfig
return ElasticCloudIndexConfig

# DB.Pinecone, DB.Chroma, DB.Redis
return EmptyDBCaseConfig

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from ..elastic_cloud.elastic_cloud import ElasticCloud
from ..elastic_cloud.config import ElasticCloudIndexConfig


class AliyunElasticsearch(ElasticCloud):
def __init__(
self,
dim: int,
db_config: dict,
db_case_config: ElasticCloudIndexConfig,
indice: str = "vdb_bench_indice", # must be lowercase
id_col_name: str = "id",
vector_col_name: str = "vector",
drop_old: bool = False,
**kwargs,
):
super().__init__(
dim=dim,
db_config=db_config,
db_case_config=db_case_config,
indice=indice,
id_col_name=id_col_name,
vector_col_name=vector_col_name,
drop_old=drop_old,
**kwargs,
)

19 changes: 19 additions & 0 deletions vectordb_bench/backend/clients/aliyun_elasticsearch/config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
from enum import Enum
from pydantic import SecretStr, BaseModel

from ..api import DBConfig, DBCaseConfig, MetricType, IndexType


class AliyunElasticsearchConfig(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()),
}
39 changes: 39 additions & 0 deletions vectordb_bench/frontend/config/dbCaseConfigs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1040,6 +1040,35 @@ class CaseConfigInput(BaseModel):
},
)

CaseConfigParamInput_EFConstruction_AliES = CaseConfigInput(
label=CaseConfigParamType.EFConstruction,
inputType=InputType.Number,
inputConfig={
"min": 8,
"max": 512,
"value": 360,
},
)

CaseConfigParamInput_M_AliES = CaseConfigInput(
label=CaseConfigParamType.M,
inputType=InputType.Number,
inputConfig={
"min": 4,
"max": 64,
"value": 30,
},
)
CaseConfigParamInput_NumCandidates_AliES = CaseConfigInput(
label=CaseConfigParamType.numCandidates,
inputType=InputType.Number,
inputConfig={
"min": 1,
"max": 10000,
"value": 100,
},
)


MilvusLoadConfig = [
CaseConfigParamInput_IndexType,
Expand Down Expand Up @@ -1206,6 +1235,12 @@ class CaseConfigInput(BaseModel):
CaseConfigParamInput_max_parallel_workers_AlloyDB,
]

AliyunElasticsearchLoadingConfig = [CaseConfigParamInput_EFConstruction_AliES, CaseConfigParamInput_M_AliES]
AliyunElasticsearchPerformanceConfig = [
CaseConfigParamInput_EFConstruction_AliES,
CaseConfigParamInput_M_AliES,
CaseConfigParamInput_NumCandidates_AliES,
]

CASE_CONFIG_MAP = {
DB.Milvus: {
Expand Down Expand Up @@ -1247,4 +1282,8 @@ class CaseConfigInput(BaseModel):
CaseLabel.Load: AlloyDBLoadConfig,
CaseLabel.Performance: AlloyDBPerformanceConfig,
},
DB.AliyunElasticsearch: {
CaseLabel.Load: AliyunElasticsearchLoadingConfig,
CaseLabel.Performance: AliyunElasticsearchPerformanceConfig,
},
}

0 comments on commit 854278a

Please sign in to comment.