diff --git a/README.md b/README.md index 0e9440b2f..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 diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index 4177c8898..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,6 +56,56 @@ def fake_run(**kwargs): assert captured["db_case_config"].use_partition_key is True +@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 = {} + + 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_args, "--dry-run"], + ) + + assert result.exit_code == 0, result.output + 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( monkeypatch: MonkeyPatch, ) -> None: diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 738248ac3..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,7 +101,7 @@ class MilvusTypedDict(TypedDict): ] -class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): ... +class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict, AutoIndexLevelTypedDict): ... @cli.command() @@ -111,14 +112,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..caa9a82e2 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 typing import Annotated, ClassVar -from pydantic import BaseModel, SecretStr +from pydantic import BaseModel, Field, SecretStr from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType @@ -51,8 +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: AutoIndexLevel | None = None def index_param(self) -> dict: return { @@ -62,9 +66,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): 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[