Skip to content

Commit

Permalink
feat: Support multiple vector indexes in a collection (#27700)
Browse files Browse the repository at this point in the history 
issue: #25639 

/kind improvement
Signed-off-by: xige-16 <[email protected]>

---------

Signed-off-by: xige-16 <[email protected]>
  • Loading branch information
@xige-16
xige-16 authored Dec 29, 2023
1 parent 55af8f6 commit 0267391
Show file tree
Hide file tree
Showing 8 changed files with 141 additions and 85 deletions.
4 changes: 2 additions & 2 deletions internal/datacoord/compaction_trigger_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2220,7 +2220,7 @@ func (s *CompactionTriggerSuite) TestHandleSignal() {
defer s.SetupTest()
tr := s.tr
s.compactionHandler.EXPECT().isFull().Return(false)
// s.allocator.EXPECT().allocTimestamp(mock.Anything).Return(10000, nil)
s.allocator.EXPECT().allocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(nil, errors.New("mocked"))
tr.handleSignal(&compactionSignal{
segmentID: 1,
Expand All @@ -2237,7 +2237,7 @@ func (s *CompactionTriggerSuite) TestHandleSignal() {
defer s.SetupTest()
tr := s.tr
s.compactionHandler.EXPECT().isFull().Return(false)
// s.allocator.EXPECT().allocTimestamp(mock.Anything).Return(10000, nil)
s.allocator.EXPECT().allocTimestamp(mock.Anything).Return(10000, nil)
s.handler.EXPECT().GetCollection(mock.Anything, int64(100)).Return(&collectionInfo{
Properties: map[string]string{
common.CollectionAutoCompactionKey: "bad_value",
Expand Down
90 changes: 44 additions & 46 deletions internal/datacoord/index_meta.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import (

"github.com/golang/protobuf/proto"
"github.com/prometheus/client_golang/prometheus"
"github.com/samber/lo"
"go.uber.org/zap"

"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
Expand Down Expand Up @@ -271,77 +272,73 @@ func (m *meta) GetIndexIDByName(collID int64, indexName string) map[int64]uint64
return indexID2CreateTs
}

type IndexState struct {
state commonpb.IndexState
failReason string
}

func (m *meta) GetSegmentIndexState(collID, segmentID UniqueID) IndexState {
func (m *meta) GetSegmentIndexState(collID, segmentID UniqueID, indexID UniqueID) *indexpb.SegmentIndexState {
m.RLock()
defer m.RUnlock()

state := IndexState{
state: commonpb.IndexState_IndexStateNone,
failReason: "",
state := &indexpb.SegmentIndexState{
SegmentID: segmentID,
State: commonpb.IndexState_IndexStateNone,
FailReason: "",
}
fieldIndexes, ok := m.indexes[collID]
if !ok {
state.failReason = fmt.Sprintf("collection not exist with ID: %d", collID)
state.FailReason = fmt.Sprintf("collection not exist with ID: %d", collID)
return state
}
segment := m.segments.GetSegment(segmentID)
if segment != nil {
for indexID, index := range fieldIndexes {
if !index.IsDeleted {
if segIdx, ok := segment.segmentIndexes[indexID]; ok {
if segIdx.IndexState != commonpb.IndexState_Finished {
state.state = segIdx.IndexState
state.failReason = segIdx.FailReason
break
}
state.state = commonpb.IndexState_Finished
continue
}
state.state = commonpb.IndexState_Unissued
break
}
if segment == nil {
state.FailReason = fmt.Sprintf("segment is not exist with ID: %d", segmentID)
return state
}

if index, ok := fieldIndexes[indexID]; ok && !index.IsDeleted {
if segIdx, ok := segment.segmentIndexes[indexID]; ok {
state.IndexName = index.IndexName
state.State = segIdx.IndexState
state.FailReason = segIdx.FailReason
return state
}
state.State = commonpb.IndexState_Unissued
return state
}
state.failReason = fmt.Sprintf("segment is not exist with ID: %d", segmentID)

state.FailReason = fmt.Sprintf("there is no index on indexID: %d", indexID)
return state
}

func (m *meta) GetSegmentIndexStateOnField(collID, segmentID, fieldID UniqueID) IndexState {
func (m *meta) GetSegmentIndexStateOnField(collID, segmentID, fieldID UniqueID) *indexpb.SegmentIndexState {
m.RLock()
defer m.RUnlock()

state := IndexState{
state: commonpb.IndexState_IndexStateNone,
failReason: "",
state := &indexpb.SegmentIndexState{
SegmentID: segmentID,
State: commonpb.IndexState_IndexStateNone,
FailReason: "",
}
fieldIndexes, ok := m.indexes[collID]
if !ok {
state.failReason = fmt.Sprintf("collection not exist with ID: %d", collID)
state.FailReason = fmt.Sprintf("collection not exist with ID: %d", collID)
return state
}
segment := m.segments.GetSegment(segmentID)
if segment != nil {
for indexID, index := range fieldIndexes {
if index.FieldID == fieldID && !index.IsDeleted {
if segIdx, ok := segment.segmentIndexes[indexID]; ok {
state.state = segIdx.IndexState
state.failReason = segIdx.FailReason
return state
}
state.state = commonpb.IndexState_Unissued
if segment == nil {
state.FailReason = fmt.Sprintf("segment is not exist with ID: %d", segmentID)
return state
}
for indexID, index := range fieldIndexes {
if index.FieldID == fieldID && !index.IsDeleted {
if segIdx, ok := segment.segmentIndexes[indexID]; ok {
state.IndexName = index.IndexName
state.State = segIdx.IndexState
state.FailReason = segIdx.FailReason
return state
}
state.State = commonpb.IndexState_Unissued
return state
}
state.failReason = fmt.Sprintf("there is no index on fieldID: %d", fieldID)
return state
}
state.failReason = fmt.Sprintf("segment is not exist with ID: %d", segmentID)
state.FailReason = fmt.Sprintf("there is no index on fieldID: %d", fieldID)
return state
}

Expand Down Expand Up @@ -716,20 +713,21 @@ func (m *meta) GetHasUnindexTaskSegments() []*SegmentInfo {
m.RLock()
defer m.RUnlock()
segments := m.segments.GetSegments()
var ret []*SegmentInfo
unindexedSegments := make(map[int64]*SegmentInfo)
for _, segment := range segments {
if !isFlush(segment) {
continue
}
if fieldIndexes, ok := m.indexes[segment.CollectionID]; ok {
for _, index := range fieldIndexes {
if _, ok := segment.segmentIndexes[index.IndexID]; !index.IsDeleted && !ok {
ret = append(ret, segment)
unindexedSegments[segment.GetID()] = segment
}
}
}
}
return ret

return lo.MapToSlice(unindexedSegments, func(_ int64, segment *SegmentInfo) *SegmentInfo { return segment })
}

func (m *meta) GetMetasByNodeID(nodeID UniqueID) []*model.SegmentIndex {
Expand Down
72 changes: 57 additions & 15 deletions internal/datacoord/index_meta_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -475,9 +475,10 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
},
}

t.Run("segment has no index", func(t *testing.T) {
state := m.GetSegmentIndexState(collID, segID)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.state)
t.Run("collection has no index", func(t *testing.T) {
state := m.GetSegmentIndexState(collID, segID, indexID)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.GetState())
assert.Contains(t, state.GetFailReason(), "collection not exist with ID")
})

t.Run("meta not saved yet", func(t *testing.T) {
Expand All @@ -496,13 +497,14 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
UserIndexParams: indexParams,
},
}
state := m.GetSegmentIndexState(collID, segID)
assert.Equal(t, commonpb.IndexState_Unissued, state.state)
state := m.GetSegmentIndexState(collID, segID, indexID)
assert.Equal(t, commonpb.IndexState_Unissued, state.GetState())
})

t.Run("segment not exist", func(t *testing.T) {
state := m.GetSegmentIndexState(collID, segID+1)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.state)
state := m.GetSegmentIndexState(collID, segID+1, indexID)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.GetState())
assert.Contains(t, state.FailReason, "segment is not exist with ID")
})

t.Run("unissued", func(t *testing.T) {
Expand All @@ -523,8 +525,8 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
IndexSize: 0,
})

state := m.GetSegmentIndexState(collID, segID)
assert.Equal(t, commonpb.IndexState_Unissued, state.state)
state := m.GetSegmentIndexState(collID, segID, indexID)
assert.Equal(t, commonpb.IndexState_Unissued, state.GetState())
})

t.Run("finish", func(t *testing.T) {
Expand All @@ -545,8 +547,8 @@ func TestMeta_GetSegmentIndexState(t *testing.T) {
IndexSize: 0,
})

state := m.GetSegmentIndexState(collID, segID)
assert.Equal(t, commonpb.IndexState_Finished, state.state)
state := m.GetSegmentIndexState(collID, segID, indexID)
assert.Equal(t, commonpb.IndexState_Finished, state.GetState())
})
}

Expand Down Expand Up @@ -643,22 +645,22 @@ func TestMeta_GetSegmentIndexStateOnField(t *testing.T) {

t.Run("success", func(t *testing.T) {
state := m.GetSegmentIndexStateOnField(collID, segID, fieldID)
assert.Equal(t, commonpb.IndexState_Finished, state.state)
assert.Equal(t, commonpb.IndexState_Finished, state.GetState())
})

t.Run("no index on field", func(t *testing.T) {
state := m.GetSegmentIndexStateOnField(collID, segID, fieldID+1)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.state)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.GetState())
})

t.Run("no index", func(t *testing.T) {
state := m.GetSegmentIndexStateOnField(collID+1, segID, fieldID+1)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.state)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.GetState())
})

t.Run("segment not exist", func(t *testing.T) {
state := m.GetSegmentIndexStateOnField(collID, segID+1, fieldID)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.state)
assert.Equal(t, commonpb.IndexState_IndexStateNone, state.GetState())
})
}

Expand Down Expand Up @@ -1230,6 +1232,19 @@ func TestMeta_GetHasUnindexTaskSegments(t *testing.T) {
IsAutoIndex: false,
UserIndexParams: nil,
},
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 1,
IndexID: indexID + 1,
IndexName: indexName + "_1",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
}
Expand All @@ -1239,6 +1254,33 @@ func TestMeta_GetHasUnindexTaskSegments(t *testing.T) {
assert.Equal(t, 1, len(segments))
assert.Equal(t, segID, segments[0].ID)
})

t.Run("segment partial field with index", func(t *testing.T) {
m.segments.segments[segID].segmentIndexes = map[UniqueID]*model.SegmentIndex{
indexID: {
CollectionID: collID,
SegmentID: segID,
IndexID: indexID,
IndexState: commonpb.IndexState_Finished,
},
}

segments := m.GetHasUnindexTaskSegments()
assert.Equal(t, 1, len(segments))
assert.Equal(t, segID, segments[0].ID)
})

t.Run("segment all vector field with index", func(t *testing.T) {
m.segments.segments[segID].segmentIndexes[indexID+1] = &model.SegmentIndex{
CollectionID: collID,
SegmentID: segID,
IndexID: indexID + 1,
IndexState: commonpb.IndexState_Finished,
}

segments := m.GetHasUnindexTaskSegments()
assert.Equal(t, 0, len(segments))
})
}

// see also: https://github.com/milvus-io/milvus/issues/21660
Expand Down
12 changes: 5 additions & 7 deletions internal/datacoord/index_service.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -337,7 +337,7 @@ func (s *Server) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegme
)
log.Info("receive GetSegmentIndexState",
zap.String("IndexName", req.GetIndexName()),
zap.Int64s("fieldID", req.GetSegmentIDs()),
zap.Int64s("segmentIDs", req.GetSegmentIDs()),
)

if err := merr.CheckHealthy(s.GetStateCode()); err != nil {
Expand All @@ -360,12 +360,10 @@ func (s *Server) GetSegmentIndexState(ctx context.Context, req *indexpb.GetSegme
}, nil
}
for _, segID := range req.GetSegmentIDs() {
state := s.meta.GetSegmentIndexState(req.GetCollectionID(), segID)
ret.States = append(ret.States, &indexpb.SegmentIndexState{
SegmentID: segID,
State: state.state,
FailReason: state.failReason,
})
for indexID := range indexID2CreateTs {
state := s.meta.GetSegmentIndexState(req.GetCollectionID(), segID, indexID)
ret.States = append(ret.States, state)
}
}
log.Info("GetSegmentIndexState successfully", zap.String("indexName", req.GetIndexName()))
return ret, nil
Expand Down
7 changes: 3 additions & 4 deletions internal/datacoord/index_service_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -755,7 +755,7 @@ func TestServer_GetIndexState(t *testing.T) {
}},
}

t.Run("index state is node", func(t *testing.T) {
t.Run("index state is none", func(t *testing.T) {
resp, err := s.GetIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
Expand All @@ -766,7 +766,6 @@ func TestServer_GetIndexState(t *testing.T) {
s.meta.indexes[collID][indexID+1] = &model.Index{
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID + 1,
IndexName: "default_idx_1",
IsDeleted: false,
Expand Down Expand Up @@ -1833,7 +1832,7 @@ func TestServer_DropIndex(t *testing.T) {
indexID + 3: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 3,
FieldID: fieldID,
IndexID: indexID + 3,
IndexName: indexName + "_3",
IsDeleted: false,
Expand All @@ -1847,7 +1846,7 @@ func TestServer_DropIndex(t *testing.T) {
indexID + 4: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 4,
FieldID: fieldID,
IndexID: indexID + 4,
IndexName: indexName + "_4",
IsDeleted: false,
Expand Down
Loading

0 comments on commit 0267391

Please sign in to comment.