Skip to content

Commit a431b05

Browse files
feat(oss-opensearch): Add KNN derived source configuration option
1 parent f1c7ad3 commit a431b05

File tree

4 files changed

+24
-0
lines changed

4 files changed

+24
-0
lines changed

vectordb_bench/backend/clients/oss_opensearch/config.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ class OSSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
7777
oversample_factor: float = 1.0
7878
quantization_type: OSSOpenSearchQuantization = OSSOpenSearchQuantization.fp32
7979
replication_type: str | None = "DOCUMENT"
80+
knn_derived_source_enabled: bool = False
8081

8182
@root_validator
8283
def validate_engine_name(cls, values: dict):
@@ -103,6 +104,7 @@ def __eq__(self, obj: any):
103104
and self.use_routing == obj.use_routing
104105
and self.quantization_type == obj.quantization_type
105106
and self.replication_type == obj.replication_type
107+
and self.knn_derived_source_enabled == obj.knn_derived_source_enabled
106108
)
107109

108110
def __hash__(self) -> int:
@@ -117,6 +119,7 @@ def __hash__(self) -> int:
117119
self.use_routing,
118120
self.quantization_type,
119121
self.replication_type,
122+
self.knn_derived_source_enabled,
120123
)
121124
)
122125

vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@
2727
"applies": lambda version, _: version >= Version("3.0"),
2828
"value": lambda _: "-1",
2929
},
30+
{
31+
"name": "knn.derived_source.enabled",
32+
"applies": lambda version, case_config: version >= Version("3.0"),
33+
"value": lambda case_config: case_config.knn_derived_source_enabled,
34+
},
3035
]
3136

3237

@@ -270,6 +275,7 @@ def _create_index(self, client: OpenSearch) -> None:
270275
log.info(f"Creating index with ef_search: {ef_search_value}")
271276
log.info(f"Creating index with number_of_replicas: {self.case_config.number_of_replicas}")
272277
log.info(f"Creating index with replication_type: {self.case_config.replication_type}")
278+
log.info(f"Creating index with knn_derived_source_enabled: {self.case_config.knn_derived_source_enabled}")
273279
log.info(f"Creating index with engine: {self.case_config.engine}")
274280
log.info(f"Creating index with metric type: {self.case_config.metric_type_name}")
275281
log.info(f"All case_config parameters: {self.case_config.__dict__}")

vectordb_bench/frontend/config/dbCaseConfigs.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1802,6 +1802,16 @@ class CaseConfigInput(BaseModel):
18021802
},
18031803
)
18041804

1805+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch = CaseConfigInput(
1806+
label=CaseConfigParamType.knn_derived_source_enabled,
1807+
displayLabel="KNN Derived Source Enabled",
1808+
inputHelp="Enable KNN derived source (OpenSearch 3.x+ only). Ignored for 2.x versions.",
1809+
inputType=InputType.Bool,
1810+
inputConfig={
1811+
"value": False,
1812+
},
1813+
)
1814+
18051815
MilvusLoadConfig = [
18061816
CaseConfigParamInput_IndexType,
18071817
CaseConfigParamInput_M,
@@ -1883,12 +1893,14 @@ class CaseConfigInput(BaseModel):
18831893
CaseConfigParamInput_EFConstruction_AWSOpensearch,
18841894
CaseConfigParamInput_M_AWSOpensearch,
18851895
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
1896+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
18861897
]
18871898
AWSOpenSearchPerformanceConfig = [
18881899
CaseConfigParamInput_EFConstruction_AWSOpensearch,
18891900
CaseConfigParamInput_M_AWSOpensearch,
18901901
CaseConfigParamInput_EF_SEARCH_AWSOpensearch,
18911902
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
1903+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
18921904
]
18931905

18941906
AliyunOpensearchLoadingConfig = []
@@ -2215,6 +2227,7 @@ class CaseConfigInput(BaseModel):
22152227
CaseConfigParamInput_EFConstruction_AWSOpensearch,
22162228
CaseConfigParamInput_NUMBER_OF_SHARDS_AWSOpensearch,
22172229
CaseConfigParamInput_NUMBER_OF_REPLICAS_AWSOpensearch,
2230+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
22182231
CaseConfigParamInput_NUMBER_OF_INDEXING_CLIENTS_AWSOpensearch,
22192232
CaseConfigParamInput_INDEX_THREAD_QTY_AWSOpensearch,
22202233
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
@@ -2230,6 +2243,7 @@ class CaseConfigInput(BaseModel):
22302243
CaseConfigParamInput_EFConstruction_AWSOpensearch,
22312244
CaseConfigParamInput_NUMBER_OF_SHARDS_AWSOpensearch,
22322245
CaseConfigParamInput_NUMBER_OF_REPLICAS_AWSOpensearch,
2246+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
22332247
CaseConfigParamInput_NUMBER_OF_INDEXING_CLIENTS_AWSOpensearch,
22342248
CaseConfigParamInput_INDEX_THREAD_QTY_AWSOpensearch,
22352249
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,

vectordb_bench/models.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ class CaseConfigParamType(Enum):
128128
use_routing = "use_routing"
129129
replication_type = "replication_type"
130130
cache_size = "cache_size"
131+
knn_derived_source_enabled = "knn_derived_source_enabled"
131132

132133
dataset_with_size_type = "dataset_with_size_type"
133134
filter_rate = "filter_rate"

0 commit comments

Comments
 (0)