mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
test: improve force merge target size coverage (#51350)
issue: #51248 ## What changed - correct the omitted-`target_size` case to assert ordinary manual compaction semantics - add a physical Force Merge target-size case using persistent segment IDs and MinIO insert-log sizes - add deterministic Go boundary coverage for grouping selection at the configured threshold and threshold + 1 - add target-size memory clamp coverage for cluster and standalone co-location modes - extend Optimize format, signed-int64 boundary, loaded-segment refresh, and live async lifecycle coverage - add the async Optimize wrapper and update the PyMilvus test dependency to 3.1.0rc64 ## Why The previous cases did not prove that an explicit Force Merge target affected physical output, assumed the wrong grouping threshold, and depended on manual log inspection for algorithm selection. Loaded refresh and the public async Optimize workflow also lacked live end-to-end coverage. The async case exposed milvus-io/pymilvus#3680 on PyMilvus 3.1.0rc62. Version 3.1.0rc64 contains the tuple-unpacking fix and passes the unmodified test. ## Validation - all 10 planned case IDs pass; 0 failed and 0 blocked - `milvus-dev-cli` Go UT job `go-ut-local-zhuwenxi-zhuwenxing-i-4082040-30413568` succeeded; both test functions and all five subtests passed - physical target-size L3 case passed against a real cluster and MinIO in 200.68s - loaded refresh, target format/boundary, and manual compaction cases passed against the same server version - async Optimize L3 passed unmodified with PyMilvus 3.1.0rc64 in 56.91s - 56 Python nodes collect successfully; Python compilation, Ruff, and `git diff --check` pass --------- Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
This commit is contained in:
zhuwenxing
@@ -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)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
return await self.async_milvus_client.close(**kwargs)
|
||||
|
||||
@@ -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)")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user