fix: add lazy load field to mark segment load type (#31591)

issue: https://github.com/milvus-io/milvus/issues/31673

---------

Signed-off-by: sunby <sunbingyi1992@gmail.com>

internal/querynodev2/segments/mock_segment.go

@@ -459,6 +459,47 @@ func (_c *MockSegment_InsertCount_Call) RunAndReturn(run func() int64) *MockSegm
 	return _c
 }
 
+// IsLazyLoad provides a mock function with given fields:
+func (_m *MockSegment) IsLazyLoad() bool {
+	ret := _m.Called()
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func() bool); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
+// MockSegment_IsLazyLoad_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsLazyLoad'
+type MockSegment_IsLazyLoad_Call struct {
+	*mock.Call
+}
+
+// IsLazyLoad is a helper method to define mock.On call
+func (_e *MockSegment_Expecter) IsLazyLoad() *MockSegment_IsLazyLoad_Call {
+	return &MockSegment_IsLazyLoad_Call{Call: _e.mock.On("IsLazyLoad")}
+}
+
+func (_c *MockSegment_IsLazyLoad_Call) Run(run func()) *MockSegment_IsLazyLoad_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockSegment_IsLazyLoad_Call) Return(_a0 bool) *MockSegment_IsLazyLoad_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *MockSegment_IsLazyLoad_Call) RunAndReturn(run func() bool) *MockSegment_IsLazyLoad_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // LastDeltaTimestamp provides a mock function with given fields:
 func (_m *MockSegment) LastDeltaTimestamp() uint64 {
 	ret := _m.Called()

internal/querynodev2/segments/retrieve.go

@@ -67,7 +67,7 @@ func retrieveOnSegments(ctx context.Context, mgr *Manager, segments []Segment, s
 			defer wg.Done()
 
 			var err error
-			if seg.LoadStatus() == LoadStatusMeta {
+			if seg.IsLazyLoad() {
 				err = mgr.DiskCache.Do(seg.ID(), retriever)
 			} else {
 				err = retriever(seg)

internal/querynodev2/segments/search.go

@@ -77,7 +77,7 @@ func searchSegments(ctx context.Context, mgr *Manager, segments []Segment, segTy
 				mu.Unlock()
 			}
 			var err error
-			if seg.LoadStatus() == LoadStatusMeta {
+			if seg.IsLazyLoad() {
 				err = mgr.DiskCache.Do(seg.ID(), searcher)
 			} else {
 				err = searcher(seg)

internal/querynodev2/segments/segment.go

@@ -89,6 +89,7 @@ type baseSegment struct {
 	segmentType    SegmentType
 	bloomFilterSet *pkoracle.BloomFilterSet
 	loadInfo       *querypb.SegmentLoadInfo
+	isLazyLoad     bool
 
 	resourceUsageCache *atomic.Pointer[ResourceUsage]
 }
@@ -192,6 +193,8 @@ func (s *baseSegment) ResourceUsageEstimate() ResourceUsage {
 	return *usage
 }
 
+func (s *baseSegment) IsLazyLoad() bool { return s.isLazyLoad }
+
 type FieldInfo struct {
 	datapb.FieldBinlog
 	RowCount int64

internal/querynodev2/segments/segment_interface.go

@@ -92,4 +92,5 @@ type Segment interface {
 	// Read operations
 	Search(ctx context.Context, searchReq *SearchRequest) (*SearchResult, error)
 	Retrieve(ctx context.Context, plan *RetrievePlan) (*segcorepb.RetrieveResults, error)
+	IsLazyLoad() bool
 }

internal/querynodev2/segments/segment_loader.go

@@ -1033,6 +1033,7 @@ func (loader *segmentLoader) LoadSegment(ctx context.Context,
 
 	if segment.Type() == SegmentTypeSealed {
 		if loadStatus == LoadStatusMeta {
+			segment.baseSegment.isLazyLoad = true
 			segment.baseSegment.loadInfo = loadInfo
 		}
 		if err := loader.loadSealedSegment(ctx, loadInfo, segment, collection, loadStatus); err != nil {