milvus/internal/indexcoord/meta_table_test.go

332 lines
11 KiB
Go
Raw Normal View History

// 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 indexcoord
import (
"strconv"
"testing"
"github.com/milvus-io/milvus/internal/proto/schemapb"
"github.com/golang/protobuf/proto"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/proto/commonpb"
"github.com/milvus-io/milvus/internal/proto/indexpb"
"github.com/stretchr/testify/assert"
)
func TestMetaTable(t *testing.T) {
Params.Init()
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
assert.Nil(t, err)
req := &indexpb.BuildIndexRequest{
IndexBuildID: 1,
IndexName: "test_index",
IndexID: 0,
DataPaths: []string{"DataPath-1-1", "DataPath-1-2"},
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-2", Value: "IndexParam-1-2"}},
NumRows: 100,
FieldSchema: &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
},
}
indexMeta1 := &indexpb.IndexMeta{
IndexBuildID: 1,
State: commonpb.IndexState_Finished,
Req: req,
IndexFilePaths: []string{"IndexFilePath-1-1", "IndexFilePath-1-2"},
NodeID: 0,
Version: 10,
Recycled: false,
}
value, err := proto.Marshal(indexMeta1)
assert.Nil(t, err)
key := "indexes/" + strconv.FormatInt(indexMeta1.IndexBuildID, 10)
err = etcdKV.Save(key, string(value))
assert.Nil(t, err)
metaTable, err := NewMetaTable(etcdKV)
assert.Nil(t, err)
assert.NotNil(t, metaTable)
t.Run("saveIndexMeta", func(t *testing.T) {
meta := &Meta{
indexMeta: indexMeta1,
revision: 10,
}
err = metaTable.saveIndexMeta(meta)
assert.NotNil(t, err)
})
t.Run("reloadMeta", func(t *testing.T) {
indexBuildID := UniqueID(3)
meta3, err := metaTable.reloadMeta(indexBuildID)
assert.NotNil(t, err)
assert.Nil(t, meta3)
})
t.Run("AddIndex", func(t *testing.T) {
req := &indexpb.BuildIndexRequest{
IndexBuildID: 1,
}
err = metaTable.AddIndex(req.IndexBuildID, req)
assert.NotNil(t, err)
})
t.Run("GetIndexMetaByIndexBuildID", func(t *testing.T) {
indexMeta := metaTable.GetIndexMetaByIndexBuildID(1)
assert.NotNil(t, indexMeta)
indexMeta2 := metaTable.GetIndexMetaByIndexBuildID(20)
assert.Nil(t, indexMeta2)
})
t.Run("BuildIndex", func(t *testing.T) {
err = metaTable.BuildIndex(UniqueID(4), 1)
assert.NotNil(t, err)
indexMeta1.NodeID = 2
value, err = proto.Marshal(indexMeta1)
assert.Nil(t, err)
key = "indexes/" + strconv.FormatInt(indexMeta1.IndexBuildID, 10)
err = etcdKV.Save(key, string(value))
assert.Nil(t, err)
err = metaTable.BuildIndex(indexMeta1.IndexBuildID, 1)
assert.Nil(t, err)
})
t.Run("UpdateVersion", func(t *testing.T) {
err = metaTable.UpdateVersion(UniqueID(5))
assert.NotNil(t, err)
indexMeta1.Version = indexMeta1.Version + 1
value, err = proto.Marshal(indexMeta1)
assert.Nil(t, err)
key = "indexes/" + strconv.FormatInt(indexMeta1.IndexBuildID, 10)
err = etcdKV.Save(key, string(value))
assert.Nil(t, err)
err = metaTable.UpdateVersion(indexMeta1.IndexBuildID)
assert.Nil(t, err)
})
t.Run("MarkIndexAsDeleted", func(t *testing.T) {
indexMeta1.Version = indexMeta1.Version + 1
value, err = proto.Marshal(indexMeta1)
assert.Nil(t, err)
key = "indexes/" + strconv.FormatInt(indexMeta1.IndexBuildID, 10)
err = etcdKV.Save(key, string(value))
assert.Nil(t, err)
err = metaTable.MarkIndexAsDeleted(indexMeta1.Req.IndexID)
assert.Nil(t, err)
})
t.Run("GetIndexState", func(t *testing.T) {
indexInfos := metaTable.GetIndexStates([]UniqueID{0, 1})
assert.Equal(t, 2, len(indexInfos))
assert.Equal(t, "index 0 not exists", indexInfos[0].Reason)
assert.Equal(t, "index 1 has been deleted", indexInfos[1].Reason)
})
t.Run("GetIndexFilePathInfo", func(t *testing.T) {
indexFilePathInfo, err := metaTable.GetIndexFilePathInfo(0)
assert.Nil(t, indexFilePathInfo)
assert.NotNil(t, err)
indexFilePathInfo2, err := metaTable.GetIndexFilePathInfo(1)
assert.Nil(t, indexFilePathInfo2)
assert.NotNil(t, err)
})
t.Run("UpdateRecycleState", func(t *testing.T) {
indexMeta1.Version = indexMeta1.Version + 1
value, err = proto.Marshal(indexMeta1)
assert.Nil(t, err)
key = "indexes/" + strconv.FormatInt(indexMeta1.IndexBuildID, 10)
err = etcdKV.Save(key, string(value))
assert.Nil(t, err)
err = metaTable.UpdateRecycleState(indexMeta1.IndexBuildID)
assert.Nil(t, err)
err = metaTable.UpdateRecycleState(indexMeta1.IndexBuildID)
assert.Nil(t, err)
err = metaTable.UpdateRecycleState(5)
assert.NotNil(t, err)
})
t.Run("HasSameReq", func(t *testing.T) {
req2 := &indexpb.BuildIndexRequest{
IndexBuildID: 6,
IndexName: "test_index",
IndexID: 2,
DataPaths: []string{"DataPath-1-1", "DataPath-1-2"},
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-2", Value: "IndexParam-1-2"}},
NumRows: 100,
FieldSchema: &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
},
}
err = metaTable.AddIndex(6, req2)
assert.Nil(t, err)
req3 := &indexpb.BuildIndexRequest{
IndexBuildID: 6,
IndexName: "test_index",
IndexID: 3,
DataPaths: []string{"DataPath-1-1", "DataPath-1-2"},
TypeParams: []*commonpb.KeyValuePair{{Key: "TypeParam-1-1", Value: "TypeParam-1-1"}, {Key: "dim", Value: "256"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-2", Value: "IndexParam-1-2"}},
NumRows: 10,
FieldSchema: &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
},
}
has, err := metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.IndexID = 2
req3.IndexName = "test_index1"
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.IndexName = "test_index"
req3.DataPaths = []string{"DataPath-1-1", "DataPath-1-2", "DataPath-1-3"}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.DataPaths = []string{"DataPath-1-1", "DataPath-1-3"}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.DataPaths = []string{"DataPath-1-1", "DataPath-1-2"}
req3.TypeParams = []*commonpb.KeyValuePair{{Key: "TypeParam-1-1", Value: "TypeParam-1-1"}}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.TypeParams = []*commonpb.KeyValuePair{{Key: "TypeParam-1-1", Value: "TypeParam-1-1"}, {Key: "TypeParam-1-3", Value: "TypeParam-1-3"}}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.TypeParams = []*commonpb.KeyValuePair{{Key: "TypeParam-1-1", Value: "TypeParam-1-1"}, {Key: "TypeParam-1-2", Value: "TypeParam-1-3"}}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.TypeParams = []*commonpb.KeyValuePair{{Key: "TypeParam-1-1", Value: "TypeParam-1-1"}, {Key: "TypeParam-1-2", Value: "TypeParam-1-2"}}
req3.IndexParams = []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.IndexParams = []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-3", Value: "IndexParam-1-3"}}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
req3.IndexParams = []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-2", Value: "IndexParam-1-3"}}
has, err = metaTable.HasSameReq(req3)
assert.Equal(t, false, has)
assert.NotNil(t, err)
})
t.Run("LoadMetaFromETCD", func(t *testing.T) {
req4 := &indexpb.BuildIndexRequest{
IndexBuildID: 7,
IndexName: "test_index",
IndexID: 4,
DataPaths: []string{"DataPath-1-1", "DataPath-1-2"},
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "128"}, {Key: "TypeParam-1-2", Value: "TypeParam-1-2"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-2", Value: "IndexParam-1-2"}},
NumRows: 10,
FieldSchema: &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
},
}
err = metaTable.AddIndex(7, req4)
assert.Nil(t, err)
ok := metaTable.LoadMetaFromETCD(8, 0)
assert.Equal(t, false, ok)
key = "indexes/" + strconv.FormatInt(req4.IndexBuildID, 10)
err = etcdKV.RemoveWithPrefix(key)
assert.Nil(t, err)
ok = metaTable.LoadMetaFromETCD(req4.IndexBuildID, 10)
assert.Equal(t, false, ok)
})
t.Run("GetNodeTaskStats", func(t *testing.T) {
req5 := &indexpb.BuildIndexRequest{
IndexBuildID: 9,
IndexName: "test_index",
IndexID: 5,
DataPaths: []string{"DataPath-1-1", "DataPath-1-2"},
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "10"}, {Key: "TypeParam-1-2", Value: "TypeParam-1-2"}},
IndexParams: []*commonpb.KeyValuePair{{Key: "IndexParam-1-1", Value: "IndexParam-1-1"}, {Key: "IndexParam-1-2", Value: "IndexParam-1-2"}},
NumRows: 10,
FieldSchema: &schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
},
}
err = metaTable.AddIndex(req5.IndexBuildID, req5)
assert.Nil(t, err)
err = metaTable.BuildIndex(req5.IndexBuildID, 4)
assert.Nil(t, err)
priorities := metaTable.GetNodeTaskStats()
assert.Equal(t, 1, priorities[4])
})
err = etcdKV.RemoveWithPrefix("indexes/")
assert.Nil(t, err)
}
func TestMetaTable_Error(t *testing.T) {
Params.Init()
etcdKV, err := etcdkv.NewEtcdKV(Params.EtcdEndpoints, Params.MetaRootPath)
assert.Nil(t, err)
t.Run("reloadFromKV error", func(t *testing.T) {
value := "indexMeta-1"
key := "indexes/" + strconv.FormatInt(2, 10)
err = etcdKV.Save(key, value)
assert.Nil(t, err)
meta, err := NewMetaTable(etcdKV)
assert.NotNil(t, err)
assert.Nil(t, meta)
err = etcdKV.RemoveWithPrefix(key)
assert.Nil(t, err)
})
err = etcdKV.RemoveWithPrefix("indexes/")
assert.Nil(t, err)
}