Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions vectordb_bench/backend/clients/zilliz_cloud/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions vectordb_bench/backend/clients/zilliz_cloud/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}


Expand Down
7 changes: 6 additions & 1 deletion vectordb_bench/backend/task_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down