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/tests/test_milvus.py b/tests/test_milvus.py index 83fabf957..55128a1b2 100644 --- a/tests/test_milvus.py +++ b/tests/test_milvus.py @@ -128,6 +128,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, @@ -135,12 +152,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 @@ -180,6 +203,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 ): @@ -353,6 +394,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 6db46e9c7..d01628f10 100644 --- a/tests/test_milvus_zilliz_cli.py +++ b/tests/test_milvus_zilliz_cli.py @@ -57,6 +57,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 + + @pytest.mark.parametrize( ("level_args", "expected_level"), [(["--level", "2"], 2), ([], None)], @@ -134,6 +201,69 @@ 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 = {} + + 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 + + def test_milvus_autoindex_cli_nq(monkeypatch: MonkeyPatch) -> None: captured = {} diff --git a/vectordb_bench/backend/clients/milvus/cli.py b/vectordb_bench/backend/clients/milvus/cli.py index ec11015c8..4df345ae7 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 @@ -17,6 +17,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: @@ -24,8 +37,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: @@ -99,6 +118,37 @@ 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, + callback=_validate_positive_int_or_none, + help=( + "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." + ), + ), + ] class MilvusAutoIndexTypedDict(CommonTypedDict, MilvusTypedDict, AutoIndexLevelTypedDict): ... @@ -112,7 +162,7 @@ def MilvusAutoIndex(**parameters: Unpack[MilvusAutoIndexTypedDict]): run( db=DBTYPE, db_config=_build_milvus_config(parameters), - db_case_config=_with_partition_key(AutoIndexConfig(level=parameters["level"]), parameters), + db_case_config=_apply_milvus_case_defaults(AutoIndexConfig(level=parameters["level"]), parameters), **parameters, ) @@ -128,7 +178,7 @@ def MilvusFlat(**parameters: Unpack[MilvusFlatTypedDict]): 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, ) @@ -144,7 +194,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"], @@ -205,7 +255,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"], @@ -245,7 +295,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"], @@ -282,7 +332,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"], @@ -309,7 +359,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"], @@ -328,7 +378,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"], @@ -386,7 +436,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"], @@ -413,7 +463,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"], ), @@ -439,7 +489,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"], @@ -471,7 +521,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 @@ -557,7 +607,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"], @@ -580,7 +630,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"], @@ -617,7 +667,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"], @@ -651,7 +701,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"], @@ -688,7 +738,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"], @@ -740,6 +790,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 caa9a82e2..1bc3b569c 100644 --- a/vectordb_bench/backend/clients/milvus/config.py +++ b/vectordb_bench/backend/clients/milvus/config.py @@ -1,6 +1,6 @@ from typing import Annotated, ClassVar -from pydantic import BaseModel, Field, SecretStr +from pydantic import BaseModel, Field, 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: @@ -542,6 +552,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 dcb7e9a5d..4296e6740 100644 --- a/vectordb_bench/backend/clients/milvus/milvus.py +++ b/vectordb_bench/backend/clients/milvus/milvus.py @@ -341,10 +341,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: @@ -361,7 +365,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): @@ -377,8 +384,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 0a127b249..9f2002987 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 ( AutoIndexLevelTypedDict, CommonTypedDict, @@ -83,6 +84,37 @@ 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, + callback=_validate_positive_int_or_none, + help=( + "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." + ), + ), + ] class ZillizAutoIndexTypedDict(ZillizTypedDict, AutoIndexLevelTypedDict): ... @@ -108,6 +140,8 @@ def ZillizAutoIndex(**parameters: Unpack[ZillizAutoIndexTypedDict]): level=parameters["level"] if parameters["level"] is not None 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/cli/cli.py b/vectordb_bench/cli/cli.py index 1d405be79..511a537f5 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", ) diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index 678b53288..f35f38109 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 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={ + "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 21eea9e95..c9cf95ed2 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"