fix ut

Signed-off-by: bigsheeper <[email protected]>
milvus-io · Dec 5, 2024 · 3a2dbaa · 3a2dbaa
1 parent a44699b
commit 3a2dbaa
Showing 1 changed file with 114 additions and 76 deletions.
diff --git a/internal/datacoord/compaction_trigger_test.go b/internal/datacoord/compaction_trigger_test.go
@@ -122,24 +122,34 @@ func Test_compactionTrigger_force_without_index(t *testing.T) {
 		},
 	}
 
+	segInfo := &datapb.SegmentInfo{
+		ID:             1,
+		CollectionID:   collectionID,
+		PartitionID:    1,
+		LastExpireTime: 100,
+		NumOfRows:      100,
+		MaxRowNum:      300,
+		InsertChannel:  "ch1",
+		State:          commonpb.SegmentState_Flushed,
+		Binlogs:        binlogs,
+		Deltalogs:      deltaLogs,
+		IsSorted:       true,
+	}
 	m := &meta{
 		catalog:    catalog,
 		channelCPs: newChannelCps(),
 		segments: &SegmentsInfo{
 			segments: map[int64]*SegmentInfo{
 				1: {
-					SegmentInfo: &datapb.SegmentInfo{
-						ID:             1,
-						CollectionID:   collectionID,
-						PartitionID:    1,
-						LastExpireTime: 100,
-						NumOfRows:      100,
-						MaxRowNum:      300,
-						InsertChannel:  "ch1",
-						State:          commonpb.SegmentState_Flushed,
-						Binlogs:        binlogs,
-						Deltalogs:      deltaLogs,
-						IsSorted:       true,
+					SegmentInfo: segInfo,
+				},
+			},
+			secondaryIndexes: segmentInfoIndexes{
+				coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
+					collectionID: {
+						1: {
+							SegmentInfo: segInfo,
+						},
 					},
 				},
 			},
@@ -214,6 +224,76 @@ func Test_compactionTrigger_force(t *testing.T) {
 
 	mock0Allocator := newMock0Allocator(t)
 
+	seg1 := &SegmentInfo{
+		SegmentInfo: &datapb.SegmentInfo{
+			ID:             1,
+			CollectionID:   2,
+			PartitionID:    1,
+			LastExpireTime: 100,
+			NumOfRows:      100,
+			MaxRowNum:      300,
+			InsertChannel:  "ch1",
+			State:          commonpb.SegmentState_Flushed,
+			Binlogs: []*datapb.FieldBinlog{
+				{
+					Binlogs: []*datapb.Binlog{
+						{EntriesNum: 5, LogID: 1},
+					},
+				},
+			},
+			Deltalogs: []*datapb.FieldBinlog{
+				{
+					Binlogs: []*datapb.Binlog{
+						{EntriesNum: 5, LogID: 1},
+					},
+				},
+			},
+			IsSorted: true,
+		},
+	}
+
+	seg2 := &SegmentInfo{
+		SegmentInfo: &datapb.SegmentInfo{
+			ID:             2,
+			CollectionID:   2,
+			PartitionID:    1,
+			LastExpireTime: 100,
+			NumOfRows:      100,
+			MaxRowNum:      300,
+			InsertChannel:  "ch1",
+			State:          commonpb.SegmentState_Flushed,
+			Binlogs: []*datapb.FieldBinlog{
+				{
+					Binlogs: []*datapb.Binlog{
+						{EntriesNum: 5, LogID: 2},
+					},
+				},
+			},
+			Deltalogs: []*datapb.FieldBinlog{
+				{
+					Binlogs: []*datapb.Binlog{
+						{EntriesNum: 5, LogID: 2},
+					},
+				},
+			},
+			IsSorted: true,
+		},
+	}
+
+	seg3 := &SegmentInfo{
+		SegmentInfo: &datapb.SegmentInfo{
+			ID:             3,
+			CollectionID:   1111,
+			PartitionID:    1,
+			LastExpireTime: 100,
+			NumOfRows:      100,
+			MaxRowNum:      300,
+			InsertChannel:  "ch1",
+			State:          commonpb.SegmentState_Flushed,
+			IsSorted:       true,
+		},
+	}
+
 	tests := []struct {
 		name         string
 		fields       fields
@@ -230,71 +310,18 @@ func Test_compactionTrigger_force(t *testing.T) {
 					channelCPs: newChannelCps(),
 					segments: &SegmentsInfo{
 						segments: map[int64]*SegmentInfo{
-							1: {
-								SegmentInfo: &datapb.SegmentInfo{
-									ID:             1,
-									CollectionID:   2,
-									PartitionID:    1,
-									LastExpireTime: 100,
-									NumOfRows:      100,
-									MaxRowNum:      300,
-									InsertChannel:  "ch1",
-									State:          commonpb.SegmentState_Flushed,
-									Binlogs: []*datapb.FieldBinlog{
-										{
-											Binlogs: []*datapb.Binlog{
-												{EntriesNum: 5, LogID: 1},
-											},
-										},
-									},
-									Deltalogs: []*datapb.FieldBinlog{
-										{
-											Binlogs: []*datapb.Binlog{
-												{EntriesNum: 5, LogID: 1},
-											},
-										},
-									},
-									IsSorted: true,
-								},
-							},
-							2: {
-								SegmentInfo: &datapb.SegmentInfo{
-									ID:             2,
-									CollectionID:   2,
-									PartitionID:    1,
-									LastExpireTime: 100,
-									NumOfRows:      100,
-									MaxRowNum:      300,
-									InsertChannel:  "ch1",
-									State:          commonpb.SegmentState_Flushed,
-									Binlogs: []*datapb.FieldBinlog{
-										{
-											Binlogs: []*datapb.Binlog{
-												{EntriesNum: 5, LogID: 2},
-											},
-										},
-									},
-									Deltalogs: []*datapb.FieldBinlog{
-										{
-											Binlogs: []*datapb.Binlog{
-												{EntriesNum: 5, LogID: 2},
-											},
-										},
-									},
-									IsSorted: true,
+							1: seg1,
+							2: seg2,
+							3: seg3,
+						},
+						secondaryIndexes: segmentInfoIndexes{
+							coll2Segments: map[UniqueID]map[UniqueID]*SegmentInfo{
+								2: {
+									seg1.GetID(): seg1,
+									seg2.GetID(): seg2,
 								},
-							},
-							3: {
-								SegmentInfo: &datapb.SegmentInfo{
-									ID:             3,
-									CollectionID:   1111,
-									PartitionID:    1,
-									LastExpireTime: 100,
-									NumOfRows:      100,
-									MaxRowNum:      300,
-									InsertChannel:  "ch1",
-									State:          commonpb.SegmentState_Flushed,
-									IsSorted:       true,
+								1111: {
+									seg3.GetID(): seg3,
 								},
 							},
 						},
@@ -617,7 +644,13 @@ func Test_compactionTrigger_force(t *testing.T) {
 		t.Run(tt.name+" with DiskANN index", func(t *testing.T) {
 			for _, segment := range tt.fields.meta.segments.GetSegments() {
 				// Collection 1000 means it has DiskANN index
+				delete(tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()], segment.GetID())
 				segment.CollectionID = 1000
+				_, ok := tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()]
+				if !ok {
+					tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()] = make(map[UniqueID]*SegmentInfo)
+				}
+				tt.fields.meta.segments.secondaryIndexes.coll2Segments[segment.GetCollectionID()][segment.GetID()] = segment
 			}
 			tr := &compactionTrigger{
 				meta:                         tt.fields.meta,
@@ -725,6 +758,9 @@ func Test_compactionTrigger_force_maxSegmentLimit(t *testing.T) {
 	vecFieldID := int64(201)
 	segmentInfos := &SegmentsInfo{
 		segments: make(map[UniqueID]*SegmentInfo),
+		secondaryIndexes: segmentInfoIndexes{
+			coll2Segments: make(map[UniqueID]map[UniqueID]*SegmentInfo),
+		},
 	}
 
 	indexMeta := newSegmentIndexMeta(nil)
@@ -751,6 +787,7 @@ func Test_compactionTrigger_force_maxSegmentLimit(t *testing.T) {
 		},
 	}
 
+	segmentInfos.secondaryIndexes.coll2Segments[2] = make(map[UniqueID]*SegmentInfo)
 	for i := UniqueID(0); i < 50; i++ {
 		info := &SegmentInfo{
 			SegmentInfo: &datapb.SegmentInfo{
@@ -793,6 +830,7 @@ func Test_compactionTrigger_force_maxSegmentLimit(t *testing.T) {
 		})
 
 		segmentInfos.segments[i] = info
+		segmentInfos.secondaryIndexes.coll2Segments[2][i] = info
 	}
 
 	mock0Allocator := newMockAllocator(t)