milvus/internal/datacoord/index_service_test.go
congqixia 41af0a98fa
Use go-api/v2 for milvus-proto (#24770)
Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2023-06-09 01:28:37 +08:00

1611 lines
44 KiB
Go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package datacoord
import (
"context"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore/kv/datacoord"
catalogmocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/metastore/model"
"github.com/milvus-io/milvus/internal/proto/datapb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/common"
)
func TestServerId(t *testing.T) {
s := &Server{session: &sessionutil.Session{ServerID: 0}}
assert.Equal(t, int64(0), s.serverID())
}
func TestServer_CreateIndex(t *testing.T) {
var (
collID = UniqueID(1)
fieldID = UniqueID(10)
//indexID = UniqueID(100)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
req = &indexpb.CreateIndexRequest{
CollectionID: collID,
FieldID: fieldID,
IndexName: indexName,
TypeParams: typeParams,
IndexParams: indexParams,
Timestamp: 100,
IsAutoIndex: false,
UserIndexParams: indexParams,
}
ctx = context.Background()
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("success", func(t *testing.T) {
resp, err := s.CreateIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
})
t.Run("server not healthy", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Abnormal)
resp, err := s.CreateIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetErrorCode())
})
t.Run("index not consistent", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Healthy)
req.FieldID++
resp, err := s.CreateIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
t.Run("alloc ID fail", func(t *testing.T) {
req.FieldID = fieldID
s.allocator = &FailsAllocator{allocIDSucceed: false}
s.meta.indexes = map[UniqueID]map[UniqueID]*model.Index{}
resp, err := s.CreateIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
t.Run("not support disk index", func(t *testing.T) {
s.allocator = newMockAllocator()
s.meta.indexes = map[UniqueID]map[UniqueID]*model.Index{}
req.IndexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "DISKANN",
},
}
s.indexNodeManager = NewNodeManager(ctx, defaultIndexNodeCreatorFunc)
resp, err := s.CreateIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
t.Run("save index fail", func(t *testing.T) {
s.meta.indexes = map[UniqueID]map[UniqueID]*model.Index{}
s.meta.catalog = &datacoord.Catalog{MetaKv: &saveFailKV{}}
req.IndexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
resp, err := s.CreateIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
}
func TestServer_GetIndexState(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.GetIndexStateRequest{
CollectionID: collID,
IndexName: "",
}
)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Initializing)
resp, err := s.GetIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetStatus().GetErrorCode())
})
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("index not exist", func(t *testing.T) {
resp, err := s.GetIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_IndexNotExist, resp.GetStatus().GetErrorCode())
})
s.meta = &meta{
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
InsertChannel: "",
NumOfRows: 10250,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS - 1,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS - 1,
},
},
segmentIndexes: nil,
currRows: 0,
allocations: nil,
lastFlushTime: time.Time{},
isCompacting: false,
lastWrittenTime: time.Time{},
},
}},
}
t.Run("index state is unissued", func(t *testing.T) {
resp, err := s.GetIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, commonpb.IndexState_InProgress, resp.GetState())
})
s.meta = &meta{
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
InsertChannel: "",
NumOfRows: 10250,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS - 1,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS - 1,
},
},
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 3000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_IndexStateNone,
FailReason: "",
IsDeleted: false,
CreateTime: 0,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
},
currRows: 0,
allocations: nil,
lastFlushTime: time.Time{},
isCompacting: false,
lastWrittenTime: time.Time{},
},
}},
}
t.Run("index state is node", func(t *testing.T) {
resp, err := s.GetIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, commonpb.IndexState_IndexStateNone, resp.GetState())
})
t.Run("ambiguous index name", func(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,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
}
resp, err := s.GetIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetStatus().GetErrorCode())
})
}
func TestServer_GetSegmentIndexState(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.GetSegmentIndexStateRequest{
CollectionID: collID,
IndexName: "",
SegmentIDs: []UniqueID{segID},
}
)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{}},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server is not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Abnormal)
resp, err := s.GetSegmentIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetStatus().GetErrorCode())
})
t.Run("no indexes", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Healthy)
resp, err := s.GetSegmentIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_IndexNotExist, resp.GetStatus().GetErrorCode())
})
t.Run("unfinished", func(t *testing.T) {
s.meta.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
}
s.meta.segments.segments[segID] = &SegmentInfo{
SegmentInfo: nil,
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10250,
IndexID: indexID,
BuildID: 10,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: []string{"file1", "file2"},
IndexSize: 1025,
WriteHandoff: false,
},
},
currRows: 0,
allocations: nil,
lastFlushTime: time.Time{},
isCompacting: false,
lastWrittenTime: time.Time{},
}
resp, err := s.GetSegmentIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
})
t.Run("finish", func(t *testing.T) {
s.meta.segments.segments[segID] = &SegmentInfo{
SegmentInfo: nil,
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10250,
IndexID: indexID,
BuildID: 10,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: []string{"file1", "file2"},
IndexSize: 1025,
WriteHandoff: false,
},
},
currRows: 0,
allocations: nil,
lastFlushTime: time.Time{},
isCompacting: false,
lastWrittenTime: time.Time{},
}
resp, err := s.GetSegmentIndexState(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
})
}
func TestServer_GetIndexBuildProgress(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.GetIndexBuildProgressRequest{
CollectionID: collID,
IndexName: "",
}
)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{}},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Initializing)
resp, err := s.GetIndexBuildProgress(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetStatus().GetErrorCode())
})
t.Run("no indexes", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Healthy)
resp, err := s.GetIndexBuildProgress(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_IndexNotExist, resp.GetStatus().GetErrorCode())
})
t.Run("unissued", func(t *testing.T) {
s.meta.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
}
s.meta.segments = &SegmentsInfo{
segments: map[UniqueID]*SegmentInfo{
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
InsertChannel: "",
NumOfRows: 10250,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS,
},
},
segmentIndexes: nil,
currRows: 10250,
allocations: nil,
lastFlushTime: time.Time{},
isCompacting: false,
lastWrittenTime: time.Time{},
},
},
}
resp, err := s.GetIndexBuildProgress(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, int64(10250), resp.GetTotalRows())
assert.Equal(t, int64(0), resp.GetIndexedRows())
})
t.Run("finish", func(t *testing.T) {
s.meta.segments = &SegmentsInfo{
segments: map[UniqueID]*SegmentInfo{
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
InsertChannel: "",
NumOfRows: 10250,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS,
},
},
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10250,
IndexID: indexID,
BuildID: 10,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: []string{"file1", "file2"},
IndexSize: 0,
WriteHandoff: false,
},
},
currRows: 10250,
allocations: nil,
lastFlushTime: time.Time{},
isCompacting: false,
lastWrittenTime: time.Time{},
},
},
}
resp, err := s.GetIndexBuildProgress(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, int64(10250), resp.GetTotalRows())
assert.Equal(t, int64(10250), resp.GetIndexedRows())
})
t.Run("multiple index", func(t *testing.T) {
s.meta.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 1,
IndexID: indexID + 1,
IndexName: "_default_idx_102",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
}
resp, err := s.GetIndexBuildProgress(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetStatus().GetErrorCode())
})
}
func TestServer_DescribeIndex(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
invalidSegID = UniqueID(1001)
buildID = UniqueID(10000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.DescribeIndexRequest{
CollectionID: collID,
IndexName: "",
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// deleted
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 1,
IndexID: indexID + 1,
IndexName: indexName + "_1",
IsDeleted: true,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// unissued
indexID + 2: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 2,
IndexID: indexID + 2,
IndexName: indexName + "_2",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// inProgress
indexID + 3: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 3,
IndexID: indexID + 3,
IndexName: indexName + "_3",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// failed
indexID + 4: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 4,
IndexID: indexID + 4,
IndexName: indexName + "_4",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// unissued
indexID + 5: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 5,
IndexID: indexID + 5,
IndexName: indexName + "_5",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{
invalidSegID: {
SegmentInfo: &datapb.SegmentInfo{
ID: invalidSegID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
// timesamp > index start time, will be filtered out
Timestamp: createTS + 1,
},
},
},
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS,
},
CreatedByCompaction: true,
CompactionFrom: []int64{segID - 1},
},
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
},
},
segID - 1: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Dropped,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS,
},
},
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID - 1,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreateTime: createTS,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreateTime: createTS,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
CreateTime: createTS,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
CreateTime: createTS,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
CreateTime: createTS,
},
},
},
}},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Initializing)
resp, err := s.DescribeIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetStatus().GetErrorCode())
})
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("success", func(t *testing.T) {
resp, err := s.DescribeIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, 5, len(resp.GetIndexInfos()))
})
t.Run("describe after drop index", func(t *testing.T) {
status, err := s.DropIndex(ctx, &indexpb.DropIndexRequest{
CollectionID: collID,
PartitionIDs: nil,
IndexName: "",
DropAll: true,
})
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
resp, err := s.DescribeIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_IndexNotExist, resp.GetStatus().GetErrorCode())
})
}
func TestServer_GetIndexStatistics(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
invalidSegID = UniqueID(1001)
buildID = UniqueID(10000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.GetIndexStatisticsRequest{
CollectionID: collID,
IndexName: "",
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// deleted
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 1,
IndexID: indexID + 1,
IndexName: indexName + "_1",
IsDeleted: true,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// unissued
indexID + 2: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 2,
IndexID: indexID + 2,
IndexName: indexName + "_2",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// inProgress
indexID + 3: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 3,
IndexID: indexID + 3,
IndexName: indexName + "_3",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// failed
indexID + 4: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 4,
IndexID: indexID + 4,
IndexName: indexName + "_4",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// unissued
indexID + 5: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 5,
IndexID: indexID + 5,
IndexName: indexName + "_5",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{
invalidSegID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
// timesamp > index start time, will be filtered out
Timestamp: createTS + 1,
},
},
},
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
StartPosition: &msgpb.MsgPosition{
Timestamp: createTS,
},
},
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 1: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 1,
BuildID: buildID + 1,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 3: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 3,
BuildID: buildID + 3,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_InProgress,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 4: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 4,
BuildID: buildID + 4,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Failed,
FailReason: "mock failed",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
indexID + 5: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID + 5,
BuildID: buildID + 5,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
},
},
}},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Initializing)
resp, err := s.GetIndexStatistics(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_NotReadyServe, resp.GetStatus().GetErrorCode())
})
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("success", func(t *testing.T) {
resp, err := s.GetIndexStatistics(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, 5, len(resp.GetIndexInfos()))
})
t.Run("describe after drop index", func(t *testing.T) {
status, err := s.DropIndex(ctx, &indexpb.DropIndexRequest{
CollectionID: collID,
PartitionIDs: nil,
IndexName: "",
DropAll: true,
})
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, status.GetErrorCode())
resp, err := s.GetIndexStatistics(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_IndexNotExist, resp.GetStatus().GetErrorCode())
})
}
func TestServer_DropIndex(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.DropIndexRequest{
CollectionID: collID,
IndexName: indexName,
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
s := &Server{
meta: &meta{
catalog: catalog,
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// deleted
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 1,
IndexID: indexID + 1,
IndexName: indexName + "_1",
IsDeleted: true,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// unissued
indexID + 2: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 2,
IndexID: indexID + 2,
IndexName: indexName + "_2",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// inProgress
indexID + 3: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 3,
IndexID: indexID + 3,
IndexName: indexName + "_3",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
// failed
indexID + 4: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 4,
IndexID: indexID + 4,
IndexName: indexName + "_4",
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
segments: &SegmentsInfo{map[UniqueID]*SegmentInfo{
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
},
segmentIndexes: nil,
},
}},
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Initializing)
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetErrorCode())
})
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("drop fail", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
s.meta.catalog = catalog
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
t.Run("drop one index", func(t *testing.T) {
s.meta.catalog = catalog
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
})
t.Run("drop one without indexName", func(t *testing.T) {
req = &indexpb.DropIndexRequest{
CollectionID: collID,
PartitionIDs: nil,
IndexName: "",
DropAll: false,
}
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_UnexpectedError, resp.GetErrorCode())
})
t.Run("drop all indexes", func(t *testing.T) {
req = &indexpb.DropIndexRequest{
CollectionID: collID,
PartitionIDs: nil,
IndexName: "",
DropAll: true,
}
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
})
t.Run("drop not exist index", func(t *testing.T) {
req = &indexpb.DropIndexRequest{
CollectionID: collID,
PartitionIDs: nil,
IndexName: "",
DropAll: true,
}
resp, err := s.DropIndex(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
})
}
func TestServer_GetIndexInfos(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
fieldID = UniqueID(10)
indexID = UniqueID(100)
segID = UniqueID(1000)
buildID = UniqueID(10000)
indexName = "default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
}
createTS = uint64(1000)
ctx = context.Background()
req = &indexpb.GetIndexInfoRequest{
CollectionID: collID,
SegmentIDs: []UniqueID{segID},
IndexName: indexName,
}
)
chunkManagerFactory := storage.NewChunkManagerFactoryWithParam(Params)
cli, err := chunkManagerFactory.NewPersistentStorageChunkManager(ctx)
assert.NoError(t, err)
s := &Server{
meta: &meta{
catalog: &datacoord.Catalog{MetaKv: mocks.NewMetaKv(t)},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
//finished
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: createTS,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
segments: &SegmentsInfo{
map[UniqueID]*SegmentInfo{
segID: {
SegmentInfo: &datapb.SegmentInfo{
ID: segID,
CollectionID: collID,
PartitionID: partID,
NumOfRows: 10000,
State: commonpb.SegmentState_Flushed,
MaxRowNum: 65536,
LastExpireTime: createTS,
},
segmentIndexes: map[UniqueID]*model.SegmentIndex{
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10000,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: createTS,
IndexFileKeys: nil,
IndexSize: 0,
WriteHandoff: false,
},
},
},
},
},
chunkManager: cli,
},
allocator: newMockAllocator(),
notifyIndexChan: make(chan UniqueID, 1),
}
t.Run("server not available", func(t *testing.T) {
s.stateCode.Store(commonpb.StateCode_Initializing)
resp, err := s.GetIndexInfos(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_DataCoordNA, resp.GetStatus().GetErrorCode())
})
s.stateCode.Store(commonpb.StateCode_Healthy)
t.Run("get segment index infos", func(t *testing.T) {
resp, err := s.GetIndexInfos(ctx, req)
assert.NoError(t, err)
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
assert.Equal(t, 1, len(resp.GetSegmentInfo()))
})
}