Skip to content

Commit 9bdeda5

Browse files
feat(oss-opensearch): Add KNN derived source configuration option
1 parent 46c8538 commit 9bdeda5

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, _: 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
@@ -1848,6 +1848,16 @@ class CaseConfigInput(BaseModel):
18481848
},
18491849
)
18501850

1851+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch = CaseConfigInput(
1852+
label=CaseConfigParamType.knn_derived_source_enabled,
1853+
displayLabel="KNN Derived Source Enabled",
1854+
inputHelp="Enable KNN derived source (OpenSearch 3.x+ only). Ignored for 2.x versions.",
1855+
inputType=InputType.Bool,
1856+
inputConfig={
1857+
"value": False,
1858+
},
1859+
)
1860+
18511861
MilvusLoadConfig = [
18521862
CaseConfigParamInput_IndexType,
18531863
CaseConfigParamInput_M,
@@ -1929,12 +1939,14 @@ class CaseConfigInput(BaseModel):
19291939
CaseConfigParamInput_EFConstruction_AWSOpensearch,
19301940
CaseConfigParamInput_M_AWSOpensearch,
19311941
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
1942+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
19321943
]
19331944
AWSOpenSearchPerformanceConfig = [
19341945
CaseConfigParamInput_EFConstruction_AWSOpensearch,
19351946
CaseConfigParamInput_M_AWSOpensearch,
19361947
CaseConfigParamInput_EF_SEARCH_AWSOpensearch,
19371948
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
1949+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
19381950
]
19391951

19401952
AliyunOpensearchLoadingConfig = []
@@ -2274,6 +2286,7 @@ class CaseConfigInput(BaseModel):
22742286
CaseConfigParamInput_EFConstruction_AWSOpensearch,
22752287
CaseConfigParamInput_NUMBER_OF_SHARDS_AWSOpensearch,
22762288
CaseConfigParamInput_NUMBER_OF_REPLICAS_AWSOpensearch,
2289+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
22772290
CaseConfigParamInput_NUMBER_OF_INDEXING_CLIENTS_AWSOpensearch,
22782291
CaseConfigParamInput_INDEX_THREAD_QTY_AWSOpensearch,
22792292
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
@@ -2289,6 +2302,7 @@ class CaseConfigInput(BaseModel):
22892302
CaseConfigParamInput_EFConstruction_AWSOpensearch,
22902303
CaseConfigParamInput_NUMBER_OF_SHARDS_AWSOpensearch,
22912304
CaseConfigParamInput_NUMBER_OF_REPLICAS_AWSOpensearch,
2305+
CaseConfigParamInput_KNN_DERIVED_SOURCE_ENABLED_AWSOpensearch,
22922306
CaseConfigParamInput_NUMBER_OF_INDEXING_CLIENTS_AWSOpensearch,
22932307
CaseConfigParamInput_INDEX_THREAD_QTY_AWSOpensearch,
22942308
CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,

vectordb_bench/models.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ class CaseConfigParamType(Enum):
127127
oversample_ratio = "oversample_ratio"
128128
use_routing = "use_routing"
129129
replication_type = "replication_type"
130+
knn_derived_source_enabled = "knn_derived_source_enabled"
130131

131132
# CockroachDB parameters
132133
min_partition_size = "min_partition_size"

0 commit comments

Comments
 (0)