milvus/internal/datacoord/index_meta_test.go
yiwangdr 037de8e4d3
enhance: speed up minor functions calls in datacoord (#32389)
Related to https://github.com/milvus-io/milvus/issues/32165

1. nodeid based channel store access should use map access instead of
iteration.

2. The join-ish functions calls are slow when # collections/segments
increases (e.g. 10k).
e.g.
getNumRowsOfCollectionUnsafe is O(num_segments); GetAllCollectionNumRows
is of O(num_collections*num_segments).

Signed-off-by: yiwangdr <yiwangdr@gmail.com>
2024-04-20 07:55:21 +08:00

1347 lines
34 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 contains core functions in datacoord
package datacoord
import (
"context"
"sync"
"testing"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
mockkv "github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore"
"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/indexpb"
"github.com/milvus-io/milvus/pkg/common"
)
func TestReloadFromKV(t *testing.T) {
t.Run("ListIndexes_fail", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().ListIndexes(mock.Anything).Return(nil, errors.New("mock"))
_, err := newIndexMeta(context.TODO(), catalog)
assert.Error(t, err)
})
t.Run("ListSegmentIndexes_fails", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().ListIndexes(mock.Anything).Return([]*model.Index{}, nil)
catalog.EXPECT().ListSegmentIndexes(mock.Anything).Return(nil, errors.New("mock"))
_, err := newIndexMeta(context.TODO(), catalog)
assert.Error(t, err)
})
t.Run("ok", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().ListIndexes(mock.Anything).Return([]*model.Index{
{
CollectionID: 1,
IndexID: 1,
IndexName: "dix",
CreateTime: 1,
},
}, nil)
catalog.EXPECT().ListSegmentIndexes(mock.Anything).Return([]*model.SegmentIndex{
{
SegmentID: 1,
IndexID: 1,
},
}, nil)
meta, err := newIndexMeta(context.TODO(), catalog)
assert.NoError(t, err)
assert.NotNil(t, meta)
})
}
func TestMeta_CanCreateIndex(t *testing.T) {
var (
collID = UniqueID(1)
// partID = UniqueID(2)
indexID = UniqueID(10)
fieldID = UniqueID(100)
indexName = "_default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "FLAT",
},
}
)
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(nil)
m := newSegmentIndexMeta(catalog)
req := &indexpb.CreateIndexRequest{
CollectionID: collID,
FieldID: fieldID,
IndexName: indexName,
TypeParams: typeParams,
IndexParams: indexParams,
Timestamp: 0,
IsAutoIndex: false,
UserIndexParams: indexParams,
}
t.Run("can create index", func(t *testing.T) {
tmpIndexID, err := m.CanCreateIndex(req)
assert.NoError(t, err)
assert.Equal(t, int64(0), tmpIndexID)
index := &model.Index{
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
}
err = m.CreateIndex(index)
assert.NoError(t, err)
tmpIndexID, err = m.CanCreateIndex(req)
assert.NoError(t, err)
assert.Equal(t, indexID, tmpIndexID)
})
t.Run("params not consistent", func(t *testing.T) {
req.TypeParams = append(req.TypeParams, &commonpb.KeyValuePair{Key: "primary_key", Value: "false"})
tmpIndexID, err := m.CanCreateIndex(req)
assert.Error(t, err)
assert.Equal(t, int64(0), tmpIndexID)
req.TypeParams = []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "64"}}
tmpIndexID, err = m.CanCreateIndex(req)
assert.Error(t, err)
assert.Equal(t, int64(0), tmpIndexID)
req.TypeParams = typeParams
req.UserIndexParams = append(indexParams, &commonpb.KeyValuePair{Key: "metrics_type", Value: "L2"})
tmpIndexID, err = m.CanCreateIndex(req)
assert.Error(t, err)
assert.Equal(t, int64(0), tmpIndexID)
req.IndexParams = []*commonpb.KeyValuePair{{Key: common.IndexTypeKey, Value: "HNSW"}}
req.UserIndexParams = req.IndexParams
tmpIndexID, err = m.CanCreateIndex(req)
assert.Error(t, err)
assert.Equal(t, int64(0), tmpIndexID)
req.IndexParams = indexParams
req.UserIndexParams = indexParams
req.FieldID++
tmpIndexID, err = m.CanCreateIndex(req)
assert.Error(t, err)
assert.Equal(t, int64(0), tmpIndexID)
})
t.Run("multiple indexes", func(t *testing.T) {
req.IndexName = "_default_idx_2"
req.FieldID = fieldID
tmpIndexID, err := m.CanCreateIndex(req)
assert.Error(t, err)
assert.Equal(t, int64(0), tmpIndexID)
})
t.Run("index has been deleted", func(t *testing.T) {
m.indexes[collID][indexID].IsDeleted = true
tmpIndexID, err := m.CanCreateIndex(req)
assert.NoError(t, err)
assert.Equal(t, int64(0), tmpIndexID)
})
}
func TestMeta_HasSameReq(t *testing.T) {
var (
collID = UniqueID(1)
// partID = UniqueID(2)
indexID = UniqueID(10)
fieldID = UniqueID(100)
indexName = "_default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "FLAT",
},
}
)
m := newSegmentIndexMeta(catalogmocks.NewDataCoordCatalog(t))
req := &indexpb.CreateIndexRequest{
CollectionID: collID,
FieldID: fieldID,
IndexName: indexName,
TypeParams: typeParams,
IndexParams: indexParams,
Timestamp: 0,
IsAutoIndex: false,
UserIndexParams: indexParams,
}
t.Run("no indexes", func(t *testing.T) {
has, _ := m.HasSameReq(req)
assert.False(t, has)
})
t.Run("has same req", func(t *testing.T) {
m.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 10,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
},
}
has, _ := m.HasSameReq(req)
assert.True(t, has)
})
t.Run("params not consistent", func(t *testing.T) {
req.TypeParams = []*commonpb.KeyValuePair{{}}
has, _ := m.HasSameReq(req)
assert.False(t, has)
})
t.Run("index has been deleted", func(t *testing.T) {
m.indexes[collID][indexID].IsDeleted = true
has, _ := m.HasSameReq(req)
assert.False(t, has)
})
}
func newSegmentIndexMeta(catalog metastore.DataCoordCatalog) *indexMeta {
return &indexMeta{
RWMutex: sync.RWMutex{},
ctx: context.Background(),
catalog: catalog,
indexes: make(map[UniqueID]map[UniqueID]*model.Index),
buildID2SegmentIndex: make(map[UniqueID]*model.SegmentIndex),
segmentIndexes: make(map[UniqueID]map[UniqueID]*model.SegmentIndex),
}
}
func TestMeta_CreateIndex(t *testing.T) {
indexParams := []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "FLAT",
},
}
index := &model.Index{
TenantID: "",
CollectionID: 1,
FieldID: 2,
IndexID: 3,
IndexName: "_default_idx",
IsDeleted: false,
CreateTime: 12,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
}
t.Run("success", func(t *testing.T) {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(nil)
m := newSegmentIndexMeta(sc)
err := m.CreateIndex(index)
assert.NoError(t, err)
})
t.Run("save fail", func(t *testing.T) {
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("CreateIndex",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m := newSegmentIndexMeta(ec)
err := m.CreateIndex(index)
assert.Error(t, err)
})
}
func TestMeta_AddSegmentIndex(t *testing.T) {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("CreateSegmentIndex",
mock.Anything,
mock.Anything,
).Return(nil)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("CreateSegmentIndex",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m := newSegmentIndexMeta(ec)
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: make(map[UniqueID]*model.SegmentIndex, 0),
}
segmentIndex := &model.SegmentIndex{
SegmentID: 1,
CollectionID: 2,
PartitionID: 3,
NumRows: 10240,
IndexID: 4,
BuildID: 5,
NodeID: 6,
IndexVersion: 0,
IndexState: 0,
FailReason: "",
IsDeleted: false,
CreateTime: 12,
IndexFileKeys: nil,
IndexSize: 0,
}
t.Run("save meta fail", func(t *testing.T) {
err := m.AddSegmentIndex(segmentIndex)
assert.Error(t, err)
})
t.Run("success", func(t *testing.T) {
m.catalog = sc
err := m.AddSegmentIndex(segmentIndex)
assert.NoError(t, err)
})
}
func TestMeta_GetIndexIDByName(t *testing.T) {
var (
collID = UniqueID(1)
// partID = UniqueID(2)
indexID = UniqueID(10)
fieldID = UniqueID(100)
indexName = "_default_idx"
typeParams = []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
}
indexParams = []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "FLAT",
},
}
)
metakv := mockkv.NewMetaKv(t)
metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
m := newSegmentIndexMeta(&datacoord.Catalog{MetaKv: metakv})
t.Run("no indexes", func(t *testing.T) {
indexID2CreateTS := m.GetIndexIDByName(collID, indexName)
assert.Equal(t, 0, len(indexID2CreateTS))
})
t.Run("success", func(t *testing.T) {
m.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 12,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
},
}
indexID2CreateTS := m.GetIndexIDByName(collID, indexName)
assert.Contains(t, indexID2CreateTS, indexID)
})
}
func TestMeta_GetSegmentIndexState(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
indexID = UniqueID(10)
fieldID = 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: "FLAT",
},
}
)
metakv := mockkv.NewMetaKv(t)
metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
metakv.EXPECT().LoadWithPrefix(mock.Anything).Return(nil, nil, nil).Maybe()
m := newSegmentIndexMeta(&datacoord.Catalog{MetaKv: metakv})
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{
segID: make(map[UniqueID]*model.SegmentIndex, 0),
}
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) {
m.indexes[collID] = map[UniqueID]*model.Index{
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 12,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
},
}
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, indexID)
assert.Equal(t, commonpb.IndexState_Unissued, state.GetState())
assert.Contains(t, state.FailReason, "segment index not exist with ID")
})
t.Run("unissued", func(t *testing.T) {
m.updateSegmentIndex(&model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10250,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 0,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreateTime: 12,
IndexFileKeys: nil,
IndexSize: 0,
})
state := m.GetSegmentIndexState(collID, segID, indexID)
assert.Equal(t, commonpb.IndexState_Unissued, state.GetState())
})
t.Run("finish", func(t *testing.T) {
m.updateSegmentIndex(&model.SegmentIndex{
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 10250,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 0,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: 12,
IndexFileKeys: nil,
IndexSize: 0,
})
state := m.GetSegmentIndexState(collID, segID, indexID)
assert.Equal(t, commonpb.IndexState_Finished, state.GetState())
})
}
func TestMeta_GetIndexedSegment(t *testing.T) {
var (
collID = UniqueID(1)
partID = UniqueID(2)
indexID = UniqueID(10)
fieldID = 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: "FLAT",
},
}
)
m := newSegmentIndexMeta(nil)
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: nodeID,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: 10,
IndexFileKeys: nil,
IndexSize: 0,
},
},
}
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 10,
TypeParams: typeParams,
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
},
},
}
m.buildID2SegmentIndex = map[UniqueID]*model.SegmentIndex{
buildID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: nodeID,
IndexVersion: 1,
IndexState: commonpb.IndexState_Finished,
FailReason: "",
IsDeleted: false,
CreateTime: 10,
IndexFileKeys: nil,
IndexSize: 0,
},
}
t.Run("success", func(t *testing.T) {
segments := m.GetIndexedSegments(collID, []int64{segID}, []int64{fieldID})
assert.Len(t, segments, 1)
})
t.Run("no index on field", func(t *testing.T) {
segments := m.GetIndexedSegments(collID, []int64{segID}, []int64{fieldID + 1})
assert.Len(t, segments, 0)
})
t.Run("no index", func(t *testing.T) {
segments := m.GetIndexedSegments(collID+1, []int64{segID}, []int64{fieldID})
assert.Len(t, segments, 0)
})
}
func TestMeta_MarkIndexAsDeleted(t *testing.T) {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m := newSegmentIndexMeta(sc)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 10,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 1,
IndexID: indexID + 1,
IndexName: "_default_idx_102",
IsDeleted: true,
CreateTime: 1,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
t.Run("fail", func(t *testing.T) {
m.catalog = ec
err := m.MarkIndexAsDeleted(collID, []UniqueID{indexID, indexID + 1, indexID + 2})
assert.Error(t, err)
})
t.Run("success", func(t *testing.T) {
m.catalog = sc
err := m.MarkIndexAsDeleted(collID, []UniqueID{indexID, indexID + 1, indexID + 2})
assert.NoError(t, err)
err = m.MarkIndexAsDeleted(collID, []UniqueID{indexID, indexID + 1, indexID + 2})
assert.NoError(t, err)
err = m.MarkIndexAsDeleted(collID+1, []UniqueID{indexID, indexID + 1, indexID + 2})
assert.NoError(t, err)
})
}
func TestMeta_GetSegmentIndexes(t *testing.T) {
m := createMetaTable(&datacoord.Catalog{MetaKv: mockkv.NewMetaKv(t)})
t.Run("success", func(t *testing.T) {
segIndexes := m.indexMeta.getSegmentIndexes(segID)
assert.Equal(t, 1, len(segIndexes))
})
t.Run("segment not exist", func(t *testing.T) {
segIndexes := m.indexMeta.getSegmentIndexes(segID + 100)
assert.Equal(t, 0, len(segIndexes))
})
t.Run("no index exist- segment index empty", func(t *testing.T) {
m := newSegmentIndexMeta(nil)
segIndexes := m.GetSegmentIndexes(collID, segID)
assert.Equal(t, 0, len(segIndexes))
})
t.Run("no index exist- field index empty", func(t *testing.T) {
m := newSegmentIndexMeta(nil)
m.segmentIndexes = map[UniqueID]map[UniqueID]*model.SegmentIndex{
1: {
1: &model.SegmentIndex{},
},
}
segIndexes := m.GetSegmentIndexes(collID, 1)
assert.Equal(t, 0, len(segIndexes))
})
t.Run("index exists", func(t *testing.T) {
m := &indexMeta{
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
segID: {
indexID: &model.SegmentIndex{
CollectionID: collID,
SegmentID: segID,
IndexID: indexID,
IndexState: commonpb.IndexState_Finished,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
}
segIndexes := m.GetSegmentIndexes(collID, segID)
assert.Equal(t, 1, len(segIndexes))
segIdx, ok := segIndexes[indexID]
assert.True(t, ok)
assert.NotNil(t, segIdx)
})
}
func TestMeta_GetFieldIDByIndexID(t *testing.T) {
m := newSegmentIndexMeta(nil)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
t.Run("success", func(t *testing.T) {
fID := m.GetFieldIDByIndexID(collID, indexID)
assert.Equal(t, fieldID, fID)
})
t.Run("fail", func(t *testing.T) {
fID := m.GetFieldIDByIndexID(collID, indexID+1)
assert.Equal(t, UniqueID(0), fID)
})
}
func TestMeta_GetIndexNameByID(t *testing.T) {
m := newSegmentIndexMeta(nil)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
t.Run("success", func(t *testing.T) {
iName := m.GetIndexNameByID(collID, indexID)
assert.Equal(t, indexName, iName)
})
t.Run("fail", func(t *testing.T) {
iName := m.GetIndexNameByID(collID, indexID+1)
assert.Equal(t, "", iName)
})
}
func TestMeta_GetTypeParams(t *testing.T) {
indexParams := []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
}
m := newSegmentIndexMeta(nil)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
},
},
}
t.Run("success", func(t *testing.T) {
tp := m.GetTypeParams(collID, indexID)
assert.Equal(t, 1, len(tp))
})
t.Run("not exist", func(t *testing.T) {
tp := m.GetTypeParams(collID, indexID+1)
assert.Equal(t, 0, len(tp))
tp = m.GetTypeParams(collID+1, indexID)
assert.Equal(t, 0, len(tp))
})
}
func TestMeta_GetIndexParams(t *testing.T) {
indexParams := []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "HNSW",
},
}
m := newSegmentIndexMeta(nil)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: []*commonpb.KeyValuePair{
{
Key: common.DimKey,
Value: "128",
},
},
IndexParams: indexParams,
IsAutoIndex: false,
UserIndexParams: indexParams,
},
},
}
t.Run("success", func(t *testing.T) {
ip := m.GetIndexParams(collID, indexID)
assert.Equal(t, 1, len(ip))
})
t.Run("not exist", func(t *testing.T) {
ip := m.GetIndexParams(collID, indexID+1)
assert.Equal(t, 0, len(ip))
ip = m.GetIndexParams(collID+1, indexID)
assert.Equal(t, 0, len(ip))
})
}
func TestMeta_GetIndexJob(t *testing.T) {
m := newSegmentIndexMeta(nil)
m.buildID2SegmentIndex = map[UniqueID]*model.SegmentIndex{
buildID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: 1,
IndexVersion: 1,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreateTime: 0,
IndexFileKeys: nil,
IndexSize: 0,
},
}
t.Run("exist", func(t *testing.T) {
segIndex, exist := m.GetIndexJob(buildID)
assert.True(t, exist)
assert.NotNil(t, segIndex)
})
t.Run("not exist", func(t *testing.T) {
segIndex, exist := m.GetIndexJob(buildID + 1)
assert.False(t, exist)
assert.Nil(t, segIndex)
})
}
func TestMeta_IsIndexExist(t *testing.T) {
m := newSegmentIndexMeta(nil)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID + 1,
IndexName: "index2",
IsDeleted: true,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
t.Run("exist", func(t *testing.T) {
exist := m.IsIndexExist(collID, indexID)
assert.True(t, exist)
})
t.Run("not exist", func(t *testing.T) {
exist := m.IsIndexExist(collID, indexID+1)
assert.False(t, exist)
exist = m.IsIndexExist(collID, indexID+2)
assert.False(t, exist)
exist = m.IsIndexExist(collID+1, indexID)
assert.False(t, exist)
})
}
func updateSegmentIndexMeta(t *testing.T) *indexMeta {
sc := catalogmocks.NewDataCoordCatalog(t)
sc.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(nil)
return &indexMeta{
catalog: sc,
segmentIndexes: map[UniqueID]map[UniqueID]*model.SegmentIndex{
segID: {
indexID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 0,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreateTime: 0,
IndexFileKeys: nil,
IndexSize: 0,
},
},
},
indexes: map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
},
buildID2SegmentIndex: map[UniqueID]*model.SegmentIndex{
buildID: {
SegmentID: segID,
CollectionID: collID,
PartitionID: partID,
NumRows: 1025,
IndexID: indexID,
BuildID: buildID,
NodeID: 0,
IndexVersion: 0,
IndexState: commonpb.IndexState_Unissued,
FailReason: "",
IsDeleted: false,
CreateTime: 0,
IndexFileKeys: nil,
IndexSize: 0,
},
},
}
}
func TestMeta_UpdateVersion(t *testing.T) {
m := updateSegmentIndexMeta(t)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
t.Run("success", func(t *testing.T) {
err := m.UpdateVersion(buildID, nodeID)
assert.NoError(t, err)
})
t.Run("fail", func(t *testing.T) {
m.catalog = ec
err := m.UpdateVersion(buildID, nodeID)
assert.Error(t, err)
})
t.Run("not exist", func(t *testing.T) {
err := m.UpdateVersion(buildID+1, nodeID)
assert.Error(t, err)
})
}
func TestMeta_FinishTask(t *testing.T) {
m := updateSegmentIndexMeta(t)
t.Run("success", func(t *testing.T) {
err := m.FinishTask(&indexpb.IndexTaskInfo{
BuildID: buildID,
State: commonpb.IndexState_Finished,
IndexFileKeys: []string{"file1", "file2"},
SerializedSize: 1024,
FailReason: "",
})
assert.NoError(t, err)
})
t.Run("fail", func(t *testing.T) {
metakv := mockkv.NewMetaKv(t)
metakv.EXPECT().Save(mock.Anything, mock.Anything).Return(errors.New("failed")).Maybe()
metakv.EXPECT().MultiSave(mock.Anything).Return(errors.New("failed")).Maybe()
m.catalog = &datacoord.Catalog{
MetaKv: metakv,
}
err := m.FinishTask(&indexpb.IndexTaskInfo{
BuildID: buildID,
State: commonpb.IndexState_Finished,
IndexFileKeys: []string{"file1", "file2"},
SerializedSize: 1024,
FailReason: "",
})
assert.Error(t, err)
})
t.Run("not exist", func(t *testing.T) {
err := m.FinishTask(&indexpb.IndexTaskInfo{
BuildID: buildID + 1,
State: commonpb.IndexState_Finished,
IndexFileKeys: []string{"file1", "file2"},
SerializedSize: 1024,
FailReason: "",
})
assert.NoError(t, err)
})
}
func TestMeta_BuildIndex(t *testing.T) {
m := updateSegmentIndexMeta(t)
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
t.Run("success", func(t *testing.T) {
err := m.BuildIndex(buildID)
assert.NoError(t, err)
})
t.Run("fail", func(t *testing.T) {
m.catalog = ec
err := m.BuildIndex(buildID)
assert.Error(t, err)
})
t.Run("not exist", func(t *testing.T) {
err := m.BuildIndex(buildID + 1)
assert.Error(t, err)
})
}
// see also: https://github.com/milvus-io/milvus/issues/21660
func TestUpdateSegmentIndexNotExists(t *testing.T) {
m := newSegmentIndexMeta(nil)
assert.NotPanics(t, func() {
m.updateSegmentIndex(&model.SegmentIndex{
SegmentID: 1,
IndexID: 2,
})
})
assert.Equal(t, 1, len(m.segmentIndexes))
segmentIdx := m.segmentIndexes[1]
assert.Equal(t, 1, len(segmentIdx))
_, ok := segmentIdx[2]
assert.True(t, ok)
}
func TestMeta_DeleteTask_Error(t *testing.T) {
m := newSegmentIndexMeta(nil)
t.Run("segment index not found", func(t *testing.T) {
err := m.DeleteTask(buildID)
assert.NoError(t, err)
})
t.Run("segment update failed", func(t *testing.T) {
ec := catalogmocks.NewDataCoordCatalog(t)
ec.On("AlterSegmentIndexes",
mock.Anything,
mock.Anything,
).Return(errors.New("fail"))
m.catalog = ec
m.buildID2SegmentIndex[buildID] = &model.SegmentIndex{
SegmentID: segID,
PartitionID: partID,
CollectionID: collID,
}
err := m.DeleteTask(buildID)
assert.Error(t, err)
})
}
func TestMeta_GetFieldIndexes(t *testing.T) {
m := newSegmentIndexMeta(nil)
m.indexes = map[UniqueID]map[UniqueID]*model.Index{
collID: {
indexID: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID,
IndexName: indexName,
IsDeleted: true,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
indexID + 1: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID,
IndexID: indexID + 1,
IndexName: indexName,
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
indexID + 2: {
TenantID: "",
CollectionID: collID,
FieldID: fieldID + 2,
IndexID: indexID + 2,
IndexName: indexName + "2",
IsDeleted: false,
CreateTime: 0,
TypeParams: nil,
IndexParams: nil,
IsAutoIndex: false,
UserIndexParams: nil,
},
},
}
indexes := m.GetFieldIndexes(collID, fieldID, "")
assert.Equal(t, 1, len(indexes))
assert.Equal(t, indexName, indexes[0].IndexName)
}
func TestRemoveIndex(t *testing.T) {
t.Run("drop index fail", func(t *testing.T) {
expectedErr := errors.New("error")
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().
DropIndex(mock.Anything, mock.Anything, mock.Anything).
Return(expectedErr)
m := newSegmentIndexMeta(catalog)
err := m.RemoveIndex(collID, indexID)
assert.Error(t, err)
assert.EqualError(t, err, "error")
})
t.Run("remove index ok", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().
DropIndex(mock.Anything, mock.Anything, mock.Anything).
Return(nil)
m := &indexMeta{
catalog: catalog,
indexes: map[int64]map[int64]*model.Index{
collID: {
indexID: &model.Index{},
},
},
}
err := m.RemoveIndex(collID, indexID)
assert.NoError(t, err)
assert.Equal(t, len(m.indexes), 0)
})
}
func TestRemoveSegmentIndex(t *testing.T) {
t.Run("drop segment index fail", func(t *testing.T) {
expectedErr := errors.New("error")
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().
DropSegmentIndex(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(expectedErr)
m := newSegmentIndexMeta(catalog)
err := m.RemoveSegmentIndex(0, 0, 0, 0, 0)
assert.Error(t, err)
assert.EqualError(t, err, "error")
})
t.Run("remove segment index ok", func(t *testing.T) {
catalog := catalogmocks.NewDataCoordCatalog(t)
catalog.EXPECT().
DropSegmentIndex(mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil)
m := &indexMeta{
catalog: catalog,
segmentIndexes: map[int64]map[int64]*model.SegmentIndex{
segID: {
indexID: &model.SegmentIndex{},
},
},
buildID2SegmentIndex: map[int64]*model.SegmentIndex{
buildID: {},
},
}
err := m.RemoveSegmentIndex(collID, partID, segID, indexID, buildID)
assert.NoError(t, err)
assert.Equal(t, len(m.segmentIndexes), 0)
assert.Equal(t, len(m.buildID2SegmentIndex), 0)
})
}
func TestIndexMeta_GetUnindexedSegments(t *testing.T) {
m := createMetaTable(&datacoord.Catalog{MetaKv: mockkv.NewMetaKv(t)})
// normal case
segmentIDs := make([]int64, 0, 11)
for i := 0; i <= 10; i++ {
segmentIDs = append(segmentIDs, segID+int64(i))
}
unindexed := m.indexMeta.GetUnindexedSegments(collID, segmentIDs)
assert.Equal(t, 8, len(unindexed))
// no index
unindexed = m.indexMeta.GetUnindexedSegments(collID+1, segmentIDs)
assert.Equal(t, 0, len(unindexed))
}