Retruen segment state when get query segment info (#10730)

Signed-off-by: Nxige-16 <xi.ge@zilliz.com>

internal/core/src/pb/milvus.pb.h

@@ -10359,6 +10359,8 @@ class QuerySegmentInfo :
     kMemSizeFieldNumber = 4,
     kNumRowsFieldNumber = 5,
     kIndexIDFieldNumber = 7,
+    kNodeIDFieldNumber = 8,
+    kStateFieldNumber = 9,
   };
   // string index_name = 6;
   void clear_index_name();
@@ -10401,6 +10403,16 @@ class QuerySegmentInfo :
   ::PROTOBUF_NAMESPACE_ID::int64 indexid() const;
   void set_indexid(::PROTOBUF_NAMESPACE_ID::int64 value);
 
+  // int64 nodeID = 8;
+  void clear_nodeid();
+  ::PROTOBUF_NAMESPACE_ID::int64 nodeid() const;
+  void set_nodeid(::PROTOBUF_NAMESPACE_ID::int64 value);
+
+  // .milvus.proto.common.SegmentState state = 9;
+  void clear_state();
+  ::milvus::proto::common::SegmentState state() const;
+  void set_state(::milvus::proto::common::SegmentState value);
+
   // @@protoc_insertion_point(class_scope:milvus.proto.milvus.QuerySegmentInfo)
  private:
   class _Internal;
@@ -10413,6 +10425,8 @@ class QuerySegmentInfo :
   ::PROTOBUF_NAMESPACE_ID::int64 mem_size_;
   ::PROTOBUF_NAMESPACE_ID::int64 num_rows_;
   ::PROTOBUF_NAMESPACE_ID::int64 indexid_;
+  ::PROTOBUF_NAMESPACE_ID::int64 nodeid_;
+  int state_;
   mutable ::PROTOBUF_NAMESPACE_ID::internal::CachedSize _cached_size_;
   friend struct ::TableStruct_milvus_2eproto;
 };
@@ -21505,6 +21519,34 @@ inline void QuerySegmentInfo::set_indexid(::PROTOBUF_NAMESPACE_ID::int64 value)
   // @@protoc_insertion_point(field_set:milvus.proto.milvus.QuerySegmentInfo.indexID)
 }
 
+// int64 nodeID = 8;
+inline void QuerySegmentInfo::clear_nodeid() {
+  nodeid_ = PROTOBUF_LONGLONG(0);
+}
+inline ::PROTOBUF_NAMESPACE_ID::int64 QuerySegmentInfo::nodeid() const {
+  // @@protoc_insertion_point(field_get:milvus.proto.milvus.QuerySegmentInfo.nodeID)
+  return nodeid_;
+}
+inline void QuerySegmentInfo::set_nodeid(::PROTOBUF_NAMESPACE_ID::int64 value) {
+  
+  nodeid_ = value;
+  // @@protoc_insertion_point(field_set:milvus.proto.milvus.QuerySegmentInfo.nodeID)
+}
+
+// .milvus.proto.common.SegmentState state = 9;
+inline void QuerySegmentInfo::clear_state() {
+  state_ = 0;
+}
+inline ::milvus::proto::common::SegmentState QuerySegmentInfo::state() const {
+  // @@protoc_insertion_point(field_get:milvus.proto.milvus.QuerySegmentInfo.state)
+  return static_cast< ::milvus::proto::common::SegmentState >(state_);
+}
+inline void QuerySegmentInfo::set_state(::milvus::proto::common::SegmentState value) {
+  
+  state_ = value;
+  // @@protoc_insertion_point(field_set:milvus.proto.milvus.QuerySegmentInfo.state)
+}
+
 // -------------------------------------------------------------------
 
 // GetQuerySegmentInfoRequest

internal/proto/milvus.proto

@@ -663,6 +663,8 @@ message QuerySegmentInfo {
   int64 num_rows = 5;
   string index_name = 6;
   int64 indexID = 7;
+  int64 nodeID = 8;
+  common.SegmentState state = 9;
 }
 
 message GetQuerySegmentInfoRequest {

internal/proto/query_coord.proto

@@ -161,9 +161,10 @@ message SegmentInfo {
   string index_name = 7;
   int64 indexID = 8;
   string channelID = 9;
-  SegmentState segment_state = 10;
+  SegmentState segment_state = 10; // will deprecated
   repeated int64 compactionFrom = 11;
   bool createdByCompaction = 12;
+  common.SegmentState state = 13;
 }
 
 message GetSegmentInfoResponse {

internal/proxy/impl.go

@@ -2199,6 +2199,8 @@ func (node *Proxy) GetQuerySegmentInfo(ctx context.Context, req *milvuspb.GetQue
 			MemSize:      info.MemSize,
 			IndexName:    info.IndexName,
 			IndexID:      info.IndexID,
+			NodeID:       info.NodeID,
+			State:        info.State,
 		}
 	}
 	resp.Status.ErrorCode = commonpb.ErrorCode_Success

internal/querynode/impl.go

@@ -457,15 +457,17 @@ func (node *QueryNode) GetSegmentInfo(ctx context.Context, in *queryPb.GetSegmen
 	}
 	infos := make([]*queryPb.SegmentInfo, 0)
 	// TODO: remove segmentType and use queryPb.SegmentState instead
-	getSegmentStateBySegmentType := func(segType segmentType) queryPb.SegmentState {
+	getSegmentStateBySegmentType := func(segType segmentType) commonpb.SegmentState {
 		switch segType {
 		case segmentTypeGrowing:
-			return queryPb.SegmentState_Growing
+			return commonpb.SegmentState_Growing
 		case segmentTypeSealed:
-			return queryPb.SegmentState_sealed
+			return commonpb.SegmentState_Sealed
+		// TODO: remove segmentTypeIndexing
+		case segmentTypeIndexing:
+			return commonpb.SegmentState_Sealed
 		default:
-			// TODO: remove segmentTypeIndexing
-			return queryPb.SegmentState_None
+			return commonpb.SegmentState_NotExist
 		}
 	}
 	getSegmentInfo := func(segment *Segment) *queryPb.SegmentInfo {
@@ -487,7 +489,7 @@ func (node *QueryNode) GetSegmentInfo(ctx context.Context, in *queryPb.GetSegmen
 			IndexName:    indexName,
 			IndexID:      indexID,
 			ChannelID:    segment.vChannelID,
-			SegmentState: getSegmentStateBySegmentType(segment.segmentType),
+			State:        getSegmentStateBySegmentType(segment.segmentType),
 		}
 		return info
 	}

internal/querynode/impl_test.go

@@ -439,6 +439,11 @@ func TestImpl_GetSegmentInfo(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Equal(t, commonpb.ErrorCode_Success, rsp.Status.ErrorCode)
 
+		seg.setType(segmentTypeIndexing)
+		rsp, err = node.GetSegmentInfo(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, rsp.Status.ErrorCode)
+
 		seg.setType(-100)
 		rsp, err = node.GetSegmentInfo(ctx, req)
 		assert.NoError(t, err)