fix: Querycoord will trigger unexpected balance task after restart

Signed-off-by: Wei Liu <[email protected]>
milvus-io · Dec 24, 2024 · 5241ae8 · 5241ae8
1 parent afaabc2
commit 5241ae8
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/internal/proto/query_coord.proto b/internal/proto/query_coord.proto
@@ -827,6 +827,7 @@ message CheckerInfo {
 message SegmentTarget {
   int64 ID = 1;
   data.SegmentLevel level = 2; 
+  int64 num_of_rows = 3;
 }
 
 message PartitionTarget {

diff --git a/internal/querycoordv2/meta/target.go b/internal/querycoordv2/meta/target.go
@@ -57,6 +57,7 @@ func FromPbCollectionTarget(target *querypb.CollectionTarget) *CollectionTarget
 					CollectionID:  target.GetCollectionID(),
 					PartitionID:   partition.GetPartitionID(),
 					InsertChannel: t.GetChannelName(),
+					NumOfRows:     segment.GetNumOfRows(),
 				}
 			}
 			partitions = append(partitions, partition.GetPartitionID())
@@ -111,8 +112,9 @@ func (p *CollectionTarget) toPbMsg() *querypb.CollectionTarget {
 				}
 
 				partitionTarget.Segments = append(partitionTarget.Segments, &querypb.SegmentTarget{
-					ID:    info.GetID(),
-					Level: info.GetLevel(),
+					ID:        info.GetID(),
+					Level:     info.GetLevel(),
+					NumOfRows: info.GetNumOfRows(),
 				})
 			}
 		}

diff --git a/internal/querycoordv2/meta/target_manager_test.go b/internal/querycoordv2/meta/target_manager_test.go
@@ -570,11 +570,13 @@ func (suite *TargetManagerSuite) TestRecover() {
 			ID:            11,
 			PartitionID:   1,
 			InsertChannel: "channel-1",
+			NumOfRows:     100,
 		},
 		{
 			ID:            12,
 			PartitionID:   1,
 			InsertChannel: "channel-2",
+			NumOfRows:     100,
 		},
 	}
 
@@ -595,6 +597,9 @@ func (suite *TargetManagerSuite) TestRecover() {
 	suite.Len(target.GetAllDmChannelNames(), 2)
 	suite.Len(target.GetAllSegmentIDs(), 2)
 	suite.Equal(target.GetTargetVersion(), version)
+	for _, segment := range target.GetAllSegments() {
+		suite.Equal(int64(100), segment.GetNumOfRows())
+	}
 
 	// after recover, target info should be cleaned up
 	targets, err := suite.catalog.GetCollectionTargets()