GetQuerySegmentInfo returns sealed loaded segments (#18255)

Signed-off-by: yah01 <yang.cen@zilliz.com>

internal/querycoord/cluster.go

@@ -466,43 +466,10 @@ func (c *queryNodeCluster) GetSegmentInfo(ctx context.Context, in *querypb.GetSe
 		}
 	}
 
-	// Fetch growing segments
-	c.RLock()
-	var wg sync.WaitGroup
-	cnt := len(c.nodes)
-	resChan := make(chan respTuple, cnt)
-	wg.Add(cnt)
-	for _, node := range c.nodes {
-		go func(node Node) {
-			defer wg.Done()
-			res, err := node.getSegmentInfo(ctx, in)
-			resChan <- respTuple{
-				res: res,
-				err: err,
-			}
-		}(node)
-	}
-	c.RUnlock()
-	wg.Wait()
-	close(resChan)
-
-	for tuple := range resChan {
-		if tuple.err != nil {
-			return nil, tuple.err
-		}
-
-		segments := tuple.res.GetInfos()
-		for _, segment := range segments {
-			if segment.SegmentState != commonpb.SegmentState_Sealed {
-				segmentInfos = append(segmentInfos, segment)
-			}
-		}
-	}
-
-	//TODO::update meta
 	return segmentInfos, nil
 }
 
+// Deprecated
 func (c *queryNodeCluster) GetSegmentInfoByNode(ctx context.Context, nodeID int64, in *querypb.GetSegmentInfoRequest) ([]*querypb.SegmentInfo, error) {
 	c.RLock()
 	node, ok := c.nodes[nodeID]

internal/querycoord/cluster_test.go

@@ -549,43 +549,18 @@ func TestGrpcRequest(t *testing.T) {
 		assert.Nil(t, err)
 	})
 
-	t.Run("Test GetSegmentInfoByNode", func(t *testing.T) {
-		getSegmentInfoReq := &querypb.GetSegmentInfoRequest{
-			Base: &commonpb.MsgBase{
-				MsgType: commonpb.MsgType_SegmentInfo,
-			},
-			CollectionID: defaultCollectionID,
-		}
-		_, err = cluster.GetSegmentInfoByNode(baseCtx, nodeID, getSegmentInfoReq)
-		assert.Nil(t, err)
-	})
-
-	node.getSegmentInfos = returnFailedGetSegmentInfoResult
-
-	t.Run("Test GetSegmentInfoFailed", func(t *testing.T) {
+	t.Run("Test GetSegmentInfoNotExist", func(t *testing.T) {
 		getSegmentInfoReq := &querypb.GetSegmentInfoRequest{
 			Base: &commonpb.MsgBase{
 				MsgType: commonpb.MsgType_SegmentInfo,
 			},
 			CollectionID: defaultCollectionID,
+			SegmentIDs:   []UniqueID{-1},
 		}
 		_, err = cluster.GetSegmentInfo(baseCtx, getSegmentInfoReq)
-		assert.NotNil(t, err)
+		assert.Error(t, err)
 	})
 
-	t.Run("Test GetSegmentInfoByNodeFailed", func(t *testing.T) {
-		getSegmentInfoReq := &querypb.GetSegmentInfoRequest{
-			Base: &commonpb.MsgBase{
-				MsgType: commonpb.MsgType_SegmentInfo,
-			},
-			CollectionID: defaultCollectionID,
-		}
-		_, err = cluster.GetSegmentInfoByNode(baseCtx, nodeID, getSegmentInfoReq)
-		assert.NotNil(t, err)
-	})
-
-	node.getSegmentInfos = returnSuccessGetSegmentInfoResult
-
 	t.Run("Test GetNodeInfoByID", func(t *testing.T) {
 		res, err := cluster.GetNodeInfoByID(nodeID)
 		assert.Nil(t, err)

internal/querycoord/impl.go

@@ -809,9 +809,7 @@ func (qc *QueryCoord) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmen
 
 	totalMemSize := int64(0)
 	totalNumRows := int64(0)
-	//TODO::get segment infos from MetaReplica
-	//segmentIDs := req.SegmentIDs
-	//segmentInfos, err := qs.MetaReplica.getSegmentInfos(segmentIDs)
+
 	segmentInfos, err := qc.cluster.GetSegmentInfo(ctx, req)
 	if err != nil {
 		status.ErrorCode = commonpb.ErrorCode_UnexpectedError
@@ -836,7 +834,9 @@ func (qc *QueryCoord) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmen
 		zap.Int64("collectionID", req.CollectionID),
 		zap.Int64("msgID", req.Base.MsgID),
 		zap.Int64("num rows", totalNumRows),
-		zap.Int64("memory size", totalMemSize))
+		zap.Int64("memory size", totalMemSize),
+		zap.Int("segmentNum", len(segmentInfos)))
+
 	return &querypb.GetSegmentInfoResponse{
 		Status: status,
 		Infos:  segmentInfos,

tests/python_client/testcases/test_delete.py

@@ -968,7 +968,7 @@ class TestDeleteOperation(TestcaseBase):
         while True:
             time.sleep(0.5)
             segment_infos = self.utility_wrap.get_query_segment_info(collection_w.name)[0]
-            if segment_infos[0].state == SegmentState.Sealed:
+            if len(segment_infos) > 0 and segment_infos[0].state == SegmentState.Sealed:
                 break
         # query deleted id
         collection_w.query(tmp_expr, check_task=CheckTasks.check_query_empty)
@@ -1580,7 +1580,7 @@ class TestDeleteString(TestcaseBase):
         while True:
             time.sleep(0.5)
             segment_infos = self.utility_wrap.get_query_segment_info(collection_w.name)[0]
-            if segment_infos[0].state == SegmentState.Sealed:
+            if len(segment_infos) > 0 and segment_infos[0].state == SegmentState.Sealed:
                 break
         # query deleted id
         collection_w.query(default_string_expr, check_task=CheckTasks.check_query_empty)