feat(oss-opensearch): Add replication type configuration option (#627)

* feat(oss-opensearch): Add replication type configuration option

* fix: format code with linting

Author: Akhil-Pathivada

vectordb_bench/backend/clients/oss_opensearch/config.py

@@ -76,6 +76,7 @@ class OSSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     use_routing: bool = False  # for label-filter cases
     oversample_factor: float = 1.0
     quantization_type: OSSOpenSearchQuantization = OSSOpenSearchQuantization.fp32
+    replication_type: str | None = "DOCUMENT"
 
     @root_validator
     def validate_engine_name(cls, values: dict):
@@ -101,6 +102,7 @@ def __eq__(self, obj: any):
             and self.number_of_segments == obj.number_of_segments
             and self.use_routing == obj.use_routing
             and self.quantization_type == obj.quantization_type
+            and self.replication_type == obj.replication_type
         )
 
     def __hash__(self) -> int:
@@ -114,6 +116,7 @@ def __hash__(self) -> int:
                 self.number_of_segments,
                 self.use_routing,
                 self.quantization_type,
+                self.replication_type,
             )
         )
 

vectordb_bench/backend/clients/oss_opensearch/oss_opensearch.py

@@ -228,6 +228,7 @@ def _create_index(self, client: OpenSearch) -> None:
         ef_search_value = self.case_config.efSearch
         log.info(f"Creating index with ef_search: {ef_search_value}")
         log.info(f"Creating index with number_of_replicas: {self.case_config.number_of_replicas}")
+        log.info(f"Creating index with replication_type: {self.case_config.replication_type}")
         log.info(f"Creating index with engine: {self.case_config.engine}")
         log.info(f"Creating index with metric type: {self.case_config.metric_type_name}")
         log.info(f"All case_config parameters: {self.case_config.__dict__}")
@@ -247,6 +248,7 @@ def _create_index(self, client: OpenSearch) -> None:
                 "number_of_replicas": self.case_config.number_of_replicas,
                 "translog.flush_threshold_size": self.case_config.flush_threshold_size,
                 "knn.advanced.approximate_threshold": "-1",
+                "replication.type": self.case_config.replication_type,
             },
             "refresh_interval": self.case_config.refresh_interval,
         }

vectordb_bench/frontend/config/dbCaseConfigs.py

@@ -1705,6 +1705,17 @@ class CaseConfigInput(BaseModel):
     inputConfig={"value": "60s", "placeholder": "e.g. 30s, 1m"},
 )
 
+CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch = CaseConfigInput(
+    label=CaseConfigParamType.replication_type,
+    displayLabel="Replication Type",
+    inputHelp="Replication strategy: DOCUMENT (default) or SEGMENT",
+    inputType=InputType.Option,
+    inputConfig={
+        "options": ["DOCUMENT", "SEGMENT"],
+        "default": "DOCUMENT",
+    },
+)
+
 MilvusLoadConfig = [
     CaseConfigParamInput_IndexType,
     CaseConfigParamInput_M,
@@ -1785,11 +1796,13 @@ class CaseConfigInput(BaseModel):
 AWSOpensearchLoadingConfig = [
     CaseConfigParamInput_EFConstruction_AWSOpensearch,
     CaseConfigParamInput_M_AWSOpensearch,
+    CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
 ]
 AWSOpenSearchPerformanceConfig = [
     CaseConfigParamInput_EFConstruction_AWSOpensearch,
     CaseConfigParamInput_M_AWSOpensearch,
     CaseConfigParamInput_EF_SEARCH_AWSOpensearch,
+    CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
 ]
 
 AliyunOpensearchLoadingConfig = []
@@ -2095,6 +2108,7 @@ class CaseConfigInput(BaseModel):
     CaseConfigParamInput_NUMBER_OF_REPLICAS_AWSOpensearch,
     CaseConfigParamInput_NUMBER_OF_INDEXING_CLIENTS_AWSOpensearch,
     CaseConfigParamInput_INDEX_THREAD_QTY_AWSOpensearch,
+    CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
     CaseConfigParamInput_INDEX_THREAD_QTY_DURING_FORCE_MERGE_AWSOpensearch,
 ]
 
@@ -2109,6 +2123,7 @@ class CaseConfigInput(BaseModel):
     CaseConfigParamInput_NUMBER_OF_REPLICAS_AWSOpensearch,
     CaseConfigParamInput_NUMBER_OF_INDEXING_CLIENTS_AWSOpensearch,
     CaseConfigParamInput_INDEX_THREAD_QTY_AWSOpensearch,
+    CaseConfigParamInput_REPLICATION_TYPE_AWSOpensearch,
     CaseConfigParamInput_INDEX_THREAD_QTY_DURING_FORCE_MERGE_AWSOpensearch,
 ]
 

vectordb_bench/models.py

@@ -125,6 +125,7 @@ class CaseConfigParamType(Enum):
     use_rescore = "use_rescore"
     oversample_ratio = "oversample_ratio"
     use_routing = "use_routing"
+    replication_type = "replication_type"
 
     dataset_with_size_type = "dataset_with_size_type"
     filter_rate = "filter_rate"