From 112d484a00d8474ad730ca1afe001e6b18c71506 Mon Sep 17 00:00:00 2001 From: YangYanbin Date: Wed, 16 Sep 2026 10:26:05 +0800 Subject: [PATCH 1/2] feat: add search level to Milvus AutoIndex Signed-off-by: YangYanbin --- README.md | 4 +++ tests/test_milvus_zilliz_cli.py | 25 +++++++++++++++++++ vectordb_bench/backend/clients/milvus/cli.py | 20 ++++++++++++--- .../backend/clients/milvus/config.py | 10 +++++--- 4 files changed, 51 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 0e9440b2f..e0b8fb6ae 100644 --- a/README.md +++ b/README.md @@ -133,6 +133,10 @@ vectordbbench zillizautoindex \ ``` +`milvusautoindex` accepts `--level` from 1 to 10 (YAML: `level`) on servers that +support AutoIndex search levels. If omitted, no level is sent, preserving the +server's default behavior. Results record it in `task_config.db_case_config.level`. + ```text $ vectordbbench pgvectorhnsw --help Usage: vectordbbench pgvectorhnsw [OPTIONS] diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index 4177c8898..469f410e9 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -55,6 +55,31 @@ def fake_run(**kwargs): assert captured["db_case_config"].use_partition_key is True +def test_milvus_autoindex_cli_sets_search_level( + monkeypatch: MonkeyPatch, +) -> None: + captured = {} + + def fake_run(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(milvus_cli, "run", fake_run) + + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + [ + "--uri", + "http://localhost:19530", + "--level", + "2", + "--dry-run", + ], + ) + + assert result.exit_code == 0, result.output + assert captured["db_case_config"].level == 2 + + def test_zilliz_autoindex_cli_enables_partition_key_for_multitenant_case( monkeypatch: MonkeyPatch, ) -> None: diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 738248ac3..497ee86d0 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -100,7 +100,16 @@ class MilvusTypedDict(TypedDict): ] -class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): ... +class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): + level: Annotated[ + int | None, + click.option( + "--level", + type=click.IntRange(1, 10), + default=None, + help="AutoIndex search level. Requires server support; omitted by default.", + ), + ] @cli.command() @@ -111,14 +120,17 @@ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key(AutoIndexConfig(), parameters), + db_case_config=_with_partition_key(AutoIndexConfig(level=parameters["level"]), parameters), **parameters, ) +class MilvusFlatTypedDict(CommonTypedDict, MilvusTypedDict): ... + + @cli.command() -@click_parameter_decorators_from_typed_dict(MilvusAutoIndexTypedDict) -def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]): +@click_parameter_decorators_from_typed_dict(MilvusFlatTypedDict) +def MilvusFlat(**parameters: Unpack[MilvusFlatTypedDict]): from .config import FLATConfig run( diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 99d5ac8cf..9a7066701 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -1,6 +1,6 @@ from typing import ClassVar -from pydantic import BaseModel, SecretStr +from pydantic import BaseModel, Field, SecretStr from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType @@ -53,6 +53,7 @@ def parse_metric(self) -> str: class AutoIndexConfig(MilvusIndexConfig, DBCaseConfig): index: IndexType = IndexType.AUTOINDEX + level: int | None = Field(default=None, ge=1, le=10) def index_param(self) -> dict: return { @@ -62,9 +63,10 @@ def index_param(self) -> dict: } def search_param(self) -> dict: - return { - "metric_type": self.parse_metric(), - } + params = {"metric_type": self.parse_metric()} + if self.level is not None: + params["params"] = {"level": self.level} + return params class HNSWConfig(MilvusIndexConfig, DBCaseConfig): From aac80192704eeeb6ed8b2a69101b632ad11d9be5 Mon Sep 17 00:00:00 2001 From: YangYanbin Date: Wed, 16 Sep 2026 17:08:32 +0800 Subject: [PATCH 2/2] refactor: share AutoIndex level validation Reuse the validated AutoIndex level contract across Milvus and Zilliz Cloud, add focused compatibility regression coverage, and clarify the CLI defaults in the README. Co-Authored-By: Codex AI-Model: gpt-6-astra AI-Contributed/Feature: 84/84 AI-Contributed/UT: 44/44 Signed-off-by: YangYanbin --- README.md | 17 +++++-- tests/test_milvus_zilliz_cli.py | 44 +++++++++++++++---- vectordb_bench/backend/clients/milvus/cli.py | 12 +---- .../backend/clients/milvus/config.py | 7 ++- .../backend/clients/zilliz_cloud/cli.py | 14 +++--- .../backend/clients/zilliz_cloud/config.py | 22 ++++------ vectordb_bench/cli/cli.py | 12 +++++ 7 files changed, 83 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index e0b8fb6ae..3fd71fae5 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,19 @@ Commands: ``` To list the options for each command, execute `vectordbbench [command] --help` +`milvusautoindex` and `zillizautoindex` accept `--level` from 1 to 10 (YAML: +`level`) on servers that support AutoIndex search levels. If omitted, Milvus +sends no level, preserving the server's default behavior; Zilliz Cloud keeps +its existing default of 1 and always sends it. Results record the value in +`task_config.db_case_config.level`. + +```shell +vectordbbench milvusautoindex \ + --uri http://localhost:19530 \ + --level 2 \ + +``` + Use `--note` or `--note-file` to preserve deployment, resource, client, network, and constraint context in each result JSON under `task_config.db_config.note`. The options are mutually exclusive. Prefer `--note-file` for structured or multiline context, and never include credentials, tokens, or sensitive connection details. ```shell @@ -133,10 +146,6 @@ vectordbbench zillizautoindex \ ``` -`milvusautoindex` accepts `--level` from 1 to 10 (YAML: `level`) on servers that -support AutoIndex search levels. If omitted, no level is sent, preserving the -server's default behavior. Results record it in `task_config.db_case_config.level`. - ```text $ vectordbbench pgvectorhnsw --help Usage: vectordbbench pgvectorhnsw [OPTIONS] diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index 469f410e9..967e5f109 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -1,3 +1,4 @@ +import pytest from click.testing import CliRunner from pytest import MonkeyPatch @@ -55,8 +56,14 @@ def fake_run(**kwargs): assert captured["db_case_config"].use_partition_key is True -def test_milvus_autoindex_cli_sets_search_level( +@pytest.mark.parametrize( + ("level_args", "expected_level"), + [(["--level", "2"], 2), ([], None)], +) +def test_milvus_autoindex_cli_handles_search_level( monkeypatch: MonkeyPatch, + level_args: list[str], + expected_level: int | None, ) -> None: captured = {} @@ -67,17 +74,36 @@ def fake_run(**kwargs): result = CliRunner().invoke( milvus_cli.MilvusAutoIndex, - [ - "--uri", - "http://localhost:19530", - "--level", - "2", - "--dry-run", - ], + ["--uri", "http://localhost:19530", *level_args, "--dry-run"], ) assert result.exit_code == 0, result.output - assert captured["db_case_config"].level == 2 + config = captured["db_case_config"] + assert config.level == expected_level + if expected_level is None: + assert "params" not in config.search_param() + + +def test_milvus_flat_cli_rejects_search_level() -> None: + result = CliRunner().invoke( + milvus_cli.MilvusFlat, + ["--uri", "http://localhost:19530", "--level", "2", "--dry-run"], + ) + + assert result.exit_code == 2 + assert "No such option" in result.output + assert "--level" in result.output + + +@pytest.mark.parametrize("level", ["0", "11"]) +def test_milvus_autoindex_cli_rejects_out_of_range_level(level: str) -> None: + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + ["--uri", "http://localhost:19530", "--level", level, "--dry-run"], + ) + + assert result.exit_code == 2 + assert "1<=x<=10" in result.output def test_zilliz_autoindex_cli_enables_partition_key_for_multitenant_case( diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 497ee86d0..ec11015c8 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -5,6 +5,7 @@ from vectordb_bench.backend.clients import DB from vectordb_bench.cli.cli import ( + AutoIndexLevelTypedDict, CommonTypedDict, HNSWFlavor3, IVFFlatTypedDictN, @@ -100,16 +101,7 @@ class MilvusTypedDict(TypedDict): ] -class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): - level: Annotated[ - int | None, - click.option( - "--level", - type=click.IntRange(1, 10), - default=None, - help="AutoIndex search level. Requires server support; omitted by default.", - ), - ] +class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict, AutoIndexLevelTypedDict): ... @cli.command() diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 9a7066701..caa9a82e2 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -1,4 +1,4 @@ -from typing import ClassVar +from typing import Annotated, ClassVar from pydantic import BaseModel, Field, SecretStr @@ -51,9 +51,12 @@ def parse_metric(self) -> str: return self.metric_type.value +AutoIndexLevel = Annotated[int, Field(ge=1, le=10)] + + class AutoIndexConfig(MilvusIndexConfig, DBCaseConfig): index: IndexType = IndexType.AUTOINDEX - level: int | None = Field(default=None, ge=1, le=10) + level: AutoIndexLevel | None = None def index_param(self) -> dict: return { diff --git a/vectordb_bench/backend/clients/zilliz_cloud/cli.py b/vectordb_bench/backend/clients/zilliz_cloud/cli.py index 030a00661..0a127b249 100644 --- a/vectordb_bench/backend/clients/zilliz_cloud/cli.py +++ b/vectordb_bench/backend/clients/zilliz_cloud/cli.py @@ -6,6 +6,7 @@ from vectordb_bench.backend.clients import DB from vectordb_bench.cli.cli import ( + AutoIndexLevelTypedDict, CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, @@ -49,10 +50,6 @@ class ZillizTypedDict(CommonTypedDict): show_default="$ZILLIZ_TOKEN", ), ] - level: Annotated[ - str, - click.option("--level", type=str, help="Zilliz index level", required=False), - ] num_shards: Annotated[ int, click.option( @@ -88,9 +85,12 @@ class ZillizTypedDict(CommonTypedDict): ] +class ZillizAutoIndexTypedDict(ZillizTypedDict, AutoIndexLevelTypedDict): ... + + @cli.command() -@click_parameter_decorators_from_typed_dict(ZillizTypedDict) -def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]): +@click_parameter_decorators_from_typed_dict(ZillizAutoIndexTypedDict) +def ZillizAutoIndex(**parameters: Unpack[ZillizAutoIndexTypedDict]): from .config import AutoIndexConfig, ZillizCloudConfig run( @@ -105,7 +105,7 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]): collection_name=parameters["collection_name"], ), db_case_config=AutoIndexConfig( - level=int(parameters["level"]) if parameters["level"] else 1, + level=parameters["level"] if parameters["level"] is not None else 1, num_shards=parameters["num_shards"], use_partition_key=_use_partition_key(parameters), ), diff --git a/vectordb_bench/backend/clients/zilliz_cloud/config.py b/vectordb_bench/backend/clients/zilliz_cloud/config.py index ee5ecf645..9db9f2b8f 100644 --- a/vectordb_bench/backend/clients/zilliz_cloud/config.py +++ b/vectordb_bench/backend/clients/zilliz_cloud/config.py @@ -1,7 +1,12 @@ from pydantic import SecretStr -from ..api import DBCaseConfig, DBConfig -from ..milvus.config import IndexType, MilvusFtsConfig, MilvusIndexConfig +from ..api import DBConfig +from ..milvus.config import AutoIndexConfig as MilvusAutoIndexConfig +from ..milvus.config import ( + AutoIndexLevel, + IndexType, + MilvusFtsConfig, +) class ZillizCloudConfig(DBConfig): @@ -27,9 +32,8 @@ def to_dict(self) -> dict: } -class AutoIndexConfig(MilvusIndexConfig, DBCaseConfig): - index: IndexType = IndexType.AUTOINDEX - level: int = 1 +class AutoIndexConfig(MilvusAutoIndexConfig): + level: AutoIndexLevel = 1 num_shards: int = 1 def index_param(self) -> dict: @@ -39,14 +43,6 @@ def index_param(self) -> dict: "params": {"shardsNum": self.num_shards}, } - def search_param(self) -> dict: - return { - "metric_type": self.parse_metric(), - "params": { - "level": self.level, - }, - } - class ZillizCloudFtsConfig(MilvusFtsConfig): index_type: str = IndexType.AUTOINDEX.value diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index c0f3798c9..9f4b5735d 100644 --- a/vectordb_bench/cli/cli.py +++ b/vectordb_bench/cli/cli.py @@ -825,6 +825,18 @@ class HNSWBaseTypedDict(TypedDict): ] +class AutoIndexLevelTypedDict(TypedDict): + level: Annotated[ + int | None, + click.option( + "--level", + type=click.IntRange(1, 10), + default=None, + help="AutoIndex search level (1-10).", + ), + ] + + class HNSWBaseRequiredTypedDict(TypedDict): m: Annotated[int | None, click.option("--m", type=int, help="hnsw m", required=True)] ef_construction: Annotated[