From 8406e34735800bd84bfa50f8140601e0200a2f15 Mon Sep 17 00:00:00 2001 From: jamesgao-jpg Date: Tue, 15 Sep 2026 15:45:00 +0800 Subject: [PATCH 1/4] feat: add CLI control for Milvus force-merge size and toggle Expose the Milvus force-merge compaction during optimize as tunable CLI options, addressing zilliztech/VectorDBBench#825 (unconditional force merge is costly and non-reproducible at 100M scale): - --force-merge-enabled/--no-force-merge-enabled: skip the force-merge stage for a sooner ready-to-search state at the cost of segments not merged to their fullest potential. - --force-merge-target-size-mb: bound the merged segment size (e.g. 1024 for bounded, reproducible segments); defaults to the current unbounded single-segment behavior for backward compatibility. Applied to all Milvus index CLI commands and to Zilliz Cloud AutoIndex, and registered in the web UI case configs. GPU index types keep skipping force merge. Disabling force merge keeps flush, normal compaction, index-wait, and refresh-load. Signed-off-by: jamesgao-jpg --- tests/test_milvus.py | 44 ++++++++- tests/test_milvus_zilliz_cli.py | 94 +++++++++++++++++++ vectordb_bench/backend/clients/milvus/cli.py | 75 +++++++++++---- .../backend/clients/milvus/config.py | 22 ++++- .../backend/clients/milvus/milvus.py | 18 +++- .../backend/clients/zilliz_cloud/cli.py | 31 ++++++ .../frontend/config/dbCaseConfigs.py | 37 ++++++++ vectordb_bench/models.py | 2 + 8 files changed, 298 insertions(+), 25 deletions(-) diff --git a/tests/test_milvus.py b/tests/test_milvus.py index d08faf0a2..15b7922cf 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -85,6 +85,23 @@ def test_milvus_rejects_existing_collection_with_incompatible_properties(monkeyp _create_milvus_with_collection_properties(monkeypatch, collection_exists=True) +def test_milvus_force_merge_config_defaults_and_validation(): + from vectordb_bench.backend.clients.milvus.config import AutoIndexConfig, MilvusFtsConfig + + assert AutoIndexConfig().force_merge_enabled is True + assert AutoIndexConfig().force_merge_target_size_mb is None + assert AutoIndexConfig(force_merge_target_size_mb=1024).force_merge_target_size_mb == 1024 + assert MilvusFtsConfig().force_merge_enabled is True + assert MilvusFtsConfig().force_merge_target_size_mb is None + + with pytest.raises(ValueError, match="positive integer"): + AutoIndexConfig(force_merge_target_size_mb=0) + with pytest.raises(ValueError, match="positive integer"): + AutoIndexConfig(force_merge_target_size_mb=-5) + with pytest.raises(ValueError, match="positive integer"): + MilvusFtsConfig(force_merge_target_size_mb=0) + + class TestMilvusOptimize: def _milvus( self, @@ -92,12 +109,18 @@ def _milvus( compact_side_effect: Exception | None = None, is_fts: bool = False, is_gpu_index: bool = False, + force_merge_enabled: bool = True, + force_merge_target_size_mb: int | None = None, ): milvus = Milvus.__new__(Milvus) milvus.name = "Milvus" milvus.collection_name = "test_collection" milvus._is_fts = is_fts - milvus.case_config = SimpleNamespace(is_gpu_index=is_gpu_index) + milvus.case_config = SimpleNamespace( + is_gpu_index=is_gpu_index, + force_merge_enabled=force_merge_enabled, + force_merge_target_size_mb=force_merge_target_size_mb, + ) milvus.client = MagicMock() milvus.client.compact.side_effect = compact_side_effect milvus.client.compact.return_value = 42 @@ -137,6 +160,24 @@ def test_optimize_flushes_and_runs_normal_compaction_before_force_merge(self): assert milvus._wait_for_compaction.call_args_list == [call(41), call(42)] milvus.client.refresh_load.assert_called_once_with("test_collection") + def test_optimize_uses_configured_force_merge_target_size(self): + milvus = self._milvus(force_merge_target_size_mb=512) + + milvus._optimize() + + milvus.client.compact.assert_any_call("test_collection", target_size=512) + milvus.client.refresh_load.assert_called_once_with("test_collection") + + def test_optimize_skips_force_merge_when_disabled_but_keeps_normal_compaction(self): + milvus = self._milvus(force_merge_enabled=False) + + milvus._optimize() + + milvus.client.flush.assert_called_once_with("test_collection") + assert milvus.client.compact.call_args_list == [call("test_collection")] + assert milvus._wait_for_compaction.call_args_list == [call(42)] + milvus.client.refresh_load.assert_called_once_with("test_collection") + def test_optimize_retries_when_compacting_segments_are_missing_from_force_merge_plan( self, monkeypatch: pytest.MonkeyPatch ): @@ -310,6 +351,7 @@ def get_compaction_plans(self, _compaction_id: int): milvus.name = "Milvus" milvus.collection_name = "test_collection" milvus._main_index_name = "vector_idx" + milvus.case_config = SimpleNamespace() milvus.client = FakeMilvusClient() monkeypatch.setattr("vectordb_bench.backend.clients.milvus.milvus.time.sleep", lambda _seconds: None) diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index 4177c8898..027601247 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -55,6 +55,73 @@ def fake_run(**kwargs): assert captured["db_case_config"].use_partition_key is True +def test_milvus_autoindex_cli_defaults_force_merge_to_current_behavior(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", "--dry-run"], + ) + + assert result.exit_code == 0, result.output + assert captured["db_case_config"].force_merge_enabled is True + assert captured["db_case_config"].force_merge_target_size_mb is None + + +def test_milvus_autoindex_cli_accepts_force_merge_flags(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", + "--force-merge-target-size-mb", + "512", + "--no-force-merge-enabled", + "--dry-run", + ], + ) + + assert result.exit_code == 0, result.output + assert captured["db_case_config"].force_merge_enabled is False + assert captured["db_case_config"].force_merge_target_size_mb == 512 + + +def test_milvus_fts_cli_accepts_force_merge_flags(monkeypatch: MonkeyPatch) -> None: + captured = {} + + def fake_run(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(milvus_cli, "run", fake_run) + + result = CliRunner().invoke( + milvus_cli.MilvusFTS, + [ + "--uri", + "http://localhost:19530", + "--force-merge-target-size-mb", + "256", + "--dry-run", + ], + ) + + assert result.exit_code == 0, result.output + assert captured["db_case_config"].force_merge_enabled is True + assert captured["db_case_config"].force_merge_target_size_mb == 256 + + def test_zilliz_autoindex_cli_enables_partition_key_for_multitenant_case( monkeypatch: MonkeyPatch, ) -> None: @@ -80,3 +147,30 @@ def fake_run(**kwargs): assert result.exit_code == 0, result.output assert captured["db_case_config"].use_partition_key is True + + +def test_zilliz_autoindex_cli_accepts_force_merge_flags(monkeypatch: MonkeyPatch) -> None: + captured = {} + + def fake_run(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(zilliz_cli, "run", fake_run) + + result = CliRunner().invoke( + zilliz_cli.ZillizAutoIndex, + [ + "--uri", + "https://example.api.gcp-us-west1.zillizcloud.com", + "--token", + "secret", + "--force-merge-target-size-mb", + "1024", + "--no-force-merge-enabled", + "--dry-run", + ], + ) + + assert result.exit_code == 0, result.output + assert captured["db_case_config"].force_merge_enabled is False + assert captured["db_case_config"].force_merge_target_size_mb == 1024 diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 738248ac3..862bcd7af 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -23,8 +23,14 @@ def _use_partition_key(parameters: dict) -> bool: return parameters.get("case_type") == "CloudMultiTenantSearchCase" -def _with_partition_key(db_case_config: BaseModel, parameters: dict) -> BaseModel: - return db_case_config.model_copy(update={"use_partition_key": _use_partition_key(parameters)}) +def _apply_milvus_case_defaults(db_case_config: BaseModel, parameters: dict) -> BaseModel: + return db_case_config.model_copy( + update={ + "use_partition_key": _use_partition_key(parameters), + "force_merge_enabled": parameters["force_merge_enabled"], + "force_merge_target_size_mb": parameters["force_merge_target_size_mb"], + } + ) def _build_milvus_config(parameters: dict) -> BaseModel: @@ -98,6 +104,35 @@ class MilvusTypedDict(TypedDict): ), ), ] + force_merge_enabled: Annotated[ + bool, + click.option( + "--force-merge-enabled/--no-force-merge-enabled", + type=bool, + default=True, + show_default=True, + help=( + "Whether to force-merge compaction during optimize. Disable for a sooner " + "ready-to-search state at the cost of segments not merged to their fullest " + "potential (search latency may vary)." + ), + ), + ] + force_merge_target_size_mb: Annotated[ + int | None, + click.option( + "--force-merge-target-size-mb", + type=int, + required=False, + default=None, + show_default=True, + help=( + "Target merged segment size in MB for the force-merge compaction during " + "optimize. Defaults to the current unbounded single-segment behavior; set " + "e.g. 1024 for bounded, reproducible segments. Must be a positive integer." + ), + ), + ] class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict): ... @@ -111,7 +146,7 @@ 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=_apply_milvus_case_defaults(AutoIndexConfig(), parameters), **parameters, ) @@ -124,7 +159,7 @@ def MilvusFlat(**parameters: Unpack[MilvusAutoIndexTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key(FLATConfig(), parameters), + db_case_config=_apply_milvus_case_defaults(FLATConfig(), parameters), **parameters, ) @@ -140,7 +175,7 @@ def MilvusHNSW(**parameters: Unpack[MilvusHNSWTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( HNSWConfig( M=parameters["m"], efConstruction=parameters["ef_construction"], @@ -201,7 +236,7 @@ def MilvusHNSWPQ(**parameters: Unpack[MilvusHNSWPQTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( HNSWPQConfig( M=parameters["m"], efConstruction=parameters["ef_construction"], @@ -241,7 +276,7 @@ def MilvusHNSWPRQ(**parameters: Unpack[MilvusHNSWPRQTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( HNSWPRQConfig( M=parameters["m"], efConstruction=parameters["ef_construction"], @@ -278,7 +313,7 @@ def MilvusHNSWSQ(**parameters: Unpack[MilvusHNSWSQTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( HNSWSQConfig( M=parameters["m"], efConstruction=parameters["ef_construction"], @@ -305,7 +340,7 @@ def MilvusIVFFlat(**parameters: Unpack[MilvusIVFFlatTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( IVFFlatConfig( nlist=parameters["nlist"], nprobe=parameters["nprobe"], @@ -324,7 +359,7 @@ def MilvusIVFSQ8(**parameters: Unpack[MilvusIVFFlatTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( IVFSQ8Config( nlist=parameters["nlist"], nprobe=parameters["nprobe"], @@ -382,7 +417,7 @@ def MilvusIVFRabitQ(**parameters: Unpack[MilvusIVFRABITQTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( IVFRABITQConfig( nlist=parameters["nlist"], nprobe=parameters["nprobe"], @@ -409,7 +444,7 @@ def MilvusDISKANN(**parameters: Unpack[MilvusDISKANNTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( DISKANNConfig( search_list=parameters["search_list"], ), @@ -435,7 +470,7 @@ def MilvusGPUIVFFlat(**parameters: Unpack[MilvusGPUIVFTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( GPUIVFFlatConfig( nlist=parameters["nlist"], nprobe=parameters["nprobe"], @@ -467,7 +502,7 @@ def MilvusGPUBruteForce(**parameters: Unpack[MilvusGPUBruteForceTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( GPUBruteForceConfig( metric_type=parameters["metric_type"], limit=parameters["limit"], # top-k for search @@ -553,7 +588,7 @@ def MilvusSVSVamana(**parameters: Unpack[MilvusSVSVamanaTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( SVSVamanaConfig( svs_graph_max_degree=parameters["svs_graph_max_degree"], svs_construction_window_size=parameters["svs_construction_window_size"], @@ -576,7 +611,7 @@ def MilvusSVSVamanaLVQ(**parameters: Unpack[MilvusSVSVamanaTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( SVSVamanaLVQConfig( svs_graph_max_degree=parameters["svs_graph_max_degree"], svs_construction_window_size=parameters["svs_construction_window_size"], @@ -613,7 +648,7 @@ def MilvusSVSVamanaLeanVec(**parameters: Unpack[MilvusSVSVamanaLeanVecTypedDict] run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( SVSVamanaLeanVecConfig( svs_graph_max_degree=parameters["svs_graph_max_degree"], svs_construction_window_size=parameters["svs_construction_window_size"], @@ -647,7 +682,7 @@ def MilvusGPUIVFPQ(**parameters: Unpack[MilvusGPUIVFPQTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( GPUIVFPQConfig( nlist=parameters["nlist"], nprobe=parameters["nprobe"], @@ -684,7 +719,7 @@ def MilvusGPUCAGRA(**parameters: Unpack[MilvusGPUCAGRATypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key( + db_case_config=_apply_milvus_case_defaults( GPUCAGRAConfig( intermediate_graph_degree=parameters["intermediate_graph_degree"], graph_degree=parameters["graph_degree"], @@ -736,6 +771,8 @@ def MilvusFTS(**parameters: Unpack[MilvusFTSTypedDict]): db_config=_build_milvus_config(parameters), db_case_config=MilvusFtsConfig( drop_ratio_search=parameters.get("drop_ratio_search"), + force_merge_enabled=parameters["force_merge_enabled"], + force_merge_target_size_mb=parameters["force_merge_target_size_mb"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/milvus/config.py b/vectordb_bench/backend/clients/milvus/config.py index 99d5ac8cf..6c3b34a16 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, SecretStr, field_validator from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType @@ -32,6 +32,16 @@ class MilvusIndexConfig(BaseModel): index: IndexType metric_type: MetricType | None = None use_partition_key: bool = False # for label-filter + force_merge_enabled: bool = True + force_merge_target_size_mb: int | None = None + + @field_validator("force_merge_target_size_mb") + @classmethod + def _validate_force_merge_target_size(cls, value: int | None) -> int | None: + if value is not None and value <= 0: + message = "force_merge_target_size_mb must be a positive integer or None" + raise ValueError(message) + return value @property def is_gpu_index(self) -> bool: @@ -537,6 +547,16 @@ class MilvusFtsConfig(BaseModel, DBCaseConfig): analyzer_max_token_length: int | None = None analyzer_stop_words: str | None = None drop_ratio_search: float | None = None + force_merge_enabled: bool = True + force_merge_target_size_mb: int | None = None + + @field_validator("force_merge_target_size_mb") + @classmethod + def _validate_force_merge_target_size(cls, value: int | None) -> int | None: + if value is not None and value <= 0: + message = "force_merge_target_size_mb must be a positive integer or None" + raise ValueError(message) + return value def analyzer_param(self) -> dict: analyzer_params = {} diff --git a/vectordb_bench/backend/clients/milvus/milvus.py b/vectordb_bench/backend/clients/milvus/milvus.py index 7c639a977..985e5c5e3 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -340,10 +340,14 @@ def _force_merge(self, max_attempts: int = MILVUS_FORCE_MERGE_MAX_ATTEMPTS): message = "force merge max_attempts must be greater than zero" raise ValueError(message) + target_size_mb = getattr(self.case_config, "force_merge_target_size_mb", None) + if target_size_mb is None: + target_size_mb = MILVUS_FORCE_MERGE_TARGET_SIZE_MB + for attempt in range(1, max_attempts + 1): self._wait_for_compaction_ready() expected_source_ids = self._force_merge_source_segment_ids() - compaction_id = self.client.compact(self.collection_name, target_size=MILVUS_FORCE_MERGE_TARGET_SIZE_MB) + compaction_id = self.client.compact(self.collection_name, target_size=target_size_mb) if compaction_id <= 0: failure_detail = f"generated no plan for snapshot {sorted(expected_source_ids)}" else: @@ -360,7 +364,10 @@ def _force_merge(self, max_attempts: int = MILVUS_FORCE_MERGE_MAX_ATTEMPTS): message = f"{self.name} force merge failed after {max_attempts} attempts; {failure_detail}" raise RuntimeError(message) - log.info(f"{self.name} force merge attempt {attempt}/{max_attempts} {failure_detail}; retrying...") + log.info( + f"{self.name} force merge attempt {attempt}/{max_attempts} " + f"(target_size={target_size_mb}MB) {failure_detail}; retrying..." + ) time.sleep(MILVUS_FORCE_MERGE_RETRY_INTERVAL_SECONDS) def _optimize(self): @@ -376,8 +383,11 @@ def _optimize(self): compaction_id = self.client.compact(self.collection_name) if compaction_id > 0: self._wait_for_compaction(compaction_id) - self._force_merge() - log.info(f"{self.name} force merge compaction completed.") + if getattr(self.case_config, "force_merge_enabled", True): + self._force_merge() + log.info(f"{self.name} force merge compaction completed.") + else: + log.info(f"{self.name} force merge compaction disabled; skipping.") 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": diff --git a/vectordb_bench/backend/clients/zilliz_cloud/cli.py b/vectordb_bench/backend/clients/zilliz_cloud/cli.py index 030a00661..ad845e0cb 100644 --- a/vectordb_bench/backend/clients/zilliz_cloud/cli.py +++ b/vectordb_bench/backend/clients/zilliz_cloud/cli.py @@ -86,6 +86,35 @@ class ZillizTypedDict(CommonTypedDict): ), ), ] + force_merge_enabled: Annotated[ + bool, + click.option( + "--force-merge-enabled/--no-force-merge-enabled", + type=bool, + default=True, + show_default=True, + help=( + "Whether to force-merge compaction during optimize. Disable for a sooner " + "ready-to-search state at the cost of segments not merged to their fullest " + "potential (search latency may vary)." + ), + ), + ] + force_merge_target_size_mb: Annotated[ + int | None, + click.option( + "--force-merge-target-size-mb", + type=int, + required=False, + default=None, + show_default=True, + help=( + "Target merged segment size in MB for the force-merge compaction during " + "optimize. Defaults to the current unbounded single-segment behavior; set " + "e.g. 1024 for bounded, reproducible segments. Must be a positive integer." + ), + ), + ] @cli.command() @@ -108,6 +137,8 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizTypedDict]): level=int(parameters["level"]) if parameters["level"] else 1, num_shards=parameters["num_shards"], use_partition_key=_use_partition_key(parameters), + force_merge_enabled=parameters["force_merge_enabled"], + force_merge_target_size_mb=parameters["force_merge_target_size_mb"], ), **parameters, ) diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index 678b53288..e33ac373a 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -1999,6 +1999,33 @@ class CaseConfigInput(BaseModel): inputConfig={"options": [False, True]}, ) +CaseConfigParamInput_Milvus_force_merge_enabled = CaseConfigInput( + label=CaseConfigParamType.force_merge_enabled, + displayLabel="Force Merge Enabled", + inputHelp=( + "Force-merge compaction during optimize. Disable for a sooner ready-to-search " + "state at the cost of segments not merged to their fullest potential." + ), + inputType=InputType.Bool, + inputConfig={"value": True}, +) + +CaseConfigParamInput_Milvus_force_merge_target_size_mb = CaseConfigInput( + label=CaseConfigParamType.force_merge_target_size_mb, + displayLabel="Force Merge Target Size (MB)", + inputHelp=( + "Target merged segment size in MB for the force-merge compaction during optimize. " + "Leave empty to keep the current unbounded single-segment behavior; e.g. 1024 for " + "bounded, reproducible segments." + ), + inputType=InputType.Number, + inputConfig={ + "value": None, + "min": 1, + "max": MAX_STREAMLIT_INT, + }, +) + CaseConfigParamInput_MongoDBQuantizationType = CaseConfigInput( label=CaseConfigParamType.mongodb_quantization_type, @@ -2423,6 +2450,8 @@ class CaseConfigInput(BaseModel): CaseConfigParamInput_RefineType, CaseConfigParamInput_NRQ, CaseConfigParamInput_Milvus_use_partition_key, + CaseConfigParamInput_Milvus_force_merge_enabled, + CaseConfigParamInput_Milvus_force_merge_target_size_mb, ] MilvusPerformanceConfig = [ CaseConfigParamInput_IndexType, @@ -2453,6 +2482,8 @@ class CaseConfigInput(BaseModel): CaseConfigParamInput_RefineType, CaseConfigParamInput_RefineK, CaseConfigParamInput_Milvus_use_partition_key, + CaseConfigParamInput_Milvus_force_merge_enabled, + CaseConfigParamInput_Milvus_force_merge_target_size_mb, ] @@ -2466,11 +2497,15 @@ class CaseConfigInput(BaseModel): CaseConfigParamInput_FTS_analyzer_max_token_length, CaseConfigParamInput_FTS_analyzer_stop_words, CaseConfigParamInput_FTS_drop_ratio_search, + CaseConfigParamInput_Milvus_force_merge_enabled, + CaseConfigParamInput_Milvus_force_merge_target_size_mb, ] ZillizCloudFtsConfig = [ CaseConfigParamInput_IndexType_FTS, CaseConfigParamInput_ZillizLevel, + CaseConfigParamInput_Milvus_force_merge_enabled, + CaseConfigParamInput_Milvus_force_merge_target_size_mb, ] ElasticCloudFtsConfig = [] @@ -2578,6 +2613,8 @@ class CaseConfigInput(BaseModel): ZillizCloudPerformanceConfig = [ CaseConfigParamInput_ZillizLevel, + CaseConfigParamInput_Milvus_force_merge_enabled, + CaseConfigParamInput_Milvus_force_merge_target_size_mb, ] PgVectorScaleLoadingConfig = [ diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 4e471363a..ff0ef0082 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -140,6 +140,8 @@ class CaseConfigParamType(Enum): mongodb_num_candidates_ratio = "num_candidates_ratio" mongodb_exact = "exact" use_partition_key = "use_partition_key" + force_merge_enabled = "force_merge_enabled" + force_merge_target_size_mb = "force_merge_target_size_mb" refresh_interval = "refresh_interval" use_rescore = "use_rescore" oversample_ratio = "oversample_ratio" From 84beccd9bf4fe5096b1dd5e1d4d38eb597e24446 Mon Sep 17 00:00:00 2001 From: jamesgao-jpg Date: Tue, 15 Sep 2026 16:34:37 +0800 Subject: [PATCH 2/4] fix: preserve force-merge target size through FTS CLI routing copy_fts_compatible_db_case_fields() now also carries force_merge_target_size_mb when a non-FTS vector config is routed to an FTS config class, so --force-merge-target-size-mb survives a --case-type FTSBm25Performance invocation (force_merge_enabled was already whitelisted). Signed-off-by: jamesgao-jpg --- tests/test_fts_cli_user_control.py | 19 +++++++++++++++++++ vectordb_bench/cli/cli.py | 1 + 2 files changed, 20 insertions(+) diff --git a/tests/test_fts_cli_user_control.py b/tests/test_fts_cli_user_control.py index 1496b3bc8..466bba6ae 100644 --- a/tests/test_fts_cli_user_control.py +++ b/tests/test_fts_cli_user_control.py @@ -80,6 +80,25 @@ def test_cli_applies_bm25_overrides_after_routing_vector_config_to_fts(): } +def test_cli_preserves_milvus_force_merge_options_when_routing_vector_config_to_fts(): + from vectordb_bench.backend.clients.milvus.config import MilvusIndexConfig + + selected = select_cli_db_case_config( + DB.Milvus, + MilvusIndexConfig( + index=IndexType.HNSW, + force_merge_enabled=False, + force_merge_target_size_mb=512, + ), + "FTSBm25Performance", + {}, + ) + + assert isinstance(selected, MilvusFtsConfig) + assert selected.force_merge_enabled is False + assert selected.force_merge_target_size_mb == 512 + + def test_cli_leaves_fts_bm25_defaults_when_options_are_omitted(): selected = select_cli_db_case_config( DB.Vespa, diff --git a/vectordb_bench/cli/cli.py b/vectordb_bench/cli/cli.py index c0f3798c9..7cff2e7e6 100644 --- a/vectordb_bench/cli/cli.py +++ b/vectordb_bench/cli/cli.py @@ -311,6 +311,7 @@ def copy_fts_compatible_db_case_fields(source: DBCaseConfig, target: DBCaseConfi "refresh_interval", "use_force_merge", "force_merge_enabled", + "force_merge_target_size_mb", "disable_backpressure", "level", ) From 72b121962f3bc93f88e0c26bf59f263142294d7e Mon Sep 17 00:00:00 2001 From: jamesgao-jpg Date: Wed, 16 Sep 2026 15:23:25 +0800 Subject: [PATCH 3/4] fix: validate and soften wording of --force-merge-target-size-mb Addresses two review findings on the force-merge CLI feature: - Regular Milvus/Zilliz CLI commands applied the flag through model_copy(update=...) in the per-command defaults helper, which skips pydantic validation, so non-positive values bypassed the field validator. Add a click callback (_validate_positive_int_or_none) on the option that rejects non-positive integers at the CLI boundary; direct config construction stays guarded by the pydantic validator. - Soften the public wording: the server-side behavior of target_size is not probed, so describe it as a requested cap whose effective behavior depends on the Milvus server instead of claiming bounded, reproducible segments. Adds CLI rejection tests for 0 and negative values on Milvus and Zilliz AutoIndex commands. Signed-off-by: jamesgao-jpg --- tests/test_milvus_zilliz_cli.py | 36 +++++++++++++++++++ vectordb_bench/backend/clients/milvus/cli.py | 23 +++++++++--- .../backend/clients/zilliz_cloud/cli.py | 9 +++-- .../frontend/config/dbCaseConfigs.py | 6 ++-- 4 files changed, 64 insertions(+), 10 deletions(-) diff --git a/tests/test_milvus_zilliz_cli.py b/tests/test_milvus_zilliz_cli.py index 027601247..4fef1f442 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -149,6 +149,42 @@ def fake_run(**kwargs): assert captured["db_case_config"].use_partition_key is True +def test_milvus_autoindex_cli_rejects_non_positive_force_merge_target_size() -> None: + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + ["--uri", "http://localhost:19530", "--force-merge-target-size-mb", "0", "--dry-run"], + ) + + assert result.exit_code != 0 + assert "positive integer" in result.output + + result = CliRunner().invoke( + milvus_cli.MilvusAutoIndex, + ["--uri", "http://localhost:19530", "--force-merge-target-size-mb", "-5", "--dry-run"], + ) + + assert result.exit_code != 0 + assert "positive integer" in result.output + + +def test_zilliz_autoindex_cli_rejects_non_positive_force_merge_target_size() -> None: + result = CliRunner().invoke( + zilliz_cli.ZillizAutoIndex, + [ + "--uri", + "https://example.api.gcp-us-west1.zillizcloud.com", + "--token", + "secret", + "--force-merge-target-size-mb", + "0", + "--dry-run", + ], + ) + + assert result.exit_code != 0 + assert "positive integer" in result.output + + def test_zilliz_autoindex_cli_accepts_force_merge_flags(monkeypatch: MonkeyPatch) -> None: captured = {} diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 862bcd7af..0914e6f28 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -1,4 +1,4 @@ -from typing import Annotated, TypedDict, Unpack +from typing import Annotated, Any, TypedDict, Unpack import click from pydantic import BaseModel, SecretStr @@ -16,6 +16,19 @@ DBTYPE = DB.Milvus +def _validate_positive_int_or_none(ctx: Any, param: Any, value: int | None) -> int | None: + """Click callback accepting ``None`` or a positive integer. + + Guards flags whose values are applied via ``model_copy(update=...)`` in + ``_apply_milvus_case_defaults`` — pydantic's ``model_copy`` does not run + validators, so the CLI boundary must reject invalid values itself. + """ + if value is not None and value <= 0: + message = f"must be a positive integer, got {value}" + raise click.BadParameter(message) + return value + + def _use_partition_key(parameters: dict) -> bool: explicit = parameters.get("use_partition_key") if explicit is not None: @@ -126,10 +139,12 @@ class MilvusTypedDict(TypedDict): required=False, default=None, show_default=True, + callback=_validate_positive_int_or_none, help=( - "Target merged segment size in MB for the force-merge compaction during " - "optimize. Defaults to the current unbounded single-segment behavior; set " - "e.g. 1024 for bounded, reproducible segments. Must be a positive integer." + "Target merged segment size in MB requested for the force-merge compaction " + "during optimize; the effective cap depends on the Milvus server. Defaults " + "to the current unbounded single-segment behavior; set e.g. 1024 for bounded " + "segments. Must be a positive integer." ), ), ] diff --git a/vectordb_bench/backend/clients/zilliz_cloud/cli.py b/vectordb_bench/backend/clients/zilliz_cloud/cli.py index ad845e0cb..59078ae8f 100644 --- a/vectordb_bench/backend/clients/zilliz_cloud/cli.py +++ b/vectordb_bench/backend/clients/zilliz_cloud/cli.py @@ -5,6 +5,7 @@ from pydantic import SecretStr from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.milvus.cli import _validate_positive_int_or_none from vectordb_bench.cli.cli import ( CommonTypedDict, cli, @@ -108,10 +109,12 @@ class ZillizTypedDict(CommonTypedDict): required=False, default=None, show_default=True, + callback=_validate_positive_int_or_none, help=( - "Target merged segment size in MB for the force-merge compaction during " - "optimize. Defaults to the current unbounded single-segment behavior; set " - "e.g. 1024 for bounded, reproducible segments. Must be a positive integer." + "Target merged segment size in MB requested for the force-merge compaction " + "during optimize; the effective cap depends on the Milvus server. Defaults " + "to the current unbounded single-segment behavior; set e.g. 1024 for bounded " + "segments. Must be a positive integer." ), ), ] diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index e33ac373a..f35f38109 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -2014,9 +2014,9 @@ class CaseConfigInput(BaseModel): label=CaseConfigParamType.force_merge_target_size_mb, displayLabel="Force Merge Target Size (MB)", inputHelp=( - "Target merged segment size in MB for the force-merge compaction during optimize. " - "Leave empty to keep the current unbounded single-segment behavior; e.g. 1024 for " - "bounded, reproducible segments." + "Target merged segment size in MB requested for the force-merge compaction during " + "optimize; the effective cap depends on the Milvus server. Leave empty to keep the " + "current unbounded single-segment behavior; e.g. 1024 for bounded segments." ), inputType=InputType.Number, inputConfig={ From e9522fd4b2bc81f9114bb4387ada4b5d27f78ea8 Mon Sep 17 00:00:00 2001 From: jamesgao-jpg Date: Mon, 21 Sep 2026 17:22:58 +0800 Subject: [PATCH 4/4] fix: silence ruff ARG001 for unused click callback args _validate_positive_int_or_none takes the required (ctx, param) callback arguments but only uses value; underscore-prefix them so the PR's lint check (ruff ARG001) passes. Signed-off-by: jamesgao-jpg --- vectordb_bench/backend/clients/milvus/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index 82a313192..4df345ae7 100644 --- a/vectordb_bench/backend/clients/milvus/cli.py +++ b/vectordb_bench/backend/clients/milvus/cli.py @@ -17,7 +17,7 @@ DBTYPE = DB.Milvus -def _validate_positive_int_or_none(ctx: Any, param: Any, value: int | None) -> int | None: +def _validate_positive_int_or_none(_ctx: Any, _param: Any, value: int | None) -> int | None: """Click callback accepting ``None`` or a positive integer. Guards flags whose values are applied via ``model_copy(update=...)`` in