diff --git a/internal/datacoord/compaction_view_forcemerge_test.go b/internal/datacoord/compaction_view_forcemerge_test.go index fd886fa99e..a870352a79 100644 --- a/internal/datacoord/compaction_view_forcemerge_test.go +++ b/internal/datacoord/compaction_view_forcemerge_test.go @@ -400,6 +400,56 @@ func TestAdaptiveGroupSegments(t *testing.T) { }) } +func TestAdaptiveGroupSegmentsThresholdBoundary(t *testing.T) { + Params.Save(Params.DataCoordCfg.CompactionMaxFullSegmentThreshold.Key, "4") + t.Cleanup(func() { + Params.Reset(Params.DataCoordCfg.CompactionMaxFullSegmentThreshold.Key) + }) + + groupIDs := func(groups [][]*SegmentView) [][]int64 { + return lo.Map(groups, func(group []*SegmentView, _ int) []int64 { + return lo.Map(group, func(segment *SegmentView, _ int) int64 { + return segment.ID + }) + }) + } + + t.Run("at threshold selects max full grouping", func(t *testing.T) { + segments := []*SegmentView{ + {ID: 1, Size: 1}, + {ID: 2, Size: 1}, + {ID: 3, Size: 1}, + {ID: 4, Size: 8}, + } + + actual := groupIDs(adaptiveGroupSegments(segments, 10)) + exact := groupIDs(maxFullSegmentsGrouping(segments, 10)) + fallback := groupIDs(largerGroupingSegments(segments, 10)) + + assert.Equal(t, [][]int64{{1}, {2, 3, 4}}, actual) + assert.Equal(t, exact, actual) + assert.NotEqual(t, fallback, actual) + }) + + t.Run("above threshold selects larger grouping", func(t *testing.T) { + segments := []*SegmentView{ + {ID: 1, Size: 1}, + {ID: 2, Size: 1}, + {ID: 3, Size: 1}, + {ID: 4, Size: 1}, + {ID: 5, Size: 7}, + } + + actual := groupIDs(adaptiveGroupSegments(segments, 10)) + exact := groupIDs(maxFullSegmentsGrouping(segments, 10)) + fallback := groupIDs(largerGroupingSegments(segments, 10)) + + assert.Equal(t, [][]int64{{1, 2, 3, 4, 5}}, actual) + assert.Equal(t, fallback, actual) + assert.NotEqual(t, exact, actual) + }) +} + func TestLargerGroupingSegments(t *testing.T) { t.Run("empty segments", func(t *testing.T) { groups := largerGroupingSegments(nil, 5*1024*1024*1024) @@ -1043,3 +1093,75 @@ func TestCalculateTargetSizeCount_QueryNodeParallelism(t *testing.T) { assert.Equal(t, int64(1), targetCount, "targetCount should not be adjusted when totalSize/desiredCount < configMaxSize") }) } + +func TestCalculateTargetSizeCount_UserTargetAndMemoryClamp(t *testing.T) { + Params.Save(Params.DataCoordCfg.CompactionForceMergeQueryNodeMemoryFactor.Key, "4") + Params.Save(Params.DataCoordCfg.CompactionForceMergeDataNodeMemoryFactor.Key, "4") + t.Cleanup(func() { + Params.Reset(Params.DataCoordCfg.CompactionForceMergeQueryNodeMemoryFactor.Key) + Params.Reset(Params.DataCoordCfg.CompactionForceMergeDataNodeMemoryFactor.Key) + }) + + const ( + mb = float64(1024 * 1024) + gb = float64(1024 * 1024 * 1024) + ) + newView := func(expectedTargetSize float64) *ForceMergeSegmentView { + return &ForceMergeSegmentView{ + label: &CompactionGroupLabel{ + CollectionID: 1, + PartitionID: 1, + Channel: "ch1", + }, + segments: []*SegmentView{ + {ID: 1, Size: 2.5 * gb}, + {ID: 2, Size: 2.5 * gb}, + }, + triggerID: 1, + configMaxSize: 64 * mb, + expectedTargetSize: expectedTargetSize, + topology: &CollectionTopology{ + NumReplicas: 1, + NumShards: 1, + QueryNodeMemory: map[int64]uint64{ + 1: 8 * 1024 * 1024 * 1024, + 2: 16 * 1024 * 1024 * 1024, + }, + DataNodeMemory: map[int64]uint64{ + 1: 12 * 1024 * 1024 * 1024, + 2: 20 * 1024 * 1024 * 1024, + }, + }, + } + } + + t.Run("user target below safe size is preserved", func(t *testing.T) { + view := newView(1 * gb) + + targetSize, targetCount := view.calculateTargetSizeCount() + + assert.Equal(t, 1*gb, targetSize) + assert.Equal(t, int64(5), targetCount) + }) + + t.Run("user target above smallest node limit is clamped", func(t *testing.T) { + view := newView(4 * gb) + + targetSize, targetCount := view.calculateTargetSizeCount() + + // The smallest QueryNode is the limiting resource: 8 GiB / factor 4 = 2 GiB. + assert.Equal(t, 2*gb, targetSize) + assert.Equal(t, int64(3), targetCount) + }) + + t.Run("standalone co-location halves the shared memory limit", func(t *testing.T) { + view := newView(4 * gb) + view.topology.IsStandaloneMode = true + view.topology.QueryNodeMemory = map[int64]uint64{1: 8 * 1024 * 1024 * 1024} + + targetSize, targetCount := view.calculateTargetSizeCount() + + assert.Equal(t, 1*gb, targetSize) + assert.Equal(t, int64(5), targetCount) + }) +} diff --git a/internal/datacoord/server_test.go b/internal/datacoord/server_test.go index 1f1db10a16..7ee044d710 100644 --- a/internal/datacoord/server_test.go +++ b/internal/datacoord/server_test.go @@ -1839,7 +1839,7 @@ func TestGetCompactionState(t *testing.T) { func TestManualCompaction(t *testing.T) { paramtable.Get().Save(Params.DataCoordCfg.EnableCompaction.Key, "true") defer paramtable.Get().Reset(Params.DataCoordCfg.EnableCompaction.Key) - t.Run("test manual compaction successfully", func(t *testing.T) { + t.Run("target size zero routes to ordinary manual compaction", func(t *testing.T) { svr := &Server{allocator: allocator.NewMockAllocator(t)} svr.stateCode.Store(commonpb.StateCode_Healthy) svr.meta = &meta{collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()} @@ -1849,7 +1849,12 @@ func TestManualCompaction(t *testing.T) { }) mockTrigger := NewMockTrigger(t) svr.compactionTrigger = mockTrigger - mockTrigger.EXPECT().TriggerCompaction(mock.Anything, mock.Anything).Return(1, nil) + mockTrigger.EXPECT().TriggerCompaction(mock.Anything, mock.MatchedBy(func(signal *compactionSignal) bool { + return signal.collectionID == 1 && signal.isForce + })).Return(1, nil).Once() + + mockTriggerManager := NewMockTriggerManager(t) + svr.compactionTriggerManager = mockTriggerManager mockHandler := NewMockCompactionInspector(t) mockHandler.EXPECT().getCompactionTasksNumBySignalID(mock.Anything).Return(1) @@ -1857,9 +1862,11 @@ func TestManualCompaction(t *testing.T) { resp, err := svr.ManualCompaction(context.TODO(), &milvuspb.ManualCompactionRequest{ CollectionID: 1, Timetravel: 1, + TargetSize: 0, }) assert.NoError(t, err) assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) + mockTriggerManager.AssertNotCalled(t, "ManualTrigger", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything) }) t.Run("test manual l0 compaction successfully", func(t *testing.T) { diff --git a/tests/python_client/base/async_milvus_client_wrapper.py b/tests/python_client/base/async_milvus_client_wrapper.py index e3edd35bf9..ef82b7dab7 100644 --- a/tests/python_client/base/async_milvus_client_wrapper.py +++ b/tests/python_client/base/async_milvus_client_wrapper.py @@ -1,16 +1,13 @@ -import asyncio import sys -from typing import Optional, List, Union, Dict - -from pymilvus import ( - AsyncMilvusClient, - AnnSearchRequest, - RRFRanker, -) -from pymilvus.orm.types import CONSISTENCY_STRONG -from pymilvus.orm.collection import CollectionSchema from check.func_check import ResponseChecker +from pymilvus import ( + AnnSearchRequest, + AsyncMilvusClient, + RRFRanker, +) +from pymilvus.orm.collection import CollectionSchema +from pymilvus.orm.types import CONSISTENCY_STRONG from utils.api_request import api_request, logger_interceptor @@ -20,121 +17,143 @@ class AsyncMilvusClientWrapper: def __init__(self, active_trace=False): self.active_trace = active_trace - def init_async_client(self, uri: str = "http://localhost:19530", - user: str = "", - password: str = "", - db_name: str = "", - token: str = "", - timeout: Optional[float] = None, - active_trace=False, - check_task=None, check_items=None, - **kwargs): + def init_async_client( + self, + uri: str = "http://localhost:19530", + user: str = "", + password: str = "", + db_name: str = "", + token: str = "", + timeout: float | None = None, + active_trace=False, + check_task=None, + check_items=None, + **kwargs, + ): self.active_trace = active_trace """ In order to distinguish the same name of collection """ func_name = sys._getframe().f_code.co_name - res, is_succ = api_request([AsyncMilvusClient, uri, user, password, db_name, token, - timeout], **kwargs) + res, is_succ = api_request([AsyncMilvusClient, uri, user, password, db_name, token, timeout], **kwargs) self.async_milvus_client = res if is_succ else None check_result = ResponseChecker(res, func_name, check_task, check_items, is_succ, **kwargs).run() return res, check_result @logger_interceptor() - async def list_collections(self, timeout: Optional[float] = None, **kwargs): + async def list_collections(self, timeout: float | None = None, **kwargs): return await self.async_milvus_client.list_collections(timeout, **kwargs) - + @logger_interceptor() - async def has_collection(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def has_collection(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.has_collection(collection_name, timeout, **kwargs) - + @logger_interceptor() - async def has_partition(self, collection_name: str, partition_name: str, timeout: Optional[float] = None, **kwargs): + async def has_partition(self, collection_name: str, partition_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.has_partition(collection_name, partition_name, timeout, **kwargs) @logger_interceptor() - async def describe_collection(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def describe_collection(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.describe_collection(collection_name, timeout, **kwargs) @logger_interceptor() - async def list_partitions(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def list_partitions(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.list_partitions(collection_name, timeout, **kwargs) @logger_interceptor() - async def get_collection_stats(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def get_collection_stats(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.get_collection_stats(collection_name, timeout, **kwargs) @logger_interceptor() - async def flush(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def flush(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.flush(collection_name, timeout, **kwargs) @logger_interceptor() - async def get_load_state(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def get_load_state(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.get_load_state(collection_name, timeout, **kwargs) @logger_interceptor() - async def describe_index(self, collection_name: str, index_name: str, timeout: Optional[float] = None, **kwargs): + async def describe_index(self, collection_name: str, index_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.describe_index(collection_name, index_name, timeout, **kwargs) @logger_interceptor() - async def create_database(self, db_name: str, timeout: Optional[float] = None, **kwargs): + async def create_database(self, db_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.create_database(db_name, timeout, **kwargs) @logger_interceptor() - async def drop_database(self, db_name: str, timeout: Optional[float] = None, **kwargs): + async def drop_database(self, db_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.drop_database(db_name, timeout, **kwargs) @logger_interceptor() - async def list_databases(self, timeout: Optional[float] = None, **kwargs): + async def list_databases(self, timeout: float | None = None, **kwargs): return await self.async_milvus_client.list_databases(timeout, **kwargs) - + @logger_interceptor() async def list_indexes(self, collection_name: str, field_name: str = "", **kwargs): return await self.async_milvus_client.list_indexes(collection_name, field_name, **kwargs) @logger_interceptor() - async def create_collection(self, - collection_name: str, - dimension: Optional[int] = None, - primary_field_name: str = "id", # default is "id" - id_type: str = "int", # or "string", - vector_field_name: str = "vector", # default is "vector" - metric_type: str = "COSINE", - auto_id: bool = False, - timeout: Optional[float] = None, - schema: Optional[CollectionSchema] = None, - index_params=None, - **kwargs): + async def create_collection( + self, + collection_name: str, + dimension: int | None = None, + primary_field_name: str = "id", # default is "id" + id_type: str = "int", # or "string", + vector_field_name: str = "vector", # default is "vector" + metric_type: str = "COSINE", + auto_id: bool = False, + timeout: float | None = None, + schema: CollectionSchema | None = None, + index_params=None, + **kwargs, + ): kwargs["consistency_level"] = kwargs.get("consistency_level", CONSISTENCY_STRONG) - return await self.async_milvus_client.create_collection(collection_name, dimension, - primary_field_name, - id_type, vector_field_name, metric_type, - auto_id, - timeout, schema, index_params, **kwargs) + return await self.async_milvus_client.create_collection( + collection_name, + dimension, + primary_field_name, + id_type, + vector_field_name, + metric_type, + auto_id, + timeout, + schema, + index_params, + **kwargs, + ) @logger_interceptor() - async def drop_collection(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def drop_collection(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.drop_collection(collection_name, timeout, **kwargs) @logger_interceptor() - async def load_collection(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def load_collection(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.load_collection(collection_name, timeout, **kwargs) - + @logger_interceptor() async def release_collection(self, collection_name, timeout=None, **kwargs): return await self.async_milvus_client.release_collection(collection_name, timeout, **kwargs) @logger_interceptor() - async def truncate_collection(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def truncate_collection(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.truncate_collection(collection_name, timeout, **kwargs) - + @logger_interceptor() - async def list_persistent_segments(self, collection_name: str, timeout: Optional[float] = None, **kwargs): + async def list_persistent_segments(self, collection_name: str, timeout: float | None = None, **kwargs): return await self.async_milvus_client.list_persistent_segments(collection_name, timeout, **kwargs) @logger_interceptor() - async def create_index(self, collection_name: str, index_params, timeout: Optional[float] = None, - **kwargs): + async def optimize(self, collection_name: str, target_size=None, wait=True, timeout: float | None = None, **kwargs): + return await self.async_milvus_client.optimize( + collection_name, + target_size=target_size, + wait=wait, + timeout=timeout, + **kwargs, + ) + + @logger_interceptor() + async def create_index(self, collection_name: str, index_params, timeout: float | None = None, **kwargs): return await self.async_milvus_client.create_index(collection_name, index_params, timeout, **kwargs) @logger_interceptor() @@ -148,11 +167,11 @@ class AsyncMilvusClientWrapper: @logger_interceptor() async def create_partition(self, collection_name, partition_name, timeout=None, **kwargs): return await self.async_milvus_client.create_partition(collection_name, partition_name, timeout, **kwargs) - + @logger_interceptor() async def drop_partition(self, collection_name, partition_name, timeout=None, **kwargs): return await self.async_milvus_client.drop_partition(collection_name, partition_name, timeout, **kwargs) - + @logger_interceptor() async def load_partitions(self, collection_name, partition_names, timeout=None, **kwargs): return await self.async_milvus_client.load_partitions(collection_name, partition_names, timeout, **kwargs) @@ -160,97 +179,112 @@ class AsyncMilvusClientWrapper: @logger_interceptor() async def release_partitions(self, collection_name, partition_names, timeout=None, **kwargs): return await self.async_milvus_client.release_partitions(collection_name, partition_names, timeout, **kwargs) - + @logger_interceptor() - async def insert(self, - collection_name: str, - data: Union[Dict, List[Dict]], - timeout: Optional[float] = None, - partition_name: Optional[str] = "", - **kwargs): + async def insert( + self, + collection_name: str, + data: dict | list[dict], + timeout: float | None = None, + partition_name: str | None = "", + **kwargs, + ): return await self.async_milvus_client.insert(collection_name, data, timeout, partition_name, **kwargs) @logger_interceptor() - async def upsert(self, - collection_name: str, - data: Union[Dict, List[Dict]], - timeout: Optional[float] = None, - partition_name: Optional[str] = "", - **kwargs): + async def upsert( + self, + collection_name: str, + data: dict | list[dict], + timeout: float | None = None, + partition_name: str | None = "", + **kwargs, + ): return await self.async_milvus_client.upsert(collection_name, data, timeout, partition_name, **kwargs) @logger_interceptor() - async def search(self, - collection_name: str, - data: Union[List[list], list], - filter: str = "", - limit: int = 10, - output_fields: Optional[List[str]] = None, - search_params: Optional[dict] = None, - timeout: Optional[float] = None, - partition_names: Optional[List[str]] = None, - anns_field: Optional[str] = None, - **kwargs): - return await self.async_milvus_client.search(collection_name, data, - filter, - limit, output_fields, search_params, - timeout, - partition_names, anns_field, **kwargs) + async def search( + self, + collection_name: str, + data: list[list] | list, + filter: str = "", + limit: int = 10, + output_fields: list[str] | None = None, + search_params: dict | None = None, + timeout: float | None = None, + partition_names: list[str] | None = None, + anns_field: str | None = None, + **kwargs, + ): + return await self.async_milvus_client.search( + collection_name, + data, + filter, + limit, + output_fields, + search_params, + timeout, + partition_names, + anns_field, + **kwargs, + ) @logger_interceptor() - async def hybrid_search(self, - collection_name: str, - reqs: List[AnnSearchRequest], - ranker: RRFRanker, - limit: int = 10, - output_fields: Optional[List[str]] = None, - timeout: Optional[float] = None, - partition_names: Optional[List[str]] = None, - **kwargs): - return await self.async_milvus_client.hybrid_search(collection_name, reqs, - ranker, - limit, output_fields, - timeout, partition_names, **kwargs) + async def hybrid_search( + self, + collection_name: str, + reqs: list[AnnSearchRequest], + ranker: RRFRanker, + limit: int = 10, + output_fields: list[str] | None = None, + timeout: float | None = None, + partition_names: list[str] | None = None, + **kwargs, + ): + return await self.async_milvus_client.hybrid_search( + collection_name, reqs, ranker, limit, output_fields, timeout, partition_names, **kwargs + ) @logger_interceptor() - async def query(self, - collection_name: str, - filter: str = "", - output_fields: Optional[List[str]] = None, - timeout: Optional[float] = None, - ids: Optional[Union[List, str, int]] = None, - partition_names: Optional[List[str]] = None, - **kwargs): - return await self.async_milvus_client.query(collection_name, filter, - output_fields, timeout, - ids, partition_names, - **kwargs) + async def query( + self, + collection_name: str, + filter: str = "", + output_fields: list[str] | None = None, + timeout: float | None = None, + ids: list | str | int | None = None, + partition_names: list[str] | None = None, + **kwargs, + ): + return await self.async_milvus_client.query( + collection_name, filter, output_fields, timeout, ids, partition_names, **kwargs + ) @logger_interceptor() - async def get(self, - collection_name: str, - ids: Union[list, str, int], - output_fields: Optional[List[str]] = None, - timeout: Optional[float] = None, - partition_names: Optional[List[str]] = None, - **kwargs): - return await self.async_milvus_client.get(collection_name, ids, - output_fields, timeout, - partition_names, - **kwargs) + async def get( + self, + collection_name: str, + ids: list | str | int, + output_fields: list[str] | None = None, + timeout: float | None = None, + partition_names: list[str] | None = None, + **kwargs, + ): + return await self.async_milvus_client.get( + collection_name, ids, output_fields, timeout, partition_names, **kwargs + ) @logger_interceptor() - async def delete(self, - collection_name: str, - ids: Optional[Union[list, str, int]] = None, - timeout: Optional[float] = None, - filter: Optional[str] = None, - partition_name: Optional[str] = None, - **kwargs): - return await self.async_milvus_client.delete(collection_name, ids, - timeout, filter, - partition_name, - **kwargs) + async def delete( + self, + collection_name: str, + ids: list | str | int | None = None, + timeout: float | None = None, + filter: str | None = None, + partition_name: str | None = None, + **kwargs, + ): + return await self.async_milvus_client.delete(collection_name, ids, timeout, filter, partition_name, **kwargs) @classmethod def create_schema(cls, **kwargs): @@ -264,4 +298,4 @@ class AsyncMilvusClientWrapper: @logger_interceptor() async def close(self, **kwargs): - return await self.async_milvus_client.close(**kwargs) \ No newline at end of file + return await self.async_milvus_client.close(**kwargs) diff --git a/tests/python_client/milvus_client/test_milvus_client_force_merge.py b/tests/python_client/milvus_client/test_milvus_client_force_merge.py index cb9b940825..1a13b99568 100644 --- a/tests/python_client/milvus_client/test_milvus_client_force_merge.py +++ b/tests/python_client/milvus_client/test_milvus_client_force_merge.py @@ -12,7 +12,10 @@ With maxSize=64MB and auto compaction disabled, small data volumes can trigger force merge compaction manually without interference from auto compaction. """ +import math +import os import time +from collections import Counter import numpy as np import pytest @@ -21,6 +24,7 @@ from common import common_func as cf from common import common_type as ct from common.common_type import CaseLabel, CheckTasks from common.constants import * # noqa: F403 +from minio import Minio from pymilvus import DataType from utils.util_log import test_log as log from utils.util_pymilvus import * # noqa: F403 @@ -45,6 +49,44 @@ default_string_field_name = ct.default_string_field_name max_int64 = (1 << 63) - 1 auto_target_size_mb = max_int64 // (1024 * 1024) + 1 # Triggers server auto target-size mode. default_max_size_mb = 1024 # Default segment max size in MB +actual_output_size_tolerance = 0.10 + + +def minio_endpoint(minio_host): + endpoint = (minio_host or "localhost").strip() + if "://" in endpoint: + endpoint = endpoint.split("://", 1)[1] + endpoint = endpoint.split("/", 1)[0] + if ":" not in endpoint: + endpoint = f"{endpoint}:9000" + return endpoint + + +def new_minio_client(minio_host): + return Minio( + minio_endpoint(minio_host), + access_key=os.getenv("MILVUS_MINIO_ACCESS_KEY", "minioadmin"), + secret_key=os.getenv("MILVUS_MINIO_SECRET_KEY", "minioadmin"), + secure=os.getenv("MILVUS_MINIO_SECURE", "false").lower() in ["1", "true", "yes"], + ) + + +def get_insert_log_sizes(minio_client, bucket, collection_id, segment_ids): + root_path = os.getenv("MILVUS_MINIO_ROOT_PATH", "files").strip("/") + prefix = f"{root_path}/insert_log/{collection_id}/" + sizes = {str(segment_id): 0 for segment_id in segment_ids} + + for item in minio_client.list_objects(bucket, prefix=prefix, recursive=True): + relative_parts = item.object_name[len(prefix) :].split("/") + if len(relative_parts) < 3: + continue + segment_id = relative_parts[1] + if segment_id in sizes: + sizes[segment_id] += item.size + + missing = [segment_id for segment_id, size in sizes.items() if size == 0] + assert not missing, f"No insert-log objects found for segments {missing} under {prefix}" + return {int(segment_id): size for segment_id, size in sizes.items()} class TestMilvusClientForceMergeInvalid(TestMilvusClientV2Base): @@ -129,12 +171,12 @@ class TestMilvusClientForceMergeValid(TestMilvusClientV2Base): """ @pytest.mark.tags(CaseLabel.L3) - def test_force_merge_default_target_size(self): + def test_manual_compaction_without_target_size(self): """ - target: test ForceMerge with default target_size (0 or not passed) + target: test ordinary manual compaction when target_size is omitted method: create collection, insert data, flush, compact without target_size - expected: Compaction completes successfully - note: L3 - requires config change (segment.maxSize=64MB) to trigger actual force merge + expected: Ordinary manual compaction completes successfully + note: Omitting target_size does not select Force Merge """ client = self._client() collection_name = cf.gen_unique_str(prefix) @@ -152,20 +194,10 @@ class TestMilvusClientForceMergeValid(TestMilvusClientV2Base): ] self.insert(client, collection_name, rows) self.flush(client, collection_name) - # 3. compact with default target_size (not passed) + # 3. compact without target_size; this is ordinary manual compaction compact_id = self.compact(client, collection_name)[0] - # 4. wait for compaction to complete - cost = 180 - start = time.time() - while True: - time.sleep(1) - res = self.get_compaction_state(client, compact_id)[0] - log.info(f"Compaction state: {res}") - if res == "Completed": - break - if time.time() - start > cost: - raise Exception(f"Compaction cost more than {cost}s") - log.info("ForceMerge with default target_size completed successfully") + assert self.wait_for_compaction_ready(client, compact_id, timeout=180) + log.info("Manual compaction without target_size completed successfully") @pytest.mark.tags(CaseLabel.L3) def test_force_merge_explicit_target_size(self): @@ -493,12 +525,99 @@ class TestMilvusClientForceMergeValid(TestMilvusClientV2Base): segment_count_after = len(segments_after) log.info(f"Segment count after ForceMerge: {segment_count_after}") - # 7. verify segment count reduced (or at least not increased) - assert segment_count_after <= segment_count_before, ( + # 7. verify segment count reduced + assert segment_count_after < segment_count_before, ( f"Expected fewer segments after ForceMerge, got {segment_count_after} >= {segment_count_before}" ) log.info(f"ForceMerge reduced segments from {segment_count_before} to {segment_count_after}") + @pytest.mark.tags(CaseLabel.L3) + def test_force_merge_target_size_controls_output_segments(self, minio_host, minio_bucket): + """ + target: prove explicit target_size controls Force Merge output count and size + method: create a 1.25x-2x target input scope, compact, inspect plans and insert logs + expected: rows and sources are conserved; output count matches target and sizes stay within tolerance + note: L3 - requires segment.maxSize=64MB, auto compaction disabled, and MinIO access + """ + client = self._client() + collection_name = cf.gen_unique_str(prefix) + dim = 1024 + batch_size = 2000 + num_batches = 18 + total_rows = batch_size * num_batches + target_size_mb = 80 + target_size_bytes = target_size_mb * 1024 * 1024 + minio_client = new_minio_client(minio_host) + assert minio_client.bucket_exists(minio_bucket), f"MinIO bucket {minio_bucket!r} does not exist" + + self.create_collection(client, collection_name, dim) + rng = np.random.default_rng(seed=19530) + for batch in range(num_batches): + vectors = rng.random((batch_size, dim), dtype=np.float32) + rows = [ + { + default_primary_key_field_name: batch * batch_size + index, + default_vector_field_name: vectors[index].tolist(), + } + for index in range(batch_size) + ] + self.insert(client, collection_name, rows) + self.flush(client, collection_name) + + segments_before = self.list_persistent_segments(client, collection_name)[0] + source_ids = {segment.segment_id for segment in segments_before} + assert len(source_ids) >= num_batches, ( + f"Expected at least {num_batches} sealed inputs, got {len(source_ids)}: {segments_before}" + ) + description = self.describe_collection(client, collection_name)[0] + collection_id = description["collection_id"] + input_sizes = get_insert_log_sizes(minio_client, minio_bucket, collection_id, source_ids) + total_input_size = sum(input_sizes.values()) + assert target_size_bytes * 1.25 < total_input_size < target_size_bytes * 2, ( + f"Fixture must produce between 1.25x and 2x target bytes; " + f"input={total_input_size}, target={target_size_bytes}, sizes={input_sizes}" + ) + + compact_id = self.compact(client, collection_name, target_size=target_size_mb)[0] + assert self.wait_for_compaction_ready(client, compact_id, timeout=600) + + plans = client.get_compaction_plans(compact_id).plans + planned_source_counts = Counter(segment_id for plan in plans for segment_id in plan.sources) + planned_source_ids = set(planned_source_counts) + assert planned_source_ids == source_ids, ( + f"Force Merge plans must cover every input: " + f"expected={source_ids}, actual={planned_source_ids}, plans={plans}" + ) + assert all(count == 1 for count in planned_source_counts.values()), ( + f"Force Merge source IDs must occur in exactly one plan: {planned_source_counts}" + ) + + segments_after = self.list_persistent_segments(client, collection_name)[0] + output_ids = {segment.segment_id for segment in segments_after} + assert output_ids.isdisjoint(source_ids), ( + f"Completed Force Merge must replace all source segments: sources={source_ids}, outputs={output_ids}" + ) + assert sum(segment.num_rows for segment in segments_after) == total_rows + + output_sizes = get_insert_log_sizes(minio_client, minio_bucket, collection_id, output_ids) + total_output_size = sum(output_sizes.values()) + rewrite_delta = abs(total_output_size - total_input_size) / total_input_size + assert rewrite_delta <= actual_output_size_tolerance, ( + f"Rewrite changed total insert-log bytes by {rewrite_delta:.2%}: " + f"before={total_input_size}, after={total_output_size}" + ) + + expected_output_count = math.ceil(total_input_size / target_size_bytes) + assert len(output_sizes) == expected_output_count == 2, ( + f"Expected two outputs for target={target_size_bytes}, got sizes={output_sizes}" + ) + max_output_size = max(output_sizes.values()) + assert target_size_bytes * (1 - actual_output_size_tolerance) <= max_output_size + assert all(size <= target_size_bytes * (1 + actual_output_size_tolerance) for size in output_sizes.values()), ( + f"Output insert-log size exceeded target tolerance: target={target_size_bytes}, sizes={output_sizes}" + ) + log.info(f"Force Merge target-size evidence: input={input_sizes}, output={output_sizes}, plans={plans}") + @pytest.mark.tags(CaseLabel.L3) def test_force_merge_max_int64_overflow(self): """ @@ -543,95 +662,3 @@ class TestMilvusClientForceMergeValid(TestMilvusClientV2Base): if time.time() - start > cost: raise Exception(f"Compaction cost more than {cost}s") log.info("ForceMerge with max_int64 target_size completed (overflow fix verified)") - - @pytest.mark.tags(CaseLabel.L3) - def test_force_merge_algorithm_selection_under_threshold(self): - """ - target: test ForceMerge uses maxFullSegmentsGrouping when segment count <= threshold - method: create collection, insert data to create 5 segments (< threshold 10), - trigger force merge with target_size - expected: Compaction completes, algorithm selection logged as maxFullSegmentsGrouping - note: L3 - Check Loki logs for 'using maxFullSegmentsGrouping algorithm' - Requires config change (segment.maxSize=64MB) to trigger actual force merge - """ - client = self._client() - collection_name = cf.gen_unique_str(prefix) - dim = 128 - num_segments = 5 # Under threshold (default 10) - # 1. create collection - self.create_collection(client, collection_name, dim) - # 2. insert data in batches to create multiple segments - rng = np.random.default_rng(seed=19530) - for batch in range(num_segments): - rows = [ - { - default_primary_key_field_name: batch * 1000 + i, - default_vector_field_name: list(rng.random((1, dim))[0]), - } - for i in range(1000) - ] - self.insert(client, collection_name, rows) - self.flush(client, collection_name) - log.info(f"Inserted batch {batch + 1}/{num_segments}") - # 3. compact with target_size to trigger ForceMerge - target_size = 2048 - compact_id = self.compact(client, collection_name, target_size=target_size)[0] - log.info(f"ForceMerge triggered with {num_segments} segments (expect maxFullSegmentsGrouping)") - # 4. wait for compaction to complete - cost = 300 - start = time.time() - while True: - time.sleep(1) - res = self.get_compaction_state(client, compact_id)[0] - log.info(f"Compaction state: {res}") - if res == "Completed": - break - if time.time() - start > cost: - raise Exception(f"Compaction cost more than {cost}s") - log.info(f"ForceMerge with {num_segments} segments completed (check logs for algorithm)") - - @pytest.mark.tags(CaseLabel.L3) - def test_force_merge_algorithm_selection_over_threshold(self): - """ - target: test ForceMerge uses largerGroupingSegments when segment count > threshold - method: create collection, insert data to create 15 segments (> threshold 10), - trigger force merge with target_size - expected: Compaction completes, algorithm selection logged as largerGroupingSegments - note: L3 - Check Loki logs for 'using largerGroupingSegments algorithm' - Requires config change (segment.maxSize=64MB) to trigger actual force merge - """ - client = self._client() - collection_name = cf.gen_unique_str(prefix) - dim = 128 - num_segments = 15 # Over threshold (default 10) - # 1. create collection - self.create_collection(client, collection_name, dim) - # 2. insert data in batches to create multiple segments - rng = np.random.default_rng(seed=19530) - for batch in range(num_segments): - rows = [ - { - default_primary_key_field_name: batch * 1000 + i, - default_vector_field_name: list(rng.random((1, dim))[0]), - } - for i in range(1000) - ] - self.insert(client, collection_name, rows) - self.flush(client, collection_name) - log.info(f"Inserted batch {batch + 1}/{num_segments}") - # 3. compact with target_size to trigger ForceMerge - target_size = 2048 - compact_id = self.compact(client, collection_name, target_size=target_size)[0] - log.info(f"ForceMerge triggered with {num_segments} segments (expect largerGroupingSegments)") - # 4. wait for compaction to complete - cost = 600 # Longer timeout for more segments - start = time.time() - while True: - time.sleep(1) - res = self.get_compaction_state(client, compact_id)[0] - log.info(f"Compaction state: {res}") - if res == "Completed": - break - if time.time() - start > cost: - raise Exception(f"Compaction cost more than {cost}s") - log.info(f"ForceMerge with {num_segments} segments completed (check logs for algorithm)") diff --git a/tests/python_client/milvus_client/test_milvus_client_optimize.py b/tests/python_client/milvus_client/test_milvus_client_optimize.py index 8bb977d308..2c8a920b61 100644 --- a/tests/python_client/milvus_client/test_milvus_client_optimize.py +++ b/tests/python_client/milvus_client/test_milvus_client_optimize.py @@ -13,6 +13,7 @@ L3 tests require Milvus configuration changes: enableAutoCompaction: false """ +import asyncio import time import numpy as np @@ -23,6 +24,7 @@ from common import common_type as ct from common.common_type import CaseLabel, CheckTasks from common.constants import * # noqa: F403 from pymilvus import DataType +from pymilvus.milvus_client.async_optimize_task import AsyncOptimizeTask from utils.util_log import test_log as log from utils.util_pymilvus import * # noqa: F403 @@ -164,6 +166,25 @@ class TestMilvusClientOptimizeInvalid(TestMilvusClientV2Base): check_items=error, ) + @pytest.mark.tags(CaseLabel.L1) + def test_optimize_oversized_target_size(self): + """ + target: test the first target_size above the signed-int64-MB maximum + method: call optimize with 9223372036854775808MB + expected: Client-side parsing rejects the value before submission + """ + client = self._client() + collection_name = cf.gen_unique_str(prefix) + self.create_collection(client, collection_name, default_dim) + error = {ct.err_code: 1, ct.err_msg: "target size too large"} + self.optimize( + client, + collection_name, + target_size="9223372036854775808MB", + check_task=CheckTasks.err_res, + check_items=error, + ) + class TestMilvusClientOptimizeValid(TestMilvusClientV2Base): """Test cases for optimize() with valid parameters""" @@ -191,6 +212,54 @@ class TestMilvusClientOptimizeValid(TestMilvusClientV2Base): assert isinstance(result.compaction_id, int) log.info(f"Optimize on empty collection completed: {result}") + @pytest.mark.tags(CaseLabel.L1) + @pytest.mark.parametrize( + "target_size", + [ + "1073741824B", + "1048576 KB", + "1024MB", + "1GB", + "1.5 gB", + " 1 gb ", + "1TB", + "1PB", + ], + ids=["B", "KB", "MB", "GB", "decimal-mixed-case", "whitespace", "TB", "PB"], + ) + def test_optimize_valid_target_size_formats(self, target_size): + """ + target: test all supported units plus decimal, case, and whitespace handling + method: optimize an empty collection with each valid representation + expected: Each representation is accepted and preserved in OptimizeResult + """ + client = self._client() + collection_name = cf.gen_unique_str(prefix) + self.create_collection(client, collection_name, default_dim) + + result = self.optimize(client, collection_name, target_size=target_size)[0] + + assert result.status == "success" + assert result.collection_name == collection_name + assert result.target_size == target_size + + @pytest.mark.tags(CaseLabel.L1) + def test_optimize_max_target_size(self): + """ + target: test the maximum accepted signed-int64-MB target_size + method: optimize an empty collection with 9223372036854775807MB + expected: The boundary value is accepted without overflow + """ + client = self._client() + collection_name = cf.gen_unique_str(prefix) + target_size = "9223372036854775807MB" + self.create_collection(client, collection_name, default_dim) + + result = self.optimize(client, collection_name, target_size=target_size)[0] + + assert result.status == "success" + assert result.target_size == target_size + @pytest.mark.tags(CaseLabel.L3) def test_optimize_default_target_size(self): """ @@ -429,11 +498,11 @@ class TestMilvusClientOptimizeValid(TestMilvusClientV2Base): log.info("Optimize task cancelled successfully") @pytest.mark.tags(CaseLabel.L3) - def test_optimize_verify_segment_count(self): + def test_optimize_refreshes_loaded_segments(self): """ - target: test optimize reduces segment count - method: create collection, insert in batches, check segments before/after optimize - expected: Fewer segments after optimize + target: test optimize refreshes an already-loaded collection after compaction + method: record loaded IDs, optimize, then poll loaded IDs without release/load + expected: Old IDs disappear, segment count decreases, and all rows remain queryable note: L3 - requires config change (segment.maxSize=64MB) """ client = self._client() @@ -454,30 +523,42 @@ class TestMilvusClientOptimizeValid(TestMilvusClientV2Base): self.insert(client, collection_name, rows) self.flush(client, collection_name) - # Get stable segment count before optimize + # Establish the loaded precondition and a stable pre-optimize view. assert self.wait_for_index_ready(client, collection_name, default_vector_field_name, timeout=300) - self.release_collection(client, collection_name) - self.load_collection(client, collection_name) + self.refresh_load(client, collection_name, timeout=300) segments_before = client.list_loaded_segments(collection_name) - segment_count_before = len(segments_before) - log.info(f"Segment count before optimize: {segment_count_before}") + segment_ids_before = {segment.segment_id for segment in segments_before} + assert len(segment_ids_before) > 1, f"Expected multiple loaded inputs, got {segments_before}" + log.info(f"Loaded segments before optimize: {segments_before}") - # Optimize (handles compaction + index rebuild + refresh load) - result = self.optimize(client, collection_name, target_size="2GB", timeout=600)[0] + # optimize() must rebuild indexes and refresh the loaded view itself. + result = self.optimize(client, collection_name, target_size="64MB", timeout=600)[0] assert result.status == "success" + progress = {getattr(stage, "value", stage) for stage in result.progress} + assert "refreshing load" in progress, f"Loaded optimize skipped refresh_load: {result.progress}" - # Release and reload to get updated segment info - assert self.wait_for_index_ready(client, collection_name, default_vector_field_name, timeout=300) - self.release_collection(client, collection_name) - self.load_collection(client, collection_name) - segments_after = client.list_loaded_segments(collection_name) - segment_count_after = len(segments_after) - log.info(f"Segment count after optimize: {segment_count_after}") + deadline = time.time() + 300 + segments_after = [] + while time.time() < deadline: + segments_after = client.list_loaded_segments(collection_name) + segment_ids_after = {segment.segment_id for segment in segments_after} + if len(segment_ids_after) < len(segment_ids_before) and segment_ids_after.isdisjoint(segment_ids_before): + break + time.sleep(2) + else: + pytest.fail( + f"Loaded segments did not converge after optimize without release/load: " + f"before={segments_before}, after={segments_after}" + ) - assert segment_count_after <= segment_count_before, ( - f"Expected fewer segments after optimize, got {segment_count_after} >= {segment_count_before}" - ) - log.info(f"Optimize reduced segments from {segment_count_before} to {segment_count_after}") + count_result = self.query( + client, + collection_name, + filter="", + output_fields=["count(*)"], + )[0] + assert count_result[0]["count(*)"] == num_batches * batch_size + log.info(f"Loaded segments converged after optimize: before={segments_before}, after={segments_after}") @pytest.mark.tags(CaseLabel.L3) def test_optimize_numeric_target_size(self): @@ -555,3 +636,103 @@ class TestMilvusClientOptimizeValid(TestMilvusClientV2Base): result = self.optimize(client, collection_name, target_size="2GB", timeout=300)[0] assert result.status == "success" log.info(f"Optimize with clustering key completed: {result}") + + +class TestAsyncMilvusClientOptimizeValid(TestMilvusClientV2Base): + """Live AsyncMilvusClient optimize coverage.""" + + @pytest.mark.asyncio + @pytest.mark.tags(CaseLabel.L3) + async def test_async_optimize_loaded_collection(self): + """ + target: test AsyncMilvusClient.optimize wait=False against a loaded collection + method: observe the task lifecycle and loaded segment IDs without release/load + expected: The task succeeds, old loaded IDs disappear, and all rows remain queryable + """ + sync_client = self._client() + self.init_async_milvus_client() + async_client = self.async_milvus_client_wrap + collection_name = cf.gen_unique_str(prefix) + dim = 128 + num_batches = 5 + batch_size = default_nb + collection_created = False + + try: + await async_client.create_collection(collection_name, dimension=dim) + collection_created = True + rng = np.random.default_rng(seed=19530) + for batch in range(num_batches): + vectors = rng.random((batch_size, dim), dtype=np.float32) + rows = [ + { + default_primary_key_field_name: batch * batch_size + index, + default_vector_field_name: vectors[index].tolist(), + } + for index in range(batch_size) + ] + await async_client.insert(collection_name, rows) + await async_client.flush(collection_name) + + await async_client.load_collection(collection_name) + segments_before = sync_client.list_loaded_segments(collection_name) + segment_ids_before = {segment.segment_id for segment in segments_before} + assert len(segment_ids_before) > 1, f"Expected multiple loaded optimize inputs: {segments_before}" + + task, check_result = await async_client.optimize( + collection_name, + target_size="64MB", + wait=False, + timeout=600, + ) + assert check_result + assert isinstance(task, AsyncOptimizeTask) + + observed_progress = set() + deadline = time.time() + 600 + while not task.done(): + stage = task.progress() + observed_progress.add(getattr(stage, "value", stage)) + if time.time() >= deadline: + pytest.fail(f"Async optimize task did not complete; progress={observed_progress}") + await asyncio.sleep(0.5) + + result = await task.result(timeout=10) + + assert result.status == "success" + assert result.collection_name == collection_name + assert result.target_size == "64MB" + progress = {getattr(stage, "value", stage) for stage in result.progress} + assert "compacting" in progress + assert "waiting for index rebuild" in progress + assert "refreshing load" in progress + assert observed_progress - {"initializing"}, ( + f"Async task exposed no live progress beyond initialization: {observed_progress}" + ) + + refresh_deadline = time.time() + 300 + segments_after = [] + while time.time() < refresh_deadline: + segments_after = sync_client.list_loaded_segments(collection_name) + segment_ids_after = {segment.segment_id for segment in segments_after} + if len(segment_ids_after) < len(segment_ids_before) and segment_ids_after.isdisjoint( + segment_ids_before + ): + break + await asyncio.sleep(2) + else: + pytest.fail( + f"Loaded segments did not converge after async optimize without release/load: " + f"before={segments_before}, after={segments_after}" + ) + + count_result, _ = await async_client.query( + collection_name, + filter="", + output_fields=["count(*)"], + ) + assert count_result[0]["count(*)"] == num_batches * batch_size + finally: + if collection_created: + await async_client.drop_collection(collection_name) + await async_client.close() diff --git a/tests/python_client/requirements.txt b/tests/python_client/requirements.txt index 029abfb5db..4c26f51f33 100644 --- a/tests/python_client/requirements.txt +++ b/tests/python_client/requirements.txt @@ -23,8 +23,8 @@ pytest-sugar==0.9.5 pytest-random-order # pymilvus -pymilvus==3.1.0rc62 -pymilvus[bulk_writer]==3.1.0rc62 +pymilvus==3.1.0rc64 +pymilvus[bulk_writer]==3.1.0rc64 # for protobuf protobuf>=5.29.5