make collection name configurable (#671)

Author: lantingchiang

vectordb_bench/backend/clients/zilliz_cloud/cli.py

@@ -47,6 +47,17 @@ class ZillizTypedDict(CommonTypedDict):
             show_default=True,
         ),
     ]
+    collection_name: Annotated[
+        str,
+        click.option(
+            "--collection-name",
+            type=str,
+            help="Collection name for Zilliz",
+            required=False,
+            default="ZillizCloudVDBBench",
+            show_default=True,
+        ),
+    ]
 
 
 @cli.command()
@@ -62,6 +73,7 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]):
             user=parameters["user_name"],
             password=SecretStr(parameters["password"]),
             num_shards=parameters["num_shards"],
+            collection_name=parameters["collection_name"],
         ),
         db_case_config=AutoIndexConfig(
             level=int(parameters["level"]) if parameters["level"] else 1,

vectordb_bench/backend/clients/zilliz_cloud/config.py

@@ -9,13 +9,15 @@ class ZillizCloudConfig(DBConfig):
     user: str
     password: SecretStr
     num_shards: int = 1
+    collection_name: str = "ZillizCloudVDBBench"
 
     def to_dict(self) -> dict:
         return {
             "uri": self.uri.get_secret_value(),
             "user": self.user,
             "password": self.password.get_secret_value(),
             "num_shards": self.num_shards,
+            "collection_name": self.collection_name,
         }
 
 

vectordb_bench/backend/task_runner.py

@@ -117,9 +117,14 @@ def init_db(self, drop_old: bool = True) -> None:
             # If anything goes wrong, fall back silently; Doris will use its default name logic
             collection_name = None
 
+        # Check if collection_name is in the db_config (e.g., for Zilliz, Milvus)
+        db_config_dict = self.config.db_config.to_dict()
+        if "collection_name" in db_config_dict and not collection_name:
+            collection_name = db_config_dict.pop("collection_name")
+
         self.db = db_cls(
             dim=self.ca.dataset.data.dim,
-            db_config=self.config.db_config.to_dict(),
+            db_config=db_config_dict,
             db_case_config=self.config.db_case_config,
             drop_old=drop_old,
             with_scalar_labels=self.ca.with_scalar_labels,