From 4117fcd0a6efde11eb8f0d1a499fdd6d9196b85c Mon Sep 17 00:00:00 2001
From: Lanting Chiang <lanting.chiang@robinhood.com>
Date: Mon, 8 Dec 2025 14:33:15 -0500
Subject: [PATCH] make collection name configurable

---
 vectordb_bench/backend/clients/zilliz_cloud/cli.py   | 12 ++++++++++++
 .../backend/clients/zilliz_cloud/config.py           |  2 ++
 vectordb_bench/backend/task_runner.py                |  7 ++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/vectordb_bench/backend/clients/zilliz_cloud/cli.py b/vectordb_bench/backend/clients/zilliz_cloud/cli.py
index 5bc39d810..a8d177ee5 100644
--- a/vectordb_bench/backend/clients/zilliz_cloud/cli.py
+++ b/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,
diff --git a/vectordb_bench/backend/clients/zilliz_cloud/config.py b/vectordb_bench/backend/clients/zilliz_cloud/config.py
index 9f7d2e95a..f0b3fd000 100644
--- a/vectordb_bench/backend/clients/zilliz_cloud/config.py
+++ b/vectordb_bench/backend/clients/zilliz_cloud/config.py
@@ -9,6 +9,7 @@ class ZillizCloudConfig(DBConfig):
     user: str
     password: SecretStr
     num_shards: int = 1
+    collection_name: str = "ZillizCloudVDBBench"
 
     def to_dict(self) -> dict:
         return {
@@ -16,6 +17,7 @@ def to_dict(self) -> dict:
             "user": self.user,
             "password": self.password.get_secret_value(),
             "num_shards": self.num_shards,
+            "collection_name": self.collection_name,
         }
 
 
diff --git a/vectordb_bench/backend/task_runner.py b/vectordb_bench/backend/task_runner.py
index 4a7ba2b3a..8224a0415 100644
--- a/vectordb_bench/backend/task_runner.py
+++ b/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,