From 2cb22a014508effede2cedf6c70c06b92afa7238 Mon Sep 17 00:00:00 2001 From: frankleaf Date: Mon, 8 Jun 2026 12:20:18 +0800 Subject: [PATCH 1/6] feat(aliyun_milvus): add AliyunMilvus DISKANN client with opt-in search params Add an AliyunMilvus client (modeled on the VolcanoMilvus DISKANN client) that exposes DISKANN build knobs and three search-time knobs: rerank_topk_multiplier, early_termination_threshold and cross_segment_rerank. The three search params are optional and only injected into the per-query search params when explicitly specified (None = not sent, server keeps its own default). 0 is preserved as a meaningful value; a negative/"DEFAULT" sentinel from the UI is normalized to None. Wires up CLI command `aliyunmilvusdiskann`, the DB registry, and the Streamlit web UI (dbCaseConfigs + CaseConfigParamType entries). Co-authored-by: Cursor --- vectordb_bench/backend/clients/__init__.py | 16 ++ .../backend/clients/aliyun_milvus/__init__.py | 0 .../clients/aliyun_milvus/aliyun_milvus.py | 193 +++++++++++++++ .../backend/clients/aliyun_milvus/cli.py | 226 ++++++++++++++++++ .../backend/clients/aliyun_milvus/config.py | 147 ++++++++++++ vectordb_bench/cli/vectordbbench.py | 2 + .../frontend/config/dbCaseConfigs.py | 156 ++++++++++++ vectordb_bench/models.py | 15 ++ 8 files changed, 755 insertions(+) create mode 100644 vectordb_bench/backend/clients/aliyun_milvus/__init__.py create mode 100644 vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py create mode 100644 vectordb_bench/backend/clients/aliyun_milvus/cli.py create mode 100644 vectordb_bench/backend/clients/aliyun_milvus/config.py diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index 4be8d0424..3da2cf68c 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -63,6 +63,7 @@ class DB(Enum): PolarDB = "PolarDB" Pinot = "Pinot" SeekDB = "SeekDB" + AliyunMilvus = "AliyunMilvus" @property def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 @@ -269,6 +270,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return SeekDB + if self == DB.AliyunMilvus: + from .aliyun_milvus.aliyun_milvus import AliyunMilvus + + return AliyunMilvus + msg = f"Unknown DB: {self.name}" raise ValueError(msg) @@ -477,6 +483,11 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return SeekDBConfig + if self == DB.AliyunMilvus: + from .aliyun_milvus.config import AliyunMilvusConfig + + return AliyunMilvusConfig + msg = f"Unknown DB: {self.name}" raise ValueError(msg) @@ -667,6 +678,11 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 return _seekdb_case_config.get(index_type) + if self == DB.AliyunMilvus: + from .aliyun_milvus.config import _aliyun_milvus_case_config + + return _aliyun_milvus_case_config.get(index_type) + # DB.Pinecone, DB.Redis return EmptyDBCaseConfig diff --git a/vectordb_bench/backend/clients/aliyun_milvus/__init__.py b/vectordb_bench/backend/clients/aliyun_milvus/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py b/vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py new file mode 100644 index 000000000..844171daf --- /dev/null +++ b/vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py @@ -0,0 +1,193 @@ +"""AliyunMilvus client implementation. + +Extends the base Milvus client with Aliyun-specific behaviors: +- Expose ``load_reqs_size`` to tune insert batch size (avoid gRPC payload limits). +- Expose ``load_after_compaction`` to choose collection load timing: + - ``False`` (default): ``load_collection`` is called right after collection creation; + ``optimize()`` ends with ``refresh_load``. + - ``True``: ``load_collection`` is deferred until after compaction & index wait + in ``optimize()``. +- Apply ``knowhere.*`` collection properties (enable_thp / enable_prefetch) before load. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from pymilvus import DataType, MilvusClient + +from ..milvus.milvus import MILVUS_FORCE_MERGE_TARGET_SIZE_MB, MILVUS_LOAD_REQS_SIZE, Milvus + +if TYPE_CHECKING: + from .config import MilvusIndexConfig + +log = logging.getLogger(__name__) + + +class AliyunMilvus(Milvus): + # Insert concurrency is handled by the framework (ConcurrentInsertRunner). + # This client does not manage its own multiprocessing pool. + thread_safe: bool = True + + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: MilvusIndexConfig, + collection_name: str = "VDBBench", + drop_old: bool = False, + name: str = "AliyunMilvus", + with_scalar_labels: bool = False, + **kwargs, + ): + """Initialize wrapper around the aliyun milvus vector database. + + The vector index is created right after collection creation (same as the + base ``Milvus`` client). Whether ``load_collection`` is called here or + deferred to ``optimize()`` is controlled by ``load_after_compaction``. + """ + self.name = name + self.db_config = db_config + self.case_config = db_case_config + self.collection_name = collection_name + self.load_reqs_size = int(self.db_config.get("load_reqs_size", MILVUS_LOAD_REQS_SIZE)) + self.batch_size = max(1, int(self.load_reqs_size / (dim * 4))) + self.with_scalar_labels = with_scalar_labels + self.load_after_compaction = bool(self.db_config.get("load_after_compaction", False)) + + self._primary_field = "pk" + self._scalar_id_field = "id" + self._scalar_label_field = "label" + self._scalar_payload_label_field = self._scalar_label_field + self._multitenant_partition_key_field = self._scalar_label_field + self.multitenant_tenant_labels: list[str] = kwargs.get("multitenant_tenant_labels", []) + if self.multitenant_tenant_labels: + self._multitenant_partition_key_field = "labels" + if self.with_scalar_labels: + self._scalar_payload_label_field = "scalar_label" + self._vector_field = "vector" + self._vector_index_name = "vector_idx" + self._scalar_id_index_name = "id_sort_idx" + self._scalar_labels_index_name = "labels_idx" + + client = MilvusClient( + uri=self.db_config.get("uri"), + user=self.db_config.get("user"), + password=self.db_config.get("password"), + token=self.db_config.get("token", ""), + timeout=30, + ) + + if drop_old and client.has_collection(self.collection_name): + log.info(f"{self.name} client drop_old collection: {self.collection_name}") + client.drop_collection(self.collection_name) + + if not client.has_collection(self.collection_name): + schema = MilvusClient.create_schema() + schema.add_field(self._primary_field, DataType.INT64, is_primary=True) + schema.add_field(self._scalar_id_field, DataType.INT64) + schema.add_field(self._vector_field, DataType.FLOAT_VECTOR, dim=dim) + + if self.multitenant_tenant_labels: + schema.add_field( + self._multitenant_partition_key_field, + DataType.VARCHAR, + max_length=256, + is_partition_key=True, + ) + + if self.with_scalar_labels: + is_partition_key = db_case_config.use_partition_key + log.info(f"with_scalar_labels, add a new varchar field, as partition_key: {is_partition_key}") + if not self.multitenant_tenant_labels or ( + self._scalar_payload_label_field != self._multitenant_partition_key_field + ): + schema.add_field( + self._scalar_payload_label_field, + DataType.VARCHAR, + max_length=256, + is_partition_key=is_partition_key and not self.multitenant_tenant_labels, + ) + + log.info(f"{self.name} create collection: {self.collection_name}") + + index_params = self._build_index_params() + client.create_collection( + collection_name=self.collection_name, + schema=schema, + num_shards=self.db_config.get("num_shards", 1), + consistency_level="Session", + ) + client.create_index(self.collection_name, index_params) + if not self.load_after_compaction: + self._apply_load_with_properties(client) + + client.close() + + def need_normalize_cosine(self) -> bool: + """Wheather this database need to normalize dataset to support COSINE""" + log.info("cosine dataset need normalize.") + return True + + def _apply_load_with_properties(self, client: MilvusClient) -> None: + """Apply Aliyun-specific knowhere.* properties and load the collection. + + The properties (currently ``knowhere.enable_thp`` / ``knowhere.enable_prefetch``) + have to be set via ``alter_collection_properties`` before ``load_collection`` + so the segment loader picks them up. + """ + load_param = getattr(self.case_config, "load_param", None) + properties = load_param() if callable(load_param) else {} + if properties: + client.alter_collection_properties( + collection_name=self.collection_name, + properties=properties, + ) + + collection_properties = client.describe_collection(self.collection_name).get("properties") + log.info(f"set collection properties to: {collection_properties}") + + client.load_collection( + self.collection_name, + replica_number=self.db_config.get("replica_number", 1), + ) + + def _optimize(self): + """Aliyun-specific optimize. + + Overrides the parent ``_optimize`` to control collection-load timing: + - ``load_after_compaction=True``: collection is loaded here (after compaction + and index wait) via ``_apply_load_with_properties``; + - ``load_after_compaction=False`` (default): collection was already loaded in + ``__init__``, so we end with ``refresh_load`` like the upstream Milvus. + """ + log.info(f"{self.name} optimizing before search") + try: + self.client.flush(self.collection_name) + + try: + self._wait_for_segments_sorted() + compaction_id = self.client.compact( + self.collection_name, + target_size=MILVUS_FORCE_MERGE_TARGET_SIZE_MB, + ) + if compaction_id > 0: + self._wait_for_compaction(compaction_id) + log.info(f"{self.name} force merge compaction completed.") + except Exception as e: + log.warning(f"{self.name} compact or list segments error: {e}") + if getattr(getattr(e, "code", None), "name", None) == "PERMISSION_DENIED": + log.warning("Skip compact due to list segments or compact permission denied.") + else: + raise e from None + + # wait for index no matter what + self._wait_for_index() + if self.load_after_compaction: + self._apply_load_with_properties(self.client) + else: + self.client.refresh_load(self.collection_name) + except Exception as e: + log.warning(f"{self.name} optimize error: {e}") + raise e from None diff --git a/vectordb_bench/backend/clients/aliyun_milvus/cli.py b/vectordb_bench/backend/clients/aliyun_milvus/cli.py new file mode 100644 index 000000000..751ffc4e2 --- /dev/null +++ b/vectordb_bench/backend/clients/aliyun_milvus/cli.py @@ -0,0 +1,226 @@ +from typing import Annotated, Unpack + +import click +from pydantic import SecretStr + +from vectordb_bench.backend.clients import DB +from vectordb_bench.cli.cli import ( + CommonTypedDict, + cli, + click_parameter_decorators_from_typed_dict, + run, +) + +from ..milvus.cli import MilvusTypedDict + +DBTYPE = DB.AliyunMilvus + + +class AliyunMilvusTypedDict(MilvusTypedDict): + load_reqs_size: Annotated[ + int, + click.option( + "--load-reqs-size", + type=int, + help="Max request payload size (bytes) used to compute insert batch size", + required=False, + default=int(1.5 * 1024 * 1024), + show_default=True, + ), + ] + load_after_compaction: Annotated[ + bool, + click.option( + "--load-after-compaction/--no-load-after-compaction", + "load_after_compaction", + help=( + "If True, defer load_collection until after compaction & index build in optimize(); " + "if False (default), load_collection right after collection creation and call " + "refresh_load at the end of optimize()." + ), + default=False, + show_default=True, + ), + ] + + +class AliyunMilvusDISKANNTypedDict(CommonTypedDict, AliyunMilvusTypedDict): + max_degree: Annotated[ + int, + click.option( + "--max-degree", + type=int, + help="R (max degree)", + required=False, + default=56, + show_default=True, + ), + ] + search_list: Annotated[ + int, + click.option( + "--search-list", + type=int, + help="L (search list size) used during DISKANN index search", + required=False, + default=200, + show_default=True, + ), + ] + build_search_list: Annotated[ + int, + click.option( + "--build-search-list", + type=int, + help="L (search list size) used during DISKANN index build", + required=False, + default=200, + show_default=True, + ), + ] + legacy: Annotated[ + bool, + click.option( + "--legacy/--no-legacy", + "legacy", + help="Use legacy Aliyun DISKANN behavior", + default=False, + show_default=True, + ), + ] + store_strategy: Annotated[ + str, + click.option( + "--store-strategy", + type=click.Choice(["MEMORY", "DISK"], case_sensitive=False), + help="Aliyun store strategy", + required=False, + default="MEMORY", + show_default=True, + ), + ] + quant_type: Annotated[ + str, + click.option( + "--quant-type", + type=click.Choice(["RABITQ", "PQ"], case_sensitive=False), + help="Aliyun quant type", + required=False, + default="RABITQ", + show_default=True, + ), + ] + num_threads: Annotated[ + int, + click.option( + "--num-threads", + type=int, + help="Degree of parallelism", + required=False, + default=4, + show_default=True, + ), + ] + distance_strategy: Annotated[ + str, + click.option( + "--distance-strategy", + type=click.Choice( + ["FULL", "SINGLE QUANT", "QUANT THEN FULL", "QUANT THEN MORE BITS"], + case_sensitive=False, + ), + help="Aliyun distance strategy", + required=False, + default="QUANT THEN MORE BITS", + show_default=True, + ), + ] + enable_prefetch: Annotated[ + bool, + click.option( + "--enable-prefetch/--disable-prefetch", + "enable_prefetch", + help="Enable Aliyun DISKANN prefetch during search", + default=False, + show_default=True, + ), + ] + enable_thp: Annotated[ + bool, + click.option( + "--enable-thp/--disable-thp", + "enable_thp", + help="Enable Aliyun DISKANN transparent huge pages on collection load", + default=False, + show_default=True, + ), + ] + rerank_topk_multiplier: Annotated[ + int | None, + click.option( + "--rerank-topk-multiplier", + type=int, + help=( + "Search param: multiplier of topk for the rerank candidate budget " + "(0 disables rerank read). Omit to NOT send this search param." + ), + required=False, + default=None, + ), + ] + early_termination_threshold: Annotated[ + int | None, + click.option( + "--early-termination-threshold", + type=int, + help="Search param: early termination threshold (0 disables). Omit to NOT send this search param.", + required=False, + default=None, + ), + ] + cross_segment_rerank: Annotated[ + bool | None, + click.option( + "--cross-segment-rerank/--no-cross-segment-rerank", + "cross_segment_rerank", + help="Search param: enable cross-segment rerank in the reduce layer. Omit to NOT send this search param.", + default=None, + ), + ] + + +@cli.command(name="aliyunmilvusdiskann") +@click_parameter_decorators_from_typed_dict(AliyunMilvusDISKANNTypedDict) +def AliyunMilvusDISKANN(**parameters: Unpack[AliyunMilvusDISKANNTypedDict]): + from .config import AliyunMilvusConfig, AliyunMilvusDISKANNConfig + + run( + db=DBTYPE, + db_config=AliyunMilvusConfig( + db_label=parameters["db_label"], + uri=SecretStr(parameters["uri"]), + user=parameters["user_name"], + password=SecretStr(parameters["password"]) if parameters["password"] else None, + num_shards=int(parameters["num_shards"]), + replica_number=int(parameters["replica_number"]), + load_reqs_size=int(parameters["load_reqs_size"]), + load_after_compaction=bool(parameters["load_after_compaction"]), + ), + db_case_config=AliyunMilvusDISKANNConfig( + search_list=parameters["search_list"], + build_search_list=parameters["build_search_list"], + max_degree=parameters["max_degree"], + legacy=parameters["legacy"], + store_strategy=parameters["store_strategy"], + quant_type=parameters["quant_type"], + num_threads=parameters["num_threads"], + distance_strategy=parameters["distance_strategy"], + enable_prefetch=bool(parameters["enable_prefetch"]), + enable_thp=bool(parameters["enable_thp"]), + # Pass through as-is; None means "not specified" -> omitted from search params. + rerank_topk_multiplier=parameters["rerank_topk_multiplier"], + early_termination_threshold=parameters["early_termination_threshold"], + cross_segment_rerank=parameters["cross_segment_rerank"], + ), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/aliyun_milvus/config.py b/vectordb_bench/backend/clients/aliyun_milvus/config.py new file mode 100644 index 000000000..42560a7b2 --- /dev/null +++ b/vectordb_bench/backend/clients/aliyun_milvus/config.py @@ -0,0 +1,147 @@ +from typing import ClassVar + +from pydantic import SecretStr, field_validator + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType +from ..milvus.config import DISKANNConfig, MilvusIndexConfig + + +class AliyunMilvusConfig(DBConfig): + _extra_empty_skip: ClassVar[frozenset[str]] = frozenset({"user", "password"}) + + uri: SecretStr = SecretStr("http://localhost:19530") + user: str | None = None + password: SecretStr | None = None + num_shards: int = 1 + replica_number: int = 1 + # Tuning knobs for load performance + load_reqs_size: int = int(1.5 * 1024 * 1024) + # Controls when load_collection runs: + # - False (default): load immediately after collection creation, then refresh_load in optimize(). + # - True: defer load until after compaction and index build in optimize(). + load_after_compaction: bool = False + + def to_dict(self) -> dict: + return { + "uri": self.uri.get_secret_value(), + "user": self.user if self.user else None, + "password": self.password.get_secret_value() if self.password else None, + "num_shards": self.num_shards, + "replica_number": self.replica_number, + "load_reqs_size": self.load_reqs_size, + "load_after_compaction": self.load_after_compaction, + } + + +class AliyunMilvusDISKANNConfig(DISKANNConfig): + """AliyunMilvus DISKANN index config. + + Inherits from Milvus ``DISKANNConfig`` and adds Aliyun-specific fields. + Overrides ``search_param`` to inject the three search-time knobs + (``rerank_topk_multiplier`` / ``early_termination_threshold`` / + ``cross_segment_rerank``) into the search ``params`` while keeping the + original contract. + """ + + # ---- index build knobs ---- + max_degree: int = 48 + legacy: bool = False + store_strategy: str = "MEMORY" + quant_type: str = "RABITQ" + num_threads: int = 4 + distance_strategy: str = "QUANT THEN MORE BITS" + enable_prefetch: bool = False + enable_thp: bool = False + build_search_list: int = 200 + + # ---- search-time knobs ---- + # Optional on purpose: a value of None means "do not send this search param" + # (the server keeps its own default). They are only injected into search + # params when explicitly set. Note ``0`` is a meaningful value (e.g. + # rerank_topk_multiplier=0 disables rerank reads), so "unset" must be None, + # not 0. UI passes a negative sentinel for "unset", normalized to None below. + rerank_topk_multiplier: int | None = None + early_termination_threshold: int | None = None + cross_segment_rerank: bool | None = None + + @field_validator("rerank_topk_multiplier", "early_termination_threshold", mode="before") + @classmethod + def _normalize_optional_int(cls, v: object) -> int | None: + if v is None or v == "" or v == "DEFAULT": + return None + iv = int(v) + # negative value is the UI "unset" sentinel; 0 stays a real value + return None if iv < 0 else iv + + @field_validator("cross_segment_rerank", mode="before") + @classmethod + def _normalize_optional_bool(cls, v: object) -> bool | None: + if v is None or v == "" or v == "DEFAULT": + return None + if isinstance(v, bool): + return v + if isinstance(v, str): + return v.strip().lower() in ("true", "1", "yes", "on") + return bool(v) + + def parse_metric(self) -> str: + if not self.metric_type: + return "" + + if self.metric_type == MetricType.COSINE: + return MetricType.L2.value + return self.metric_type.value + + def index_param(self) -> dict: + extra_params: dict = { + "max_degree": self.max_degree, + "legacy": self.legacy, + "store_strategy": self.store_strategy, + "quant_type": self.quant_type, + "num_threads": self.num_threads, + "distance_strategy": self.distance_strategy, + "search_list_size": self.build_search_list, + } + return { + "metric_type": self.parse_metric(), + "index_type": self.index.value, + "params": extra_params, + } + + def search_param(self) -> dict: + # Only inject a search param when it was explicitly specified. + params: dict = {} + if self.search_list is not None: + params["search_list"] = self.search_list + if self.rerank_topk_multiplier is not None: + params["rerank_topk_multiplier"] = self.rerank_topk_multiplier + if self.early_termination_threshold is not None: + params["early_termination_threshold"] = self.early_termination_threshold + if self.cross_segment_rerank is not None: + params["cross_segment_rerank"] = self.cross_segment_rerank + return { + "metric_type": self.parse_metric(), + "params": params, + } + + def load_param(self) -> dict: + return { + "knowhere.enable_thp": str(self.enable_thp).lower(), + "knowhere.enable_prefetch": str(self.enable_prefetch).lower(), + } + + +# Only DISKANN is supported by AliyunMilvus today. Other index types are not +# exposed here so callers fail fast (case_config_cls returns None) instead of +# silently using the upstream Milvus index implementation. +_aliyun_milvus_case_config: dict[IndexType, type[DBCaseConfig]] = { + IndexType.DISKANN: AliyunMilvusDISKANNConfig, +} + + +__all__ = [ + "MilvusIndexConfig", + "AliyunMilvusConfig", + "AliyunMilvusDISKANNConfig", + "_aliyun_milvus_case_config", +] diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index 13c9687c7..2269abb2f 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -1,4 +1,5 @@ from ..backend.clients.alisql.cli import AliSQLHNSW +from ..backend.clients.aliyun_milvus.cli import AliyunMilvusDISKANN from ..backend.clients.alloydb.cli import AlloyDBScaNN from ..backend.clients.aws_opensearch.cli import AWSOpenSearch from ..backend.clients.chroma.cli import Chroma @@ -98,6 +99,7 @@ cli.add_command(PolarDBHNSWPQ) cli.add_command(PolarDBHNSWSQ) cli.add_command(SeekDBHNSW) +cli.add_command(AliyunMilvusDISKANN) if __name__ == "__main__": diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index d15c4e7ee..c0959f6e4 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -2168,6 +2168,157 @@ class CaseConfigInput(BaseModel): CaseConfigParamInput_Milvus_use_partition_key, ] +# ---- AliyunMilvus ---- +# Reuse Milvus UI configs for all index types except DISKANN, where +# Aliyun-specific build params and the three search-time knobs +# (rerank_topk_multiplier / early_termination_threshold / cross_segment_rerank) +# are shown in addition to Milvus's SearchList. + +CaseConfigParamInput_IndexType_AliyunMilvus = CaseConfigInput( + label=CaseConfigParamType.IndexType, + inputType=InputType.Option, + inputHelp="AliyunMilvus currently supports DISKANN only", + inputConfig={ + "options": [IndexType.DISKANN.value], + }, +) + +CaseConfigParamInput_Aliyun_max_degree = CaseConfigInput( + label=CaseConfigParamType.max_degree, + inputType=InputType.Number, + inputConfig={"min": 1, "max": 65536, "value": 56}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_legacy = CaseConfigInput( + label=CaseConfigParamType.legacy, + inputType=InputType.Bool, + displayLabel="Legacy", + inputHelp="Use legacy Aliyun DISKANN behavior", + inputConfig={"value": False}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_store_strategy = CaseConfigInput( + label=CaseConfigParamType.store_strategy, + inputType=InputType.Option, + inputConfig={"options": ["MEMORY", "DISK"], "value": "MEMORY"}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_quant_type = CaseConfigInput( + label=CaseConfigParamType.quant_type, + inputType=InputType.Option, + inputConfig={"options": ["RABITQ", "PQ"], "value": "RABITQ"}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_num_threads = CaseConfigInput( + label=CaseConfigParamType.num_threads, + inputType=InputType.Number, + inputConfig={"min": 1, "max": 1024, "value": 4}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_distance_strategy = CaseConfigInput( + label=CaseConfigParamType.distance_strategy, + inputType=InputType.Option, + inputConfig={ + "options": ["FULL", "SINGLE QUANT", "QUANT THEN FULL", "QUANT THEN MORE BITS"], + "value": "QUANT THEN MORE BITS", + }, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_enable_prefetch = CaseConfigInput( + label=CaseConfigParamType.enable_prefetch, + inputType=InputType.Bool, + displayLabel="Enable Prefetch", + inputHelp="Enable Aliyun DISKANN prefetch during search", + inputConfig={"value": False}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_enable_thp = CaseConfigInput( + label=CaseConfigParamType.enable_thp, + inputType=InputType.Bool, + displayLabel="Enable THP", + inputHelp="Enable Aliyun DISKANN transparent huge pages on collection load", + inputConfig={"value": False}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_BuildSearchList = CaseConfigInput( + label=CaseConfigParamType.BuildSearchList, + inputType=InputType.Number, + inputConfig={ + "min": 1, + "max": MAX_STREAMLIT_INT, + "value": 200, + }, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +# The three search-time knobs default to "unset" -> the param is NOT sent to the +# server (it keeps its own default). For Number inputs a negative value (-1) is +# the "unset" sentinel; for the bool, "DEFAULT" is the "unset" sentinel. The +# backend config normalizes these to None (note: 0 is a real, meaningful value). +CaseConfigParamInput_Aliyun_rerank_topk_multiplier = CaseConfigInput( + label=CaseConfigParamType.rerank_topk_multiplier, + inputType=InputType.Number, + displayLabel="Rerank TopK Multiplier", + inputHelp="Search param: multiplier of topk for rerank budget (0 disables rerank read). -1 = not specified (omit).", + inputConfig={"min": -1, "max": MAX_STREAMLIT_INT, "value": -1}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_early_termination_threshold = CaseConfigInput( + label=CaseConfigParamType.early_termination_threshold, + inputType=InputType.Number, + displayLabel="Early Termination Threshold", + inputHelp="Search param: early termination threshold (0 disables). -1 = not specified (omit).", + inputConfig={"min": -1, "max": MAX_STREAMLIT_INT, "value": -1}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +CaseConfigParamInput_Aliyun_cross_segment_rerank = CaseConfigInput( + label=CaseConfigParamType.cross_segment_rerank, + inputType=InputType.Option, + displayLabel="Cross Segment Rerank", + inputHelp="Search param: enable cross-segment rerank. DEFAULT = not specified (omit).", + inputConfig={"options": ["DEFAULT", "True", "False"]}, + isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, +) + +_AliyunMilvusDiskannBuildParams = [ + CaseConfigParamInput_Aliyun_max_degree, + CaseConfigParamInput_Aliyun_legacy, + CaseConfigParamInput_Aliyun_store_strategy, + CaseConfigParamInput_Aliyun_quant_type, + CaseConfigParamInput_Aliyun_num_threads, + CaseConfigParamInput_Aliyun_distance_strategy, + CaseConfigParamInput_Aliyun_enable_prefetch, + CaseConfigParamInput_Aliyun_enable_thp, + CaseConfigParamInput_Aliyun_BuildSearchList, +] +_AliyunMilvusDiskannSearchParams = [ + CaseConfigParamInput_Aliyun_rerank_topk_multiplier, + CaseConfigParamInput_Aliyun_early_termination_threshold, + CaseConfigParamInput_Aliyun_cross_segment_rerank, +] + +AliyunMilvusLoadConfig = [ + CaseConfigParamInput_IndexType_AliyunMilvus, + *MilvusLoadConfig[1:], + *_AliyunMilvusDiskannBuildParams, +] +AliyunMilvusPerformanceConfig = [ + CaseConfigParamInput_IndexType_AliyunMilvus, + *MilvusPerformanceConfig[1:], + *_AliyunMilvusDiskannBuildParams, + *_AliyunMilvusDiskannSearchParams, +] + WeaviateLoadConfig = [ CaseConfigParamInput_MaxConnections, CaseConfigParamInput_EFConstruction_Weaviate, @@ -3064,6 +3215,11 @@ class FilterType(Enum): CaseLabel.Load: PolarDBConfig, CaseLabel.Performance: PolarDBConfig, }, + DB.AliyunMilvus: { + CaseLabel.Load: AliyunMilvusLoadConfig, + CaseLabel.Performance: AliyunMilvusPerformanceConfig, + CaseLabel.Streaming: AliyunMilvusPerformanceConfig, + }, } diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index dc1709cc0..70f638dad 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -51,6 +51,7 @@ class CaseConfigParamType(Enum): ef_construction = "ef_construction" EF = "ef" SearchList = "search_list" + BuildSearchList = "build_search_list" ef_search = "ef_search" Nlist = "nlist" Nprobe = "nprobe" @@ -175,6 +176,20 @@ class CaseConfigParamType(Enum): exbits = "exbits" number_of_regions = "number_of_regions" + # AliyunMilvus DISKANN parameters + max_degree = "max_degree" + legacy = "legacy" + store_strategy = "store_strategy" + quant_type = "quant_type" + num_threads = "num_threads" + distance_strategy = "distance_strategy" + enable_prefetch = "enable_prefetch" + enable_thp = "enable_thp" + # AliyunMilvus DISKANN search-time parameters + rerank_topk_multiplier = "rerank_topk_multiplier" + early_termination_threshold = "early_termination_threshold" + cross_segment_rerank = "cross_segment_rerank" + class CustomizedCase(BaseModel): pass From 67bf7409d2aa32e4d723af0b3293fd0766e78b74 Mon Sep 17 00:00:00 2001 From: frankleaf Date: Mon, 8 Jun 2026 12:31:25 +0800 Subject: [PATCH 2/6] refactor(aliyun_milvus): minimize to Milvus DISKANN + opt-in search params Drop the Volcano-style scaffolding that is not needed for this use case: remove the dedicated AliyunMilvus client class and db config, and all the extra index-build knobs (max_degree/store_strategy/quant_type/ distance_strategy/enable_thp/enable_prefetch/build_search_list/legacy). AliyunMilvus now reuses the upstream Milvus client and MilvusConfig as-is; the only difference is AliyunMilvusDISKANNConfig, which extends the existing DISKANNConfig with the three opt-in search params and injects them into the search params only when explicitly specified. Co-authored-by: Cursor --- vectordb_bench/backend/clients/__init__.py | 10 +- .../clients/aliyun_milvus/aliyun_milvus.py | 193 ------------------ .../backend/clients/aliyun_milvus/cli.py | 169 +-------------- .../backend/clients/aliyun_milvus/config.py | 126 ++---------- .../frontend/config/dbCaseConfigs.py | 103 +--------- vectordb_bench/models.py | 12 +- 6 files changed, 42 insertions(+), 571 deletions(-) delete mode 100644 vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index 3da2cf68c..fdc8f64a1 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -271,9 +271,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return SeekDB if self == DB.AliyunMilvus: - from .aliyun_milvus.aliyun_milvus import AliyunMilvus + # AliyunMilvus reuses the upstream Milvus client; only the DISKANN + # case config differs (three opt-in search params). + from .milvus.milvus import Milvus - return AliyunMilvus + return Milvus msg = f"Unknown DB: {self.name}" raise ValueError(msg) @@ -484,9 +486,9 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return SeekDBConfig if self == DB.AliyunMilvus: - from .aliyun_milvus.config import AliyunMilvusConfig + from .milvus.config import MilvusConfig - return AliyunMilvusConfig + return MilvusConfig msg = f"Unknown DB: {self.name}" raise ValueError(msg) diff --git a/vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py b/vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py deleted file mode 100644 index 844171daf..000000000 --- a/vectordb_bench/backend/clients/aliyun_milvus/aliyun_milvus.py +++ /dev/null @@ -1,193 +0,0 @@ -"""AliyunMilvus client implementation. - -Extends the base Milvus client with Aliyun-specific behaviors: -- Expose ``load_reqs_size`` to tune insert batch size (avoid gRPC payload limits). -- Expose ``load_after_compaction`` to choose collection load timing: - - ``False`` (default): ``load_collection`` is called right after collection creation; - ``optimize()`` ends with ``refresh_load``. - - ``True``: ``load_collection`` is deferred until after compaction & index wait - in ``optimize()``. -- Apply ``knowhere.*`` collection properties (enable_thp / enable_prefetch) before load. -""" - -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -from pymilvus import DataType, MilvusClient - -from ..milvus.milvus import MILVUS_FORCE_MERGE_TARGET_SIZE_MB, MILVUS_LOAD_REQS_SIZE, Milvus - -if TYPE_CHECKING: - from .config import MilvusIndexConfig - -log = logging.getLogger(__name__) - - -class AliyunMilvus(Milvus): - # Insert concurrency is handled by the framework (ConcurrentInsertRunner). - # This client does not manage its own multiprocessing pool. - thread_safe: bool = True - - def __init__( - self, - dim: int, - db_config: dict, - db_case_config: MilvusIndexConfig, - collection_name: str = "VDBBench", - drop_old: bool = False, - name: str = "AliyunMilvus", - with_scalar_labels: bool = False, - **kwargs, - ): - """Initialize wrapper around the aliyun milvus vector database. - - The vector index is created right after collection creation (same as the - base ``Milvus`` client). Whether ``load_collection`` is called here or - deferred to ``optimize()`` is controlled by ``load_after_compaction``. - """ - self.name = name - self.db_config = db_config - self.case_config = db_case_config - self.collection_name = collection_name - self.load_reqs_size = int(self.db_config.get("load_reqs_size", MILVUS_LOAD_REQS_SIZE)) - self.batch_size = max(1, int(self.load_reqs_size / (dim * 4))) - self.with_scalar_labels = with_scalar_labels - self.load_after_compaction = bool(self.db_config.get("load_after_compaction", False)) - - self._primary_field = "pk" - self._scalar_id_field = "id" - self._scalar_label_field = "label" - self._scalar_payload_label_field = self._scalar_label_field - self._multitenant_partition_key_field = self._scalar_label_field - self.multitenant_tenant_labels: list[str] = kwargs.get("multitenant_tenant_labels", []) - if self.multitenant_tenant_labels: - self._multitenant_partition_key_field = "labels" - if self.with_scalar_labels: - self._scalar_payload_label_field = "scalar_label" - self._vector_field = "vector" - self._vector_index_name = "vector_idx" - self._scalar_id_index_name = "id_sort_idx" - self._scalar_labels_index_name = "labels_idx" - - client = MilvusClient( - uri=self.db_config.get("uri"), - user=self.db_config.get("user"), - password=self.db_config.get("password"), - token=self.db_config.get("token", ""), - timeout=30, - ) - - if drop_old and client.has_collection(self.collection_name): - log.info(f"{self.name} client drop_old collection: {self.collection_name}") - client.drop_collection(self.collection_name) - - if not client.has_collection(self.collection_name): - schema = MilvusClient.create_schema() - schema.add_field(self._primary_field, DataType.INT64, is_primary=True) - schema.add_field(self._scalar_id_field, DataType.INT64) - schema.add_field(self._vector_field, DataType.FLOAT_VECTOR, dim=dim) - - if self.multitenant_tenant_labels: - schema.add_field( - self._multitenant_partition_key_field, - DataType.VARCHAR, - max_length=256, - is_partition_key=True, - ) - - if self.with_scalar_labels: - is_partition_key = db_case_config.use_partition_key - log.info(f"with_scalar_labels, add a new varchar field, as partition_key: {is_partition_key}") - if not self.multitenant_tenant_labels or ( - self._scalar_payload_label_field != self._multitenant_partition_key_field - ): - schema.add_field( - self._scalar_payload_label_field, - DataType.VARCHAR, - max_length=256, - is_partition_key=is_partition_key and not self.multitenant_tenant_labels, - ) - - log.info(f"{self.name} create collection: {self.collection_name}") - - index_params = self._build_index_params() - client.create_collection( - collection_name=self.collection_name, - schema=schema, - num_shards=self.db_config.get("num_shards", 1), - consistency_level="Session", - ) - client.create_index(self.collection_name, index_params) - if not self.load_after_compaction: - self._apply_load_with_properties(client) - - client.close() - - def need_normalize_cosine(self) -> bool: - """Wheather this database need to normalize dataset to support COSINE""" - log.info("cosine dataset need normalize.") - return True - - def _apply_load_with_properties(self, client: MilvusClient) -> None: - """Apply Aliyun-specific knowhere.* properties and load the collection. - - The properties (currently ``knowhere.enable_thp`` / ``knowhere.enable_prefetch``) - have to be set via ``alter_collection_properties`` before ``load_collection`` - so the segment loader picks them up. - """ - load_param = getattr(self.case_config, "load_param", None) - properties = load_param() if callable(load_param) else {} - if properties: - client.alter_collection_properties( - collection_name=self.collection_name, - properties=properties, - ) - - collection_properties = client.describe_collection(self.collection_name).get("properties") - log.info(f"set collection properties to: {collection_properties}") - - client.load_collection( - self.collection_name, - replica_number=self.db_config.get("replica_number", 1), - ) - - def _optimize(self): - """Aliyun-specific optimize. - - Overrides the parent ``_optimize`` to control collection-load timing: - - ``load_after_compaction=True``: collection is loaded here (after compaction - and index wait) via ``_apply_load_with_properties``; - - ``load_after_compaction=False`` (default): collection was already loaded in - ``__init__``, so we end with ``refresh_load`` like the upstream Milvus. - """ - log.info(f"{self.name} optimizing before search") - try: - self.client.flush(self.collection_name) - - try: - self._wait_for_segments_sorted() - compaction_id = self.client.compact( - self.collection_name, - target_size=MILVUS_FORCE_MERGE_TARGET_SIZE_MB, - ) - if compaction_id > 0: - self._wait_for_compaction(compaction_id) - log.info(f"{self.name} force merge compaction completed.") - except Exception as e: - log.warning(f"{self.name} compact or list segments error: {e}") - if getattr(getattr(e, "code", None), "name", None) == "PERMISSION_DENIED": - log.warning("Skip compact due to list segments or compact permission denied.") - else: - raise e from None - - # wait for index no matter what - self._wait_for_index() - if self.load_after_compaction: - self._apply_load_with_properties(self.client) - else: - self.client.refresh_load(self.collection_name) - except Exception as e: - log.warning(f"{self.name} optimize error: {e}") - raise e from None diff --git a/vectordb_bench/backend/clients/aliyun_milvus/cli.py b/vectordb_bench/backend/clients/aliyun_milvus/cli.py index 751ffc4e2..314e6a10f 100644 --- a/vectordb_bench/backend/clients/aliyun_milvus/cli.py +++ b/vectordb_bench/backend/clients/aliyun_milvus/cli.py @@ -5,165 +5,28 @@ from vectordb_bench.backend.clients import DB from vectordb_bench.cli.cli import ( - CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, run, ) -from ..milvus.cli import MilvusTypedDict +from ..milvus.cli import MilvusDISKANNTypedDict DBTYPE = DB.AliyunMilvus -class AliyunMilvusTypedDict(MilvusTypedDict): - load_reqs_size: Annotated[ - int, - click.option( - "--load-reqs-size", - type=int, - help="Max request payload size (bytes) used to compute insert batch size", - required=False, - default=int(1.5 * 1024 * 1024), - show_default=True, - ), - ] - load_after_compaction: Annotated[ - bool, - click.option( - "--load-after-compaction/--no-load-after-compaction", - "load_after_compaction", - help=( - "If True, defer load_collection until after compaction & index build in optimize(); " - "if False (default), load_collection right after collection creation and call " - "refresh_load at the end of optimize()." - ), - default=False, - show_default=True, - ), - ] +class AliyunMilvusDISKANNTypedDict(MilvusDISKANNTypedDict): + """Same as Milvus DISKANN, plus three opt-in search-time params. + Each of the three options defaults to "not specified" (omit -> not sent). + """ -class AliyunMilvusDISKANNTypedDict(CommonTypedDict, AliyunMilvusTypedDict): - max_degree: Annotated[ - int, - click.option( - "--max-degree", - type=int, - help="R (max degree)", - required=False, - default=56, - show_default=True, - ), - ] - search_list: Annotated[ - int, - click.option( - "--search-list", - type=int, - help="L (search list size) used during DISKANN index search", - required=False, - default=200, - show_default=True, - ), - ] - build_search_list: Annotated[ - int, - click.option( - "--build-search-list", - type=int, - help="L (search list size) used during DISKANN index build", - required=False, - default=200, - show_default=True, - ), - ] - legacy: Annotated[ - bool, - click.option( - "--legacy/--no-legacy", - "legacy", - help="Use legacy Aliyun DISKANN behavior", - default=False, - show_default=True, - ), - ] - store_strategy: Annotated[ - str, - click.option( - "--store-strategy", - type=click.Choice(["MEMORY", "DISK"], case_sensitive=False), - help="Aliyun store strategy", - required=False, - default="MEMORY", - show_default=True, - ), - ] - quant_type: Annotated[ - str, - click.option( - "--quant-type", - type=click.Choice(["RABITQ", "PQ"], case_sensitive=False), - help="Aliyun quant type", - required=False, - default="RABITQ", - show_default=True, - ), - ] - num_threads: Annotated[ - int, - click.option( - "--num-threads", - type=int, - help="Degree of parallelism", - required=False, - default=4, - show_default=True, - ), - ] - distance_strategy: Annotated[ - str, - click.option( - "--distance-strategy", - type=click.Choice( - ["FULL", "SINGLE QUANT", "QUANT THEN FULL", "QUANT THEN MORE BITS"], - case_sensitive=False, - ), - help="Aliyun distance strategy", - required=False, - default="QUANT THEN MORE BITS", - show_default=True, - ), - ] - enable_prefetch: Annotated[ - bool, - click.option( - "--enable-prefetch/--disable-prefetch", - "enable_prefetch", - help="Enable Aliyun DISKANN prefetch during search", - default=False, - show_default=True, - ), - ] - enable_thp: Annotated[ - bool, - click.option( - "--enable-thp/--disable-thp", - "enable_thp", - help="Enable Aliyun DISKANN transparent huge pages on collection load", - default=False, - show_default=True, - ), - ] rerank_topk_multiplier: Annotated[ int | None, click.option( "--rerank-topk-multiplier", type=int, - help=( - "Search param: multiplier of topk for the rerank candidate budget " - "(0 disables rerank read). Omit to NOT send this search param." - ), + help="Search param: topk multiplier for rerank budget (0 disables rerank read). Omit to not send it.", required=False, default=None, ), @@ -173,7 +36,7 @@ class AliyunMilvusDISKANNTypedDict(CommonTypedDict, AliyunMilvusTypedDict): click.option( "--early-termination-threshold", type=int, - help="Search param: early termination threshold (0 disables). Omit to NOT send this search param.", + help="Search param: early termination threshold (0 disables). Omit to not send it.", required=False, default=None, ), @@ -183,7 +46,7 @@ class AliyunMilvusDISKANNTypedDict(CommonTypedDict, AliyunMilvusTypedDict): click.option( "--cross-segment-rerank/--no-cross-segment-rerank", "cross_segment_rerank", - help="Search param: enable cross-segment rerank in the reduce layer. Omit to NOT send this search param.", + help="Search param: enable cross-segment rerank. Omit to not send it.", default=None, ), ] @@ -192,31 +55,21 @@ class AliyunMilvusDISKANNTypedDict(CommonTypedDict, AliyunMilvusTypedDict): @cli.command(name="aliyunmilvusdiskann") @click_parameter_decorators_from_typed_dict(AliyunMilvusDISKANNTypedDict) def AliyunMilvusDISKANN(**parameters: Unpack[AliyunMilvusDISKANNTypedDict]): - from .config import AliyunMilvusConfig, AliyunMilvusDISKANNConfig + from ..milvus.config import MilvusConfig + from .config import AliyunMilvusDISKANNConfig run( db=DBTYPE, - db_config=AliyunMilvusConfig( + db_config=MilvusConfig( db_label=parameters["db_label"], uri=SecretStr(parameters["uri"]), user=parameters["user_name"], password=SecretStr(parameters["password"]) if parameters["password"] else None, num_shards=int(parameters["num_shards"]), replica_number=int(parameters["replica_number"]), - load_reqs_size=int(parameters["load_reqs_size"]), - load_after_compaction=bool(parameters["load_after_compaction"]), ), db_case_config=AliyunMilvusDISKANNConfig( search_list=parameters["search_list"], - build_search_list=parameters["build_search_list"], - max_degree=parameters["max_degree"], - legacy=parameters["legacy"], - store_strategy=parameters["store_strategy"], - quant_type=parameters["quant_type"], - num_threads=parameters["num_threads"], - distance_strategy=parameters["distance_strategy"], - enable_prefetch=bool(parameters["enable_prefetch"]), - enable_thp=bool(parameters["enable_thp"]), # Pass through as-is; None means "not specified" -> omitted from search params. rerank_topk_multiplier=parameters["rerank_topk_multiplier"], early_termination_threshold=parameters["early_termination_threshold"], diff --git a/vectordb_bench/backend/clients/aliyun_milvus/config.py b/vectordb_bench/backend/clients/aliyun_milvus/config.py index 42560a7b2..6f53323c6 100644 --- a/vectordb_bench/backend/clients/aliyun_milvus/config.py +++ b/vectordb_bench/backend/clients/aliyun_milvus/config.py @@ -1,65 +1,22 @@ -from typing import ClassVar +from pydantic import field_validator -from pydantic import SecretStr, field_validator - -from ..api import DBCaseConfig, DBConfig, IndexType, MetricType -from ..milvus.config import DISKANNConfig, MilvusIndexConfig - - -class AliyunMilvusConfig(DBConfig): - _extra_empty_skip: ClassVar[frozenset[str]] = frozenset({"user", "password"}) - - uri: SecretStr = SecretStr("http://localhost:19530") - user: str | None = None - password: SecretStr | None = None - num_shards: int = 1 - replica_number: int = 1 - # Tuning knobs for load performance - load_reqs_size: int = int(1.5 * 1024 * 1024) - # Controls when load_collection runs: - # - False (default): load immediately after collection creation, then refresh_load in optimize(). - # - True: defer load until after compaction and index build in optimize(). - load_after_compaction: bool = False - - def to_dict(self) -> dict: - return { - "uri": self.uri.get_secret_value(), - "user": self.user if self.user else None, - "password": self.password.get_secret_value() if self.password else None, - "num_shards": self.num_shards, - "replica_number": self.replica_number, - "load_reqs_size": self.load_reqs_size, - "load_after_compaction": self.load_after_compaction, - } +from ..api import DBCaseConfig, IndexType +from ..milvus.config import DISKANNConfig class AliyunMilvusDISKANNConfig(DISKANNConfig): - """AliyunMilvus DISKANN index config. + """Milvus DISKANN plus three opt-in Aliyun search-time params. - Inherits from Milvus ``DISKANNConfig`` and adds Aliyun-specific fields. - Overrides ``search_param`` to inject the three search-time knobs - (``rerank_topk_multiplier`` / ``early_termination_threshold`` / - ``cross_segment_rerank``) into the search ``params`` while keeping the - original contract. - """ + Identical to the upstream Milvus DISKANN for index build/load. The only + difference is three extra search params that are injected into the per-query + search params **only when explicitly specified**. ``None`` means "not + specified" (the param is omitted and the server keeps its own default). - # ---- index build knobs ---- - max_degree: int = 48 - legacy: bool = False - store_strategy: str = "MEMORY" - quant_type: str = "RABITQ" - num_threads: int = 4 - distance_strategy: str = "QUANT THEN MORE BITS" - enable_prefetch: bool = False - enable_thp: bool = False - build_search_list: int = 200 + Note ``0`` is a meaningful value (e.g. ``rerank_topk_multiplier=0`` disables + rerank reads), so "unset" must be ``None``, not ``0``. The web UI passes a + negative number / ``"DEFAULT"`` sentinel which is normalized to ``None``. + """ - # ---- search-time knobs ---- - # Optional on purpose: a value of None means "do not send this search param" - # (the server keeps its own default). They are only injected into search - # params when explicitly set. Note ``0`` is a meaningful value (e.g. - # rerank_topk_multiplier=0 disables rerank reads), so "unset" must be None, - # not 0. UI passes a negative sentinel for "unset", normalized to None below. rerank_topk_multiplier: int | None = None early_termination_threshold: int | None = None cross_segment_rerank: bool | None = None @@ -84,64 +41,19 @@ def _normalize_optional_bool(cls, v: object) -> bool | None: return v.strip().lower() in ("true", "1", "yes", "on") return bool(v) - def parse_metric(self) -> str: - if not self.metric_type: - return "" - - if self.metric_type == MetricType.COSINE: - return MetricType.L2.value - return self.metric_type.value - - def index_param(self) -> dict: - extra_params: dict = { - "max_degree": self.max_degree, - "legacy": self.legacy, - "store_strategy": self.store_strategy, - "quant_type": self.quant_type, - "num_threads": self.num_threads, - "distance_strategy": self.distance_strategy, - "search_list_size": self.build_search_list, - } - return { - "metric_type": self.parse_metric(), - "index_type": self.index.value, - "params": extra_params, - } - def search_param(self) -> dict: - # Only inject a search param when it was explicitly specified. - params: dict = {} - if self.search_list is not None: - params["search_list"] = self.search_list + # Reuse the base DISKANN search params (metric_type + search_list) and + # only add the three knobs that were explicitly specified. + sp = super().search_param() if self.rerank_topk_multiplier is not None: - params["rerank_topk_multiplier"] = self.rerank_topk_multiplier + sp["params"]["rerank_topk_multiplier"] = self.rerank_topk_multiplier if self.early_termination_threshold is not None: - params["early_termination_threshold"] = self.early_termination_threshold + sp["params"]["early_termination_threshold"] = self.early_termination_threshold if self.cross_segment_rerank is not None: - params["cross_segment_rerank"] = self.cross_segment_rerank - return { - "metric_type": self.parse_metric(), - "params": params, - } + sp["params"]["cross_segment_rerank"] = self.cross_segment_rerank + return sp - def load_param(self) -> dict: - return { - "knowhere.enable_thp": str(self.enable_thp).lower(), - "knowhere.enable_prefetch": str(self.enable_prefetch).lower(), - } - -# Only DISKANN is supported by AliyunMilvus today. Other index types are not -# exposed here so callers fail fast (case_config_cls returns None) instead of -# silently using the upstream Milvus index implementation. _aliyun_milvus_case_config: dict[IndexType, type[DBCaseConfig]] = { IndexType.DISKANN: AliyunMilvusDISKANNConfig, } - - -__all__ = [ - "MilvusIndexConfig", - "AliyunMilvusConfig", - "AliyunMilvusDISKANNConfig", - "_aliyun_milvus_case_config", -] diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index c0959f6e4..d5d3ad0db 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -2169,11 +2169,11 @@ class CaseConfigInput(BaseModel): ] # ---- AliyunMilvus ---- -# Reuse Milvus UI configs for all index types except DISKANN, where -# Aliyun-specific build params and the three search-time knobs -# (rerank_topk_multiplier / early_termination_threshold / cross_segment_rerank) -# are shown in addition to Milvus's SearchList. - +# Same as Milvus DISKANN, plus three opt-in search-time knobs. They default to +# "unset" -> the param is NOT sent to the server (it keeps its own default). +# For Number inputs a negative value (-1) is the "unset" sentinel; for the bool, +# "DEFAULT" is the "unset" sentinel. The backend config normalizes these to None +# (note: 0 is a real, meaningful value). CaseConfigParamInput_IndexType_AliyunMilvus = CaseConfigInput( label=CaseConfigParamType.IndexType, inputType=InputType.Option, @@ -2183,86 +2183,6 @@ class CaseConfigInput(BaseModel): }, ) -CaseConfigParamInput_Aliyun_max_degree = CaseConfigInput( - label=CaseConfigParamType.max_degree, - inputType=InputType.Number, - inputConfig={"min": 1, "max": 65536, "value": 56}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_legacy = CaseConfigInput( - label=CaseConfigParamType.legacy, - inputType=InputType.Bool, - displayLabel="Legacy", - inputHelp="Use legacy Aliyun DISKANN behavior", - inputConfig={"value": False}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_store_strategy = CaseConfigInput( - label=CaseConfigParamType.store_strategy, - inputType=InputType.Option, - inputConfig={"options": ["MEMORY", "DISK"], "value": "MEMORY"}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_quant_type = CaseConfigInput( - label=CaseConfigParamType.quant_type, - inputType=InputType.Option, - inputConfig={"options": ["RABITQ", "PQ"], "value": "RABITQ"}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_num_threads = CaseConfigInput( - label=CaseConfigParamType.num_threads, - inputType=InputType.Number, - inputConfig={"min": 1, "max": 1024, "value": 4}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_distance_strategy = CaseConfigInput( - label=CaseConfigParamType.distance_strategy, - inputType=InputType.Option, - inputConfig={ - "options": ["FULL", "SINGLE QUANT", "QUANT THEN FULL", "QUANT THEN MORE BITS"], - "value": "QUANT THEN MORE BITS", - }, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_enable_prefetch = CaseConfigInput( - label=CaseConfigParamType.enable_prefetch, - inputType=InputType.Bool, - displayLabel="Enable Prefetch", - inputHelp="Enable Aliyun DISKANN prefetch during search", - inputConfig={"value": False}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_enable_thp = CaseConfigInput( - label=CaseConfigParamType.enable_thp, - inputType=InputType.Bool, - displayLabel="Enable THP", - inputHelp="Enable Aliyun DISKANN transparent huge pages on collection load", - inputConfig={"value": False}, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -CaseConfigParamInput_Aliyun_BuildSearchList = CaseConfigInput( - label=CaseConfigParamType.BuildSearchList, - inputType=InputType.Number, - inputConfig={ - "min": 1, - "max": MAX_STREAMLIT_INT, - "value": 200, - }, - isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, -) - -# The three search-time knobs default to "unset" -> the param is NOT sent to the -# server (it keeps its own default). For Number inputs a negative value (-1) is -# the "unset" sentinel; for the bool, "DEFAULT" is the "unset" sentinel. The -# backend config normalizes these to None (note: 0 is a real, meaningful value). CaseConfigParamInput_Aliyun_rerank_topk_multiplier = CaseConfigInput( label=CaseConfigParamType.rerank_topk_multiplier, inputType=InputType.Number, @@ -2290,17 +2210,6 @@ class CaseConfigInput(BaseModel): isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.DISKANN.value, ) -_AliyunMilvusDiskannBuildParams = [ - CaseConfigParamInput_Aliyun_max_degree, - CaseConfigParamInput_Aliyun_legacy, - CaseConfigParamInput_Aliyun_store_strategy, - CaseConfigParamInput_Aliyun_quant_type, - CaseConfigParamInput_Aliyun_num_threads, - CaseConfigParamInput_Aliyun_distance_strategy, - CaseConfigParamInput_Aliyun_enable_prefetch, - CaseConfigParamInput_Aliyun_enable_thp, - CaseConfigParamInput_Aliyun_BuildSearchList, -] _AliyunMilvusDiskannSearchParams = [ CaseConfigParamInput_Aliyun_rerank_topk_multiplier, CaseConfigParamInput_Aliyun_early_termination_threshold, @@ -2310,12 +2219,10 @@ class CaseConfigInput(BaseModel): AliyunMilvusLoadConfig = [ CaseConfigParamInput_IndexType_AliyunMilvus, *MilvusLoadConfig[1:], - *_AliyunMilvusDiskannBuildParams, ] AliyunMilvusPerformanceConfig = [ CaseConfigParamInput_IndexType_AliyunMilvus, *MilvusPerformanceConfig[1:], - *_AliyunMilvusDiskannBuildParams, *_AliyunMilvusDiskannSearchParams, ] diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 70f638dad..f1c7d8cce 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -51,7 +51,6 @@ class CaseConfigParamType(Enum): ef_construction = "ef_construction" EF = "ef" SearchList = "search_list" - BuildSearchList = "build_search_list" ef_search = "ef_search" Nlist = "nlist" Nprobe = "nprobe" @@ -176,16 +175,7 @@ class CaseConfigParamType(Enum): exbits = "exbits" number_of_regions = "number_of_regions" - # AliyunMilvus DISKANN parameters - max_degree = "max_degree" - legacy = "legacy" - store_strategy = "store_strategy" - quant_type = "quant_type" - num_threads = "num_threads" - distance_strategy = "distance_strategy" - enable_prefetch = "enable_prefetch" - enable_thp = "enable_thp" - # AliyunMilvus DISKANN search-time parameters + # AliyunMilvus DISKANN opt-in search-time parameters rerank_topk_multiplier = "rerank_topk_multiplier" early_termination_threshold = "early_termination_threshold" cross_segment_rerank = "cross_segment_rerank" From 8da77115d56c97308a9ffe86192947b3b8ae936c Mon Sep 17 00:00:00 2001 From: frankleaf Date: Mon, 8 Jun 2026 14:09:41 +0800 Subject: [PATCH 3/6] refactor(aliyun_milvus): match Milvus DISKANN exactly + drop empty __init__ - Remove the empty aliyun_milvus/__init__.py: like milvus/ (and most client dirs) it is an implicit namespace package; the marker was unnecessary. - Wrap the case config with _with_partition_key, mirroring the MilvusDISKANN command, so AliyunMilvus is literally "Milvus DISKANN + 3 opt-in search params" with identical behavior otherwise. Co-authored-by: Cursor --- .../backend/clients/aliyun_milvus/__init__.py | 0 .../backend/clients/aliyun_milvus/cli.py | 17 ++++++++++------- 2 files changed, 10 insertions(+), 7 deletions(-) delete mode 100644 vectordb_bench/backend/clients/aliyun_milvus/__init__.py diff --git a/vectordb_bench/backend/clients/aliyun_milvus/__init__.py b/vectordb_bench/backend/clients/aliyun_milvus/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/vectordb_bench/backend/clients/aliyun_milvus/cli.py b/vectordb_bench/backend/clients/aliyun_milvus/cli.py index 314e6a10f..01f63527d 100644 --- a/vectordb_bench/backend/clients/aliyun_milvus/cli.py +++ b/vectordb_bench/backend/clients/aliyun_milvus/cli.py @@ -10,7 +10,7 @@ run, ) -from ..milvus.cli import MilvusDISKANNTypedDict +from ..milvus.cli import MilvusDISKANNTypedDict, _with_partition_key DBTYPE = DB.AliyunMilvus @@ -68,12 +68,15 @@ def AliyunMilvusDISKANN(**parameters: Unpack[AliyunMilvusDISKANNTypedDict]): num_shards=int(parameters["num_shards"]), replica_number=int(parameters["replica_number"]), ), - db_case_config=AliyunMilvusDISKANNConfig( - search_list=parameters["search_list"], - # Pass through as-is; None means "not specified" -> omitted from search params. - rerank_topk_multiplier=parameters["rerank_topk_multiplier"], - early_termination_threshold=parameters["early_termination_threshold"], - cross_segment_rerank=parameters["cross_segment_rerank"], + db_case_config=_with_partition_key( + AliyunMilvusDISKANNConfig( + search_list=parameters["search_list"], + # Pass through as-is; None means "not specified" -> omitted from search params. + rerank_topk_multiplier=parameters["rerank_topk_multiplier"], + early_termination_threshold=parameters["early_termination_threshold"], + cross_segment_rerank=parameters["cross_segment_rerank"], + ), + parameters, ), **parameters, ) From 2e91c03df5ee24664695239eb14e7baaa0795767 Mon Sep 17 00:00:00 2001 From: frankleaf Date: Mon, 8 Jun 2026 14:24:47 +0800 Subject: [PATCH 4/6] test(aliyun_milvus): add unit tests + list AliyunMilvus in README - Add tests/test_aliyun_milvus.py covering: reuse of the upstream Milvus client/config, opt-in search params (omitted by default, injected only when set), 0 treated as a real value, UI sentinel (-1 / "DEFAULT") -> unset, frontend exposure in Performance (not Load), and DISKANN-only case mapping. - README: add aliyunmilvus to the default (pymilvus) client row. Co-authored-by: Cursor --- README.md | 2 +- tests/test_aliyun_milvus.py | 135 ++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 tests/test_aliyun_milvus.py diff --git a/README.md b/README.md index 3cdceddc0..4a07dc2d1 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ All the database client supported | Optional database client | install command | |--------------------------|---------------------------------------------| -| pymilvus, zilliz_cloud (*default*) | `pip install vectordb-bench` | +| pymilvus, zilliz_cloud, aliyunmilvus (*default*) | `pip install vectordb-bench` | | qdrant | `pip install vectordb-bench[qdrant]` | | pinecone | `pip install vectordb-bench[pinecone]` | | weaviate | `pip install vectordb-bench[weaviate]` | diff --git a/tests/test_aliyun_milvus.py b/tests/test_aliyun_milvus.py new file mode 100644 index 000000000..20c0a5136 --- /dev/null +++ b/tests/test_aliyun_milvus.py @@ -0,0 +1,135 @@ +import pytest +from click.testing import CliRunner + +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.aliyun_milvus import cli as aliyun_milvus_cli +from vectordb_bench.backend.clients.api import IndexType +from vectordb_bench.backend.clients.milvus.config import MilvusConfig +from vectordb_bench.backend.clients.milvus.milvus import Milvus +from vectordb_bench.frontend.config.dbCaseConfigs import ( + AliyunMilvusLoadConfig, + AliyunMilvusPerformanceConfig, +) +from vectordb_bench.models import CaseConfigParamType + + +def _diskann_config_cls(): + return DB.AliyunMilvus.case_config_cls(index_type=IndexType.DISKANN) + + +def test_aliyun_milvus_reuses_upstream_milvus_client_and_config(): + """AliyunMilvus is minimal: it reuses the upstream Milvus client and config.""" + assert DB.AliyunMilvus.init_cls is Milvus + assert DB.AliyunMilvus.config_cls is MilvusConfig + + +def test_aliyun_milvus_search_params_opt_in_by_default(): + """With none of the three knobs set, only search_list is sent (same as Milvus DISKANN).""" + case_config = _diskann_config_cls()(search_list=200) + + assert case_config.search_param()["params"] == {"search_list": 200} + + +def test_aliyun_milvus_search_params_injected_when_set(): + case_config = _diskann_config_cls()( + search_list=200, + rerank_topk_multiplier=0, + early_termination_threshold=0, + cross_segment_rerank=True, + ) + + assert case_config.search_param()["params"] == { + "search_list": 200, + "rerank_topk_multiplier": 0, + "early_termination_threshold": 0, + "cross_segment_rerank": True, + } + + +def test_aliyun_milvus_zero_is_a_meaningful_value_not_unset(): + """0 is a real value (e.g. disables rerank reads) and must be sent.""" + case_config = _diskann_config_cls()(search_list=200, rerank_topk_multiplier=0) + + assert case_config.search_param()["params"]["rerank_topk_multiplier"] == 0 + + +def test_aliyun_milvus_ui_sentinels_normalize_to_unset(): + """UI 'unset' sentinels (-1 for numbers, 'DEFAULT' for the bool) -> None -> omitted.""" + case_config = _diskann_config_cls()( + search_list=200, + rerank_topk_multiplier=-1, + early_termination_threshold=-1, + cross_segment_rerank="DEFAULT", + ) + + assert case_config.search_param()["params"] == {"search_list": 200} + + +def test_aliyun_milvus_frontend_exposes_search_params_in_performance_only(): + """The three knobs are search-time only: shown in Performance, not in Load.""" + load_labels = [config.label for config in AliyunMilvusLoadConfig] + performance_labels = [config.label for config in AliyunMilvusPerformanceConfig] + + for label in ( + CaseConfigParamType.rerank_topk_multiplier, + CaseConfigParamType.early_termination_threshold, + CaseConfigParamType.cross_segment_rerank, + ): + assert label in performance_labels + assert label not in load_labels + + +def test_aliyun_milvus_load_config_has_no_search_list(): + """search_list is a query-time param, must not appear in load config.""" + load_labels = [config.label for config in AliyunMilvusLoadConfig] + assert CaseConfigParamType.SearchList not in load_labels + + +def test_aliyun_milvus_only_diskann_supported(): + """Non-DISKANN index types should not silently fall back to upstream Milvus.""" + assert DB.AliyunMilvus.case_config_cls(index_type=IndexType.DISKANN) is not None + assert DB.AliyunMilvus.case_config_cls(index_type=IndexType.HNSW) is None + assert DB.AliyunMilvus.case_config_cls(index_type=IndexType.IVFFlat) is None + + +def test_aliyun_milvus_cli_omitting_knobs_sends_only_search_list(monkeypatch: pytest.MonkeyPatch): + captured = {} + + monkeypatch.setattr(aliyun_milvus_cli, "run", lambda **kwargs: captured.update(kwargs)) + + runner = CliRunner() + result = runner.invoke( + aliyun_milvus_cli.AliyunMilvusDISKANN, + ["--uri", "http://localhost:19530", "--search-list", "200"], + ) + + assert result.exit_code == 0, result.output + case_config = captured["db_case_config"] + assert case_config.rerank_topk_multiplier is None + assert case_config.early_termination_threshold is None + assert case_config.cross_segment_rerank is None + assert case_config.search_param()["params"] == {"search_list": 200} + + +def test_aliyun_milvus_cli_passes_knobs_when_specified(monkeypatch: pytest.MonkeyPatch): + captured = {} + + monkeypatch.setattr(aliyun_milvus_cli, "run", lambda **kwargs: captured.update(kwargs)) + + runner = CliRunner() + result = runner.invoke( + aliyun_milvus_cli.AliyunMilvusDISKANN, + [ + "--uri", "http://localhost:19530", + "--search-list", "200", + "--rerank-topk-multiplier", "0", + "--early-termination-threshold", "0", + "--cross-segment-rerank", + ], + ) + + assert result.exit_code == 0, result.output + case_config = captured["db_case_config"] + assert case_config.rerank_topk_multiplier == 0 + assert case_config.early_termination_threshold == 0 + assert case_config.cross_segment_rerank is True From 3aed50daf7553c59da6fc73bc6d181fbdd1b6442 Mon Sep 17 00:00:00 2001 From: frankleaf Date: Mon, 8 Jun 2026 14:37:52 +0800 Subject: [PATCH 5/6] test(aliyun_milvus): cover bool-knob normalization (100% module cover) Add parametrized cases for cross_segment_rerank string/int coercion so config.py reaches 100% coverage. Tidy README default-client wording. Co-authored-by: Cursor --- README.md | 2 +- tests/test_aliyun_milvus.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 4a07dc2d1..8a93f0a8a 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ All the database client supported | Optional database client | install command | |--------------------------|---------------------------------------------| -| pymilvus, zilliz_cloud, aliyunmilvus (*default*) | `pip install vectordb-bench` | +| pymilvus, zilliz_cloud (*default*) , aliyunmilvus | `pip install vectordb-bench` | | qdrant | `pip install vectordb-bench[qdrant]` | | pinecone | `pip install vectordb-bench[pinecone]` | | weaviate | `pip install vectordb-bench[weaviate]` | diff --git a/tests/test_aliyun_milvus.py b/tests/test_aliyun_milvus.py index 20c0a5136..3b8922193 100644 --- a/tests/test_aliyun_milvus.py +++ b/tests/test_aliyun_milvus.py @@ -53,6 +53,29 @@ def test_aliyun_milvus_zero_is_a_meaningful_value_not_unset(): assert case_config.search_param()["params"]["rerank_topk_multiplier"] == 0 +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("true", True), + ("True", True), + ("1", True), + ("yes", True), + ("on", True), + ("false", False), + ("0", False), + ("anything", False), + (1, True), + (0, False), + ], +) +def test_aliyun_milvus_cross_segment_rerank_normalizes_non_bool_inputs(raw, expected): + """CLI/UI may pass strings or ints for the bool knob; coerce them consistently.""" + case_config = _diskann_config_cls()(search_list=200, cross_segment_rerank=raw) + + assert case_config.cross_segment_rerank is expected + assert case_config.search_param()["params"]["cross_segment_rerank"] is expected + + def test_aliyun_milvus_ui_sentinels_normalize_to_unset(): """UI 'unset' sentinels (-1 for numbers, 'DEFAULT' for the bool) -> None -> omitted.""" case_config = _diskann_config_cls()( From bfd8861e142f367590e84a8f9f5f5eb240f1697a Mon Sep 17 00:00:00 2001 From: frankleaf Date: Mon, 8 Jun 2026 15:11:32 +0800 Subject: [PATCH 6/6] test(aliyun_milvus): cover --no-cross-segment-rerank sends False Explicit --no-cross-segment-rerank must send cross_segment_rerank=False (disable), not omit it. Omission is the only "unset" path. Co-authored-by: Cursor --- tests/test_aliyun_milvus.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_aliyun_milvus.py b/tests/test_aliyun_milvus.py index 3b8922193..b3980c65b 100644 --- a/tests/test_aliyun_milvus.py +++ b/tests/test_aliyun_milvus.py @@ -156,3 +156,28 @@ def test_aliyun_milvus_cli_passes_knobs_when_specified(monkeypatch: pytest.Monke assert case_config.rerank_topk_multiplier == 0 assert case_config.early_termination_threshold == 0 assert case_config.cross_segment_rerank is True + + +def test_aliyun_milvus_cli_no_cross_segment_rerank_sends_false(monkeypatch: pytest.MonkeyPatch): + """--no-cross-segment-rerank explicitly disables it: send False, not omit.""" + captured = {} + + monkeypatch.setattr(aliyun_milvus_cli, "run", lambda **kwargs: captured.update(kwargs)) + + runner = CliRunner() + result = runner.invoke( + aliyun_milvus_cli.AliyunMilvusDISKANN, + [ + "--uri", "http://localhost:19530", + "--search-list", "200", + "--no-cross-segment-rerank", + ], + ) + + assert result.exit_code == 0, result.output + case_config = captured["db_case_config"] + assert case_config.cross_segment_rerank is False + assert case_config.search_param()["params"] == { + "search_list": 200, + "cross_segment_rerank": False, + }