vectordb-bench 1.0.5__tar.gz → 1.0.7__tar.gz

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vectordb-bench
-Version: 1.0.5
+Version: 1.0.7
 Summary: VectorDBBench is not just an offering of benchmark results for mainstream vector databases and cloud services, it's your go-to tool for the ultimate performance and cost-effectiveness comparison. Designed with ease-of-use in mind, VectorDBBench is devised to help users, even non-professionals, reproduce results or test new systems, making the hunt for the optimal choice amongst a plethora of cloud services and open-source vector databases a breeze.
 Author-email: XuanYang-cn <xuan.yang@zilliz.com>
 Project-URL: repository, https://github.com/zilliztech/VectorDBBench
@@ -153,6 +153,7 @@ All the database client supported
 | tidb                     | `pip install vectordb-bench[tidb]`          |
 | vespa                    | `pip install vectordb-bench[vespa]`         |
 | oceanbase                | `pip install vectordb-bench[oceanbase]`     |
+| hologres                 | `pip install vectordb-bench[hologres]`      |
 
 ### Run
 
@@ -269,7 +270,7 @@ vectordbbench awsopensearch --db-label awsopensearch \
 --m 16 --ef-construction 256 \
 --host search-vector-db-prod-h4f6m4of6x7yp2rz7gdmots7w4.us-west-2.es.amazonaws.com --port 443 \
 --user vector --password '<password>' \
---case-type Performance1536D5M --num-insert-workers 10  \
+--case-type Performance1536D5M --number-of-indexing-clients 10  \
 --skip-load --num-concurrency 75
 ```
 
@@ -297,7 +298,7 @@ Options:
   --force-merge-enabled BOOLEAN   Whether to perform force merge operation
   --flush-threshold-size TEXT     Size threshold for flushing the transaction
                                   log
-  --engine TEXT                   type of engine to use valid values [faiss, lucene]
+  --engine TEXT                   type of engine to use valid values [faiss, lucene, s3vector]
   # Memory Management
   --cb-threshold TEXT             k-NN Memory circuit breaker threshold
   
@@ -372,6 +373,37 @@ Options:
   --help                          Show this message and exit.                       Show this message and exit.
   ```
 
+### Run Hologres from command line
+
+Execute tests for the index types: HGraph.
+
+```shell
+vectordbbench hologreshgraph --host xxx --port xxx --user ACCESS_ID --password ACCESS_KEY --database test \
+--m 64 --ef-construction 400 --case-type Performance768D1M \
+--index-type HGraph --ef-search 51 --k 10
+```
+
+To list the options for Hologres, execute `vectordbbench hologreshgraph --help`, The following are some Hologres-specific command-line options.
+
+```text
+$ vectordbbench hologreshgraph --help
+Usage: vectordbbench hologreshgraph [OPTIONS]
+
+Options:
+  [...]
+  --host TEXT                     Hologres host
+  --user TEXT                     Hologres username  [required]
+  --password TEXT                 Hologres database password
+  --database TEXT                 Hologres database name  [required]
+  --port INTEGER                  Hologres port  [required]
+  --m INTEGER                     hnsw m  [required]
+  --ef-construction INTEGER       hnsw ef-construction  [required]
+  --ef-search INTEGER             hnsw ef-search  [required]
+  --index-type [HGraph]           Type of index to use. Supported values:
+                                  HGraph [required]
+  --help                          Show this message and exit.
+  ```
+
 #### Using a configuration file.
 
 The vectordbbench command can optionally read some or all the options from a yaml formatted configuration file.
@@ -525,9 +557,9 @@ All standard benchmark results are generated by a client running on an 8 core, 3
 2. The next step is to select the test cases you want to perform. You can select multiple cases at once, and a form to collect corresponding parameters will appear.
 3. Finally, you'll need to provide a task label to distinguish different test results. Using the same label for different tests will result in the previous results being overwritten.
 Now we can only run one task at the same time. 
-![image](fig/run_test_select_db.png)
-![image](fig/run_test_select_case.png)
-![image](fig/run_test_submit.png)
+![image](vectordb_bench/fig/run_test_select_db.png)
+![image](vectordb_bench/fig/run_test_select_case.png)
+![image](vectordb_bench/fig/run_test_submit.png)
 
 
 ## Module
@@ -557,8 +589,8 @@ Each case provides an in-depth examination of a vector database's abilities, pro
 
 Through the `/custom` page, users can customize their own performance case using local datasets. After saving, the corresponding case can be selected from the `/run_test` page to perform the test.
 
-![image](fig/custom_dataset.png)
-![image](fig/custom_case_run_test.png)
+![image](vectordb_bench/fig/custom_dataset.png)
+![image](vectordb_bench/fig/custom_case_run_test.png)
 
 We have strict requirements for the data set format, please follow them.
 - `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/README.md

@@ -58,6 +58,7 @@ All the database client supported
 | tidb                     | `pip install vectordb-bench[tidb]`          |
 | vespa                    | `pip install vectordb-bench[vespa]`         |
 | oceanbase                | `pip install vectordb-bench[oceanbase]`     |
+| hologres                 | `pip install vectordb-bench[hologres]`      |
 
 ### Run
 
@@ -174,7 +175,7 @@ vectordbbench awsopensearch --db-label awsopensearch \
 --m 16 --ef-construction 256 \
 --host search-vector-db-prod-h4f6m4of6x7yp2rz7gdmots7w4.us-west-2.es.amazonaws.com --port 443 \
 --user vector --password '<password>' \
---case-type Performance1536D5M --num-insert-workers 10  \
+--case-type Performance1536D5M --number-of-indexing-clients 10  \
 --skip-load --num-concurrency 75
 ```
 
@@ -202,7 +203,7 @@ Options:
   --force-merge-enabled BOOLEAN   Whether to perform force merge operation
   --flush-threshold-size TEXT     Size threshold for flushing the transaction
                                   log
-  --engine TEXT                   type of engine to use valid values [faiss, lucene]
+  --engine TEXT                   type of engine to use valid values [faiss, lucene, s3vector]
   # Memory Management
   --cb-threshold TEXT             k-NN Memory circuit breaker threshold
   
@@ -277,6 +278,37 @@ Options:
   --help                          Show this message and exit.                       Show this message and exit.
   ```
 
+### Run Hologres from command line
+
+Execute tests for the index types: HGraph.
+
+```shell
+vectordbbench hologreshgraph --host xxx --port xxx --user ACCESS_ID --password ACCESS_KEY --database test \
+--m 64 --ef-construction 400 --case-type Performance768D1M \
+--index-type HGraph --ef-search 51 --k 10
+```
+
+To list the options for Hologres, execute `vectordbbench hologreshgraph --help`, The following are some Hologres-specific command-line options.
+
+```text
+$ vectordbbench hologreshgraph --help
+Usage: vectordbbench hologreshgraph [OPTIONS]
+
+Options:
+  [...]
+  --host TEXT                     Hologres host
+  --user TEXT                     Hologres username  [required]
+  --password TEXT                 Hologres database password
+  --database TEXT                 Hologres database name  [required]
+  --port INTEGER                  Hologres port  [required]
+  --m INTEGER                     hnsw m  [required]
+  --ef-construction INTEGER       hnsw ef-construction  [required]
+  --ef-search INTEGER             hnsw ef-search  [required]
+  --index-type [HGraph]           Type of index to use. Supported values:
+                                  HGraph [required]
+  --help                          Show this message and exit.
+  ```
+
 #### Using a configuration file.
 
 The vectordbbench command can optionally read some or all the options from a yaml formatted configuration file.
@@ -430,9 +462,9 @@ All standard benchmark results are generated by a client running on an 8 core, 3
 2. The next step is to select the test cases you want to perform. You can select multiple cases at once, and a form to collect corresponding parameters will appear.
 3. Finally, you'll need to provide a task label to distinguish different test results. Using the same label for different tests will result in the previous results being overwritten.
 Now we can only run one task at the same time. 
-![image](fig/run_test_select_db.png)
-![image](fig/run_test_select_case.png)
-![image](fig/run_test_submit.png)
+![image](vectordb_bench/fig/run_test_select_db.png)
+![image](vectordb_bench/fig/run_test_select_case.png)
+![image](vectordb_bench/fig/run_test_submit.png)
 
 
 ## Module
@@ -462,8 +494,8 @@ Each case provides an in-depth examination of a vector database's abilities, pro
 
 Through the `/custom` page, users can customize their own performance case using local datasets. After saving, the corresponding case can be selected from the `/run_test` page to perform the test.
 
-![image](fig/custom_dataset.png)
-![image](fig/custom_case_run_test.png)
+![image](vectordb_bench/fig/custom_dataset.png)
+![image](vectordb_bench/fig/custom_case_run_test.png)
 
 We have strict requirements for the data set format, please follow them.
 - `Folder Path` - The path to the folder containing all the files. Please ensure that all files in the folder are in the `Parquet` format.

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/vectordb_bench/__init__.py

@@ -16,6 +16,7 @@ class config:
     LOG_LEVEL = env.str("LOG_LEVEL", "INFO")
 
     DEFAULT_DATASET_URL = env.str("DEFAULT_DATASET_URL", AWS_S3_URL)
+    DATASET_SOURCE = env.str("DATASET_SOURCE", "S3")  # Options "S3" or "AliyunOSS"
     DATASET_LOCAL_DIR = env.path("DATASET_LOCAL_DIR", "/tmp/vectordb_bench/dataset")
     NUM_PER_BATCH = env.int("NUM_PER_BATCH", 100)
     TIME_PER_BATCH = 1  # 1s. for streaming insertion.

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/vectordb_bench/backend/clients/__init__.py

@@ -50,6 +50,7 @@ class DB(Enum):
     LanceDB = "LanceDB"
     OceanBase = "OceanBase"
     S3Vectors = "S3Vectors"
+    Hologres = "Alibaba Cloud Hologres"
 
     @property
     def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
@@ -194,6 +195,11 @@ class DB(Enum):
 
             return S3Vectors
 
+        if self == DB.Hologres:
+            from .hologres.hologres import Hologres
+
+            return Hologres
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -340,6 +346,11 @@ class DB(Enum):
 
             return S3VectorsConfig
 
+        if self == DB.Hologres:
+            from .hologres.config import HologresConfig
+
+            return HologresConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -461,6 +472,10 @@ class DB(Enum):
             from .s3_vectors.config import S3VectorsIndexConfig
 
             return S3VectorsIndexConfig
+        if self == DB.Hologres:
+            from .hologres.config import HologresIndexConfig
+
+            return HologresIndexConfig
 
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/vectordb_bench/backend/clients/api.py

@@ -40,6 +40,8 @@ class IndexType(str, Enum):
     GPU_IVF_PQ = "GPU_IVF_PQ"
     GPU_CAGRA = "GPU_CAGRA"
     SCANN = "scann"
+    Hologres_HGraph = "HGraph"
+    Hologres_Graph = "Graph"
     NONE = "NONE"
 
 

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/vectordb_bench/backend/clients/aws_opensearch/aws_opensearch.py

@@ -65,9 +65,7 @@ class AWSOpenSearch(VectorDB):
             self._load_graphs_to_memory(client)
 
     def _create_index(self, client: OpenSearch) -> None:
-        ef_search_value = (
-            self.case_config.ef_search if self.case_config.ef_search is not None else self.case_config.efSearch
-        )
+        ef_search_value = self.case_config.ef_search
         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}")
 
@@ -81,7 +79,7 @@ class AWSOpenSearch(VectorDB):
                 "knn.memory.circuit_breaker.limit": self.case_config.cb_threshold,
             }
         }
-        client.cluster.put_settings(cluster_settings_body)
+        client.cluster.put_settings(body=cluster_settings_body)
         settings = {
             "index": {
                 "knn": True,
@@ -93,25 +91,83 @@ class AWSOpenSearch(VectorDB):
             "refresh_interval": self.case_config.refresh_interval,
         }
         settings["index"]["knn.algo_param.ef_search"] = ef_search_value
-        mappings = {
-            "_source": {"excludes": [self.vector_col_name], "recovery_source_excludes": [self.vector_col_name]},
-            "properties": {
-                self.id_col_name: {"type": "integer", "store": True},
-                self.label_col_name: {"type": "keyword"},
-                self.vector_col_name: {
-                    "type": "knn_vector",
-                    "dimension": self.dim,
-                    "method": self.case_config.index_param(),
-                },
-            },
+
+        # Get method configuration and log it for debugging
+        method_config = self.case_config.index_param()
+        log.info(f"Raw method config from index_param(): {method_config}")
+
+        # For s3vector engine, ensure method only contains engine field
+        if self.case_config.engine == AWSOS_Engine.s3vector:
+            method_config = {"engine": "s3vector"}
+            log.info(f"Cleaned method config for s3vector: {method_config}")
+
+        # Prepare vector field configuration
+        vector_field_config = {
+            "type": "knn_vector",
+            "store": True,
+            "dimension": self.dim,
+            "method": method_config,
         }
+
+        # For s3vector engine, space_type should be set at the vector field level
+        if self.case_config.engine == AWSOS_Engine.s3vector:
+            space_type = self.case_config.parse_metric()
+            vector_field_config["space_type"] = space_type
+
+            # Ensure method config is absolutely clean for s3vector - remove any potential extra fields
+            vector_field_config["method"] = {"engine": "s3vector"}
+
+            log.info(f"Setting space_type '{space_type}' at vector field level for s3vector engine")
+            log.info(f"Final vector field config for s3vector: {vector_field_config}")
+
+        # Configure mappings based on engine type
+        if self.case_config.engine == AWSOS_Engine.s3vector:
+            # For s3vector engine, use simplified mappings without _source configuration
+            mappings = {
+                "properties": {
+                    # self.id_col_name: {"type": "integer", "store": True},
+                    self.label_col_name: {"type": "keyword"},
+                    self.vector_col_name: vector_field_config,
+                },
+            }
+            log.info("Using simplified mappings for s3vector engine (no _source configuration)")
+        else:
+            # For other engines (faiss, lucene), use standard mappings with _source configuration
+            mappings = {
+                "_source": {"excludes": [self.vector_col_name], "recovery_source_excludes": [self.vector_col_name]},
+                "properties": {
+                    # self.id_col_name: {"type": "integer", "store": True},
+                    self.label_col_name: {"type": "keyword"},
+                    self.vector_col_name: vector_field_config,
+                },
+            }
+            log.info("Using standard mappings with _source configuration for non-s3vector engines")
         try:
             log.info(f"Creating index with settings: {settings}")
             log.info(f"Creating index with mappings: {mappings}")
+
+            # Additional logging for s3vector to confirm method config before sending
+            if self.case_config.engine == AWSOS_Engine.s3vector:
+                method_in_mappings = mappings["properties"][self.vector_col_name]["method"]
+                log.info(f"Final method config being sent to OpenSearch: {method_in_mappings}")
+
             client.indices.create(
                 index=self.index_name,
                 body={"settings": settings, "mappings": mappings},
             )
+
+            # For s3vector, verify the actual index configuration after creation
+            if self.case_config.engine == AWSOS_Engine.s3vector:
+                try:
+                    actual_mapping = client.indices.get_mapping(index=self.index_name)
+                    actual_method = actual_mapping[self.index_name]["mappings"]["properties"][self.vector_col_name][
+                        "method"
+                    ]
+                    log.info(f"Actual method config in created index: {actual_method}")
+
+                except Exception as e:
+                    log.warning(f"Failed to verify index configuration: {e}")
+
         except Exception as e:
             log.warning(f"Failed to create index: {self.index_name} error: {e!s}")
             raise e from None
@@ -153,12 +209,12 @@ class AWSOpenSearch(VectorDB):
         insert_data = []
         for i in range(len(embeddings)):
             index_data = {"index": {"_index": self.index_name, self.id_col_name: metadata[i]}}
-            if self.with_scalar_labels and self.case_config.use_routing:
+            if self.with_scalar_labels and self.case_config.use_routing and labels_data is not None:
                 index_data["routing"] = labels_data[i]
             insert_data.append(index_data)
 
             other_data = {self.vector_col_name: embeddings[i]}
-            if self.with_scalar_labels:
+            if self.with_scalar_labels and labels_data is not None:
                 other_data[self.label_col_name] = labels_data[i]
             insert_data.append(other_data)
 
@@ -168,7 +224,7 @@ class AWSOpenSearch(VectorDB):
         except Exception as e:
             log.warning(f"Failed to insert data: {self.index_name} error: {e!s}")
             time.sleep(10)
-            return self._insert_with_single_client(embeddings, metadata)
+            return self._insert_with_single_client(embeddings, metadata, labels_data)
 
     def _insert_with_multiple_clients(
         self,
@@ -186,7 +242,8 @@ class AWSOpenSearch(VectorDB):
 
         for i in range(0, len(embeddings_list), chunk_size):
             end = min(i + chunk_size, len(embeddings_list))
-            chunks.append((embeddings_list[i:end], metadata[i:end], labels_data[i:end]))
+            chunk_labels = labels_data[i:end] if labels_data is not None else None
+            chunks.append((embeddings_list[i:end], metadata[i:end], chunk_labels))
 
         clients = []
         for _ in range(min(num_clients, len(chunks))):
@@ -202,12 +259,12 @@ class AWSOpenSearch(VectorDB):
             insert_data = []
             for i in range(len(chunk_embeddings)):
                 index_data = {"index": {"_index": self.index_name, self.id_col_name: chunk_metadata[i]}}
-                if self.with_scalar_labels and self.case_config.use_routing:
+                if self.with_scalar_labels and self.case_config.use_routing and chunk_labels_data is not None:
                     index_data["routing"] = chunk_labels_data[i]
                 insert_data.append(index_data)
 
                 other_data = {self.vector_col_name: chunk_embeddings[i]}
-                if self.with_scalar_labels:
+                if self.with_scalar_labels and chunk_labels_data is not None:
                     other_data[self.label_col_name] = chunk_labels_data[i]
                 insert_data.append(other_data)
 
@@ -254,10 +311,7 @@ class AWSOpenSearch(VectorDB):
         return (total_count, None)
 
     def _update_ef_search_before_search(self, client: OpenSearch):
-        ef_search_value = (
-            self.case_config.ef_search if self.case_config.ef_search is not None else self.case_config.efSearch
-        )
-
+        ef_search_value = self.case_config.ef_search
         try:
             index_settings = client.indices.get_settings(index=self.index_name)
             current_ef_search = (
@@ -297,23 +351,33 @@ class AWSOpenSearch(VectorDB):
         """
         assert self.client is not None, "should self.init() first"
 
+        # Configure query based on engine type
+        if self.case_config.engine == AWSOS_Engine.s3vector:
+            # For s3vector engine, use simplified query without method_parameters
+            knn_query = {
+                "vector": query,
+                "k": k,
+                **({"filter": self.filter} if self.filter else {}),
+            }
+            log.debug("Using simplified knn query for s3vector engine (no method_parameters)")
+        else:
+            # For other engines (faiss, lucene), use standard query with method_parameters
+            knn_query = {
+                "vector": query,
+                "k": k,
+                "method_parameters": self.case_config.search_param(),
+                **({"filter": self.filter} if self.filter else {}),
+                **(
+                    {"rescore": {"oversample_factor": self.case_config.oversample_factor}}
+                    if self.case_config.use_quant
+                    else {}
+                ),
+            }
+            log.debug("Using standard knn query with method_parameters for non-s3vector engines")
+
         body = {
             "size": k,
-            "query": {
-                "knn": {
-                    self.vector_col_name: {
-                        "vector": query,
-                        "k": k,
-                        "method_parameters": self.case_config.search_param(),
-                        **({"filter": self.filter} if self.filter else {}),
-                        **(
-                            {"rescore": {"oversample_factor": self.case_config.oversample_factor}}
-                            if self.case_config.use_quant
-                            else {}
-                        ),
-                    }
-                }
-            },
+            "query": {"knn": {self.vector_col_name: knn_query}},
         }
 
         try:

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/vectordb_bench/backend/clients/aws_opensearch/cli.py

@@ -46,7 +46,7 @@ class AWSOpenSearchTypedDict(TypedDict):
         str,
         click.option(
             "--engine",
-            type=click.Choice(["nmslib", "faiss", "lucene"], case_sensitive=False),
+            type=click.Choice(["faiss", "lucene", "s3vector"], case_sensitive=False),
             help="HNSW algorithm implementation to use",
             default="faiss",
         ),
@@ -96,24 +96,44 @@ class AWSOpenSearchTypedDict(TypedDict):
         ),
     ]
 
-    quantization_type: Annotated[
-        str | None,
+    index_thread_qty_during_force_merge: Annotated[
+        int,
         click.option(
-            "--quantization-type",
-            type=click.Choice(["fp32", "fp16"]),
-            help="quantization type for vectors (in index)",
-            default="fp32",
+            "--index-thread-qty-during-force-merge",
+            type=int,
+            help="Thread count during force merge operations",
+            default=8,
+        ),
+    ]
+
+    number_of_indexing_clients: Annotated[
+        int,
+        click.option(
+            "--number-of-indexing-clients",
+            type=int,
+            help="Number of concurrent indexing clients",
+            default=1,
+        ),
+    ]
+
+    ef_construction: Annotated[
+        int | None,
+        click.option(
+            "--ef-construction",
+            type=int,
+            help="ef parameter for HNSW construction (not used for s3vector engine)",
+            default=None,
             required=False,
         ),
     ]
 
-    engine: Annotated[
+    quantization_type: Annotated[
         str | None,
         click.option(
-            "--engine",
-            type=click.Choice(["faiss", "lucene"]),
+            "--quantization-type",
+            type=click.Choice(["fp32", "fp16"]),
             help="quantization type for vectors (in index)",
-            default="faiss",
+            default="fp32",
             required=False,
         ),
     ]
@@ -127,6 +147,21 @@ class AWSOpenSearchHNSWTypedDict(CommonTypedDict, AWSOpenSearchTypedDict, HNSWFl
 def AWSOpenSearch(**parameters: Unpack[AWSOpenSearchHNSWTypedDict]):
     from .config import AWSOpenSearchConfig, AWSOpenSearchIndexConfig
 
+    # Set default values for HNSW parameters if not provided and not using s3vector
+    engine = AWSOS_Engine(parameters["engine"])
+    ef_construction = parameters.get("ef_construction")
+    ef_search = parameters.get("ef_search")
+    m = parameters.get("m")
+
+    # For non-s3vector engines, provide defaults if None
+    if engine != AWSOS_Engine.s3vector:
+        if ef_construction is None:
+            ef_construction = 200
+        if ef_search is None:
+            ef_search = 100
+        if m is None:
+            m = 16
+
     run(
         db=DB.AWSOpenSearch,
         db_config=AWSOpenSearchConfig(
@@ -144,12 +179,14 @@ def AWSOpenSearch(**parameters: Unpack[AWSOpenSearchHNSWTypedDict]):
             force_merge_enabled=parameters["force_merge_enabled"],
             flush_threshold_size=parameters["flush_threshold_size"],
             index_thread_qty_during_force_merge=parameters["index_thread_qty_during_force_merge"],
+            number_of_indexing_clients=parameters["number_of_indexing_clients"],
             cb_threshold=parameters["cb_threshold"],
-            efConstruction=parameters["ef_construction"],
-            efSearch=parameters["ef_runtime"],
-            M=parameters["m"],
-            engine=AWSOS_Engine(parameters["engine"]),
+            efConstruction=ef_construction,
+            ef_search=ef_search,
+            M=m,
+            engine=engine,
             quantization_type=AWSOSQuantization(parameters["quantization_type"]),
+            metric_type_name=parameters["metric_type"],
         ),
         **parameters,
     )

{vectordb_bench-1.0.5 → vectordb_bench-1.0.7}/vectordb_bench/backend/clients/aws_opensearch/config.py

@@ -34,6 +34,7 @@ class AWSOpenSearchConfig(DBConfig, BaseModel):
 class AWSOS_Engine(Enum):
     faiss = "faiss"
     lucene = "lucene"
+    s3vector = "s3vector"
 
 
 class AWSOSQuantization(Enum):
@@ -44,11 +45,11 @@ class AWSOSQuantization(Enum):
 class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType = MetricType.L2
     engine: AWSOS_Engine = AWSOS_Engine.faiss
-    efConstruction: int = 256
-    efSearch: int = 100
+    efConstruction: int | None = 256
+    ef_search: int | None = 100
     engine_name: str | None = None
     metric_type_name: str | None = None
-    M: int = 16
+    M: int | None = 16
     index_thread_qty: int | None = 4
     number_of_shards: int | None = 1
     number_of_replicas: int | None = 0
@@ -91,6 +92,13 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
 
     def parse_metric(self) -> str:
         log.info(f"User specified metric_type: {self.metric_type_name}")
+
+        # Handle None or empty metric_type_name
+        if self.metric_type_name is None or self.metric_type_name == "":
+            log.info("No metric_type_name specified, defaulting to l2")
+            self.metric_type = MetricType.L2
+            return "l2"
+
         self.metric_type = MetricType[self.metric_type_name.upper()]
         if self.metric_type == MetricType.IP:
             return "innerproduct"
@@ -108,20 +116,28 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
     def index_param(self) -> dict:
         log.info(f"Using engine: {self.engine} for index creation")
         log.info(f"Using metric_type: {self.metric_type_name} for index creation")
-        log.info(f"Resulting space_type: {self.parse_metric()} for index creation")
+        space_type = self.parse_metric()
+        log.info(f"Resulting space_type: {space_type} for index creation")
+
+        # Handle s3vector engine with simplified configuration
+        # For s3vector, space_type should be set at the vector field level, not in method
+        if self.engine == AWSOS_Engine.s3vector:
+            return {"engine": "s3vector"}
 
         parameters = {"ef_construction": self.efConstruction, "m": self.M}
 
-        if self.engine == AWSOS_Engine.faiss and self.faiss_use_fp16:
+        if self.engine == AWSOS_Engine.faiss and self.quantization_type == AWSOSQuantization.fp16:
             parameters["encoder"] = {"name": "sq", "parameters": {"type": "fp16"}}
 
+        # For other engines (faiss, lucene), space_type is set at method level
         return {
             "name": "hnsw",
             "engine": self.engine.value,
+            "space_type": space_type,
             "parameters": {
                 "ef_construction": self.efConstruction,
                 "m": self.M,
-                "ef_search": self.efSearch,
+                "ef_search": self.ef_search,
                 **(
                     {"encoder": {"name": "sq", "parameters": {"type": self.quantization_type.fp16.value}}}
                     if self.use_quant
@@ -131,4 +147,8 @@ class AWSOpenSearchIndexConfig(BaseModel, DBCaseConfig):
         }
 
     def search_param(self) -> dict:
-        return {"ef_search": self.efSearch}
+        # s3vector engine doesn't use ef_search parameter
+        if self.engine == AWSOS_Engine.s3vector:
+            return {}
+
+        return {"ef_search": self.ef_search}