mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-11-30 02:48:45 +08:00
365 lines
10 KiB
Go
365 lines
10 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 syncmgr
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math/rand"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/mock"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"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-proto/go-api/v2/schemapb"
|
|
milvus_storage "github.com/milvus-io/milvus-storage/go/storage"
|
|
"github.com/milvus-io/milvus-storage/go/storage/options"
|
|
"github.com/milvus-io/milvus-storage/go/storage/schema"
|
|
"github.com/milvus-io/milvus/internal/datanode/metacache"
|
|
"github.com/milvus-io/milvus/internal/proto/datapb"
|
|
"github.com/milvus-io/milvus/internal/storage"
|
|
"github.com/milvus-io/milvus/pkg/common"
|
|
"github.com/milvus-io/milvus/pkg/util/paramtable"
|
|
"github.com/milvus-io/milvus/pkg/util/tsoutil"
|
|
)
|
|
|
|
type StorageV2SerializerSuite struct {
|
|
suite.Suite
|
|
|
|
collectionID int64
|
|
partitionID int64
|
|
segmentID int64
|
|
channelName string
|
|
|
|
schema *schemapb.CollectionSchema
|
|
storageCache *metacache.StorageV2Cache
|
|
mockCache *metacache.MockMetaCache
|
|
mockMetaWriter *MockMetaWriter
|
|
|
|
serializer *storageV2Serializer
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) SetupSuite() {
|
|
paramtable.Get().Init(paramtable.NewBaseTable())
|
|
|
|
s.collectionID = rand.Int63n(100) + 1000
|
|
s.partitionID = rand.Int63n(100) + 2000
|
|
s.segmentID = rand.Int63n(1000) + 10000
|
|
s.channelName = fmt.Sprintf("by-dev-rootcoord-dml0_%d_v1", s.collectionID)
|
|
s.schema = &schemapb.CollectionSchema{
|
|
Name: "sync_task_test_col",
|
|
Fields: []*schemapb.FieldSchema{
|
|
{FieldID: common.RowIDField, DataType: schemapb.DataType_Int64, Name: common.RowIDFieldName},
|
|
{FieldID: common.TimeStampField, DataType: schemapb.DataType_Int64, Name: common.TimeStampFieldName},
|
|
{
|
|
FieldID: 100,
|
|
Name: "pk",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPrimaryKey: true,
|
|
},
|
|
{
|
|
FieldID: 101,
|
|
Name: "vector",
|
|
DataType: schemapb.DataType_FloatVector,
|
|
TypeParams: []*commonpb.KeyValuePair{
|
|
{Key: common.DimKey, Value: "128"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
s.mockCache = metacache.NewMockMetaCache(s.T())
|
|
s.mockMetaWriter = NewMockMetaWriter(s.T())
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) SetupTest() {
|
|
storageCache, err := metacache.NewStorageV2Cache(s.schema)
|
|
s.Require().NoError(err)
|
|
s.storageCache = storageCache
|
|
|
|
s.mockCache.EXPECT().Collection().Return(s.collectionID)
|
|
s.mockCache.EXPECT().Schema().Return(s.schema)
|
|
|
|
s.serializer, err = NewStorageV2Serializer(storageCache, s.mockCache, s.mockMetaWriter)
|
|
s.Require().NoError(err)
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getSpace() *milvus_storage.Space {
|
|
tmpDir := s.T().TempDir()
|
|
space, err := milvus_storage.Open(fmt.Sprintf("file:///%s", tmpDir), options.NewSpaceOptionBuilder().
|
|
SetSchema(schema.NewSchema(s.storageCache.ArrowSchema(), &schema.SchemaOptions{
|
|
PrimaryColumn: "pk", VectorColumn: "vector", VersionColumn: common.TimeStampFieldName,
|
|
})).Build())
|
|
s.Require().NoError(err)
|
|
return space
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getBasicPack() *SyncPack {
|
|
pack := &SyncPack{}
|
|
|
|
pack.WithCollectionID(s.collectionID).
|
|
WithPartitionID(s.partitionID).
|
|
WithSegmentID(s.segmentID).
|
|
WithChannelName(s.channelName).
|
|
WithCheckpoint(&msgpb.MsgPosition{
|
|
Timestamp: 1000,
|
|
ChannelName: s.channelName,
|
|
})
|
|
|
|
return pack
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getEmptyInsertBuffer() *storage.InsertData {
|
|
buf, err := storage.NewInsertData(s.schema)
|
|
s.Require().NoError(err)
|
|
|
|
return buf
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getInsertBuffer() *storage.InsertData {
|
|
buf := s.getEmptyInsertBuffer()
|
|
|
|
// generate data
|
|
for i := 0; i < 10; i++ {
|
|
data := make(map[storage.FieldID]any)
|
|
data[common.RowIDField] = int64(i + 1)
|
|
data[common.TimeStampField] = int64(i + 1)
|
|
data[100] = int64(i + 1)
|
|
vector := lo.RepeatBy(128, func(_ int) float32 {
|
|
return rand.Float32()
|
|
})
|
|
data[101] = vector
|
|
err := buf.Append(data)
|
|
s.Require().NoError(err)
|
|
}
|
|
return buf
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getDeleteBuffer() *storage.DeleteData {
|
|
buf := &storage.DeleteData{}
|
|
for i := 0; i < 10; i++ {
|
|
pk := storage.NewInt64PrimaryKey(int64(i + 1))
|
|
ts := tsoutil.ComposeTSByTime(time.Now(), 0)
|
|
buf.Append(pk, ts)
|
|
}
|
|
return buf
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getDeleteBufferZeroTs() *storage.DeleteData {
|
|
buf := &storage.DeleteData{}
|
|
for i := 0; i < 10; i++ {
|
|
pk := storage.NewInt64PrimaryKey(int64(i + 1))
|
|
buf.Append(pk, 0)
|
|
}
|
|
return buf
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) getBfs() *metacache.BloomFilterSet {
|
|
bfs := metacache.NewBloomFilterSet()
|
|
fd, err := storage.NewFieldData(schemapb.DataType_Int64, &schemapb.FieldSchema{
|
|
FieldID: 101,
|
|
Name: "ID",
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
})
|
|
s.Require().NoError(err)
|
|
|
|
ids := []int64{1, 2, 3, 4, 5, 6, 7}
|
|
for _, id := range ids {
|
|
err = fd.AppendRow(id)
|
|
s.Require().NoError(err)
|
|
}
|
|
|
|
bfs.UpdatePKRange(fd)
|
|
return bfs
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) TestSerializeInsert() {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
s.storageCache.SetSpace(s.segmentID, s.getSpace())
|
|
|
|
s.Run("no_data", func() {
|
|
pack := s.getBasicPack()
|
|
pack.WithTimeRange(50, 100)
|
|
pack.WithDrop()
|
|
|
|
task, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.NoError(err)
|
|
taskV1, ok := task.(*SyncTaskV2)
|
|
s.Require().True(ok)
|
|
s.Equal(s.collectionID, taskV1.collectionID)
|
|
s.Equal(s.partitionID, taskV1.partitionID)
|
|
s.Equal(s.channelName, taskV1.channelName)
|
|
s.Equal(&msgpb.MsgPosition{
|
|
Timestamp: 1000,
|
|
ChannelName: s.channelName,
|
|
}, taskV1.checkpoint)
|
|
s.EqualValues(50, taskV1.tsFrom)
|
|
s.EqualValues(100, taskV1.tsTo)
|
|
s.True(taskV1.isDrop)
|
|
})
|
|
|
|
s.Run("empty_insert_data", func() {
|
|
pack := s.getBasicPack()
|
|
pack.WithTimeRange(50, 100)
|
|
pack.WithInsertData(s.getEmptyInsertBuffer()).WithBatchSize(0)
|
|
|
|
_, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("with_normal_data", func() {
|
|
pack := s.getBasicPack()
|
|
pack.WithTimeRange(50, 100)
|
|
pack.WithInsertData(s.getInsertBuffer()).WithBatchSize(10)
|
|
|
|
s.mockCache.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Return().Once()
|
|
|
|
task, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.NoError(err)
|
|
|
|
taskV2, ok := task.(*SyncTaskV2)
|
|
s.Require().True(ok)
|
|
s.Equal(s.collectionID, taskV2.collectionID)
|
|
s.Equal(s.partitionID, taskV2.partitionID)
|
|
s.Equal(s.channelName, taskV2.channelName)
|
|
s.Equal(&msgpb.MsgPosition{
|
|
Timestamp: 1000,
|
|
ChannelName: s.channelName,
|
|
}, taskV2.checkpoint)
|
|
s.EqualValues(50, taskV2.tsFrom)
|
|
s.EqualValues(100, taskV2.tsTo)
|
|
s.NotNil(taskV2.reader)
|
|
s.NotNil(taskV2.statsBlob)
|
|
})
|
|
|
|
s.Run("with_flush_segment_not_found", func() {
|
|
pack := s.getBasicPack()
|
|
pack.WithFlush()
|
|
|
|
s.mockCache.EXPECT().GetSegmentByID(s.segmentID).Return(nil, false).Once()
|
|
_, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("with_flush", func() {
|
|
pack := s.getBasicPack()
|
|
pack.WithTimeRange(50, 100)
|
|
pack.WithInsertData(s.getInsertBuffer()).WithBatchSize(10)
|
|
pack.WithFlush()
|
|
|
|
bfs := s.getBfs()
|
|
segInfo := metacache.NewSegmentInfo(&datapb.SegmentInfo{}, bfs)
|
|
metacache.UpdateNumOfRows(1000)(segInfo)
|
|
metacache.CompactTo(metacache.NullSegment)(segInfo)
|
|
s.mockCache.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) {
|
|
action(segInfo)
|
|
}).Return().Once()
|
|
s.mockCache.EXPECT().GetSegmentByID(s.segmentID).Return(segInfo, true).Once()
|
|
|
|
task, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.NoError(err)
|
|
|
|
taskV2, ok := task.(*SyncTaskV2)
|
|
s.Require().True(ok)
|
|
s.Equal(s.collectionID, taskV2.collectionID)
|
|
s.Equal(s.partitionID, taskV2.partitionID)
|
|
s.Equal(s.channelName, taskV2.channelName)
|
|
s.Equal(&msgpb.MsgPosition{
|
|
Timestamp: 1000,
|
|
ChannelName: s.channelName,
|
|
}, taskV2.checkpoint)
|
|
s.EqualValues(50, taskV2.tsFrom)
|
|
s.EqualValues(100, taskV2.tsTo)
|
|
s.NotNil(taskV2.mergedStatsBlob)
|
|
})
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) TestSerializeDelete() {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
s.Run("serialize_failed", func() {
|
|
pkField := s.serializer.pkField
|
|
s.serializer.pkField = &schemapb.FieldSchema{}
|
|
defer func() {
|
|
s.serializer.pkField = pkField
|
|
}()
|
|
pack := s.getBasicPack()
|
|
pack.WithDeleteData(s.getDeleteBufferZeroTs())
|
|
pack.WithTimeRange(50, 100)
|
|
|
|
_, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("serialize_failed_bad_pk", func() {
|
|
pkField := s.serializer.pkField
|
|
s.serializer.pkField = &schemapb.FieldSchema{
|
|
DataType: schemapb.DataType_Array,
|
|
}
|
|
defer func() {
|
|
s.serializer.pkField = pkField
|
|
}()
|
|
pack := s.getBasicPack()
|
|
pack.WithDeleteData(s.getDeleteBufferZeroTs())
|
|
pack.WithTimeRange(50, 100)
|
|
|
|
_, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("serialize_normal", func() {
|
|
pack := s.getBasicPack()
|
|
pack.WithDeleteData(s.getDeleteBuffer())
|
|
pack.WithTimeRange(50, 100)
|
|
|
|
task, err := s.serializer.EncodeBuffer(ctx, pack)
|
|
s.NoError(err)
|
|
|
|
taskV2, ok := task.(*SyncTaskV2)
|
|
s.Require().True(ok)
|
|
s.Equal(s.collectionID, taskV2.collectionID)
|
|
s.Equal(s.partitionID, taskV2.partitionID)
|
|
s.Equal(s.channelName, taskV2.channelName)
|
|
s.Equal(&msgpb.MsgPosition{
|
|
Timestamp: 1000,
|
|
ChannelName: s.channelName,
|
|
}, taskV2.checkpoint)
|
|
s.EqualValues(50, taskV2.tsFrom)
|
|
s.EqualValues(100, taskV2.tsTo)
|
|
s.NotNil(taskV2.deleteReader)
|
|
})
|
|
}
|
|
|
|
func (s *StorageV2SerializerSuite) TestBadSchema() {
|
|
mockCache := metacache.NewMockMetaCache(s.T())
|
|
mockCache.EXPECT().Collection().Return(s.collectionID).Once()
|
|
mockCache.EXPECT().Schema().Return(&schemapb.CollectionSchema{}).Once()
|
|
_, err := NewStorageV2Serializer(s.storageCache, mockCache, s.mockMetaWriter)
|
|
s.Error(err)
|
|
}
|
|
|
|
func TestStorageV2Serializer(t *testing.T) {
|
|
suite.Run(t, new(StorageV2SerializerSuite))
|
|
}
|