milvus/internal/storage/binlog_iterator_test.go

259 lines
6.8 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 storage
import (
"testing"
"github.com/milvus-io/milvus-proto/go-api/schemapb"
"github.com/milvus-io/milvus/internal/common"
"github.com/milvus-io/milvus/internal/proto/etcdpb"
"github.com/stretchr/testify/assert"
)
func generateTestData(t *testing.T, num int) []*Blob {
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: common.TimeStampField, Name: "ts", DataType: schemapb.DataType_Int64},
{FieldID: common.RowIDField, Name: "rowid", DataType: schemapb.DataType_Int64},
{FieldID: 101, Name: "int32", DataType: schemapb.DataType_Int32},
{FieldID: 102, Name: "floatVector", DataType: schemapb.DataType_FloatVector},
{FieldID: 103, Name: "binaryVector", DataType: schemapb.DataType_BinaryVector},
}}
insertCodec := NewInsertCodec(&etcdpb.CollectionMeta{ID: 1, Schema: schema})
var (
field0 []int64
field1 []int64
field101 []int32
field102 []float32
field103 []byte
)
for i := 1; i <= num; i++ {
field0 = append(field0, int64(i))
field1 = append(field1, int64(i))
field101 = append(field101, int32(i))
f102 := make([]float32, 8)
for j := range f102 {
f102[j] = float32(i)
}
field102 = append(field102, f102...)
field103 = append(field103, byte(i))
}
data := &InsertData{Data: map[FieldID]FieldData{
common.RowIDField: &Int64FieldData{Data: field0},
common.TimeStampField: &Int64FieldData{Data: field1},
101: &Int32FieldData{Data: field101},
102: &FloatVectorFieldData{
NumRows: []int64{int64(num)},
Data: field102,
Dim: 8,
},
103: &BinaryVectorFieldData{
NumRows: []int64{int64(num)},
Data: field103,
Dim: 8,
},
}}
blobs, _, err := insertCodec.Serialize(1, 1, data)
assert.Nil(t, err)
return blobs
}
func TestInsertlogIterator(t *testing.T) {
t.Run("empty iterator", func(t *testing.T) {
itr := &InsertBinlogIterator{
data: &InsertData{},
}
assert.False(t, itr.HasNext())
_, err := itr.Next()
assert.Equal(t, ErrNoMoreRecord, err)
})
t.Run("test dispose", func(t *testing.T) {
blobs := generateTestData(t, 1)
itr, err := NewInsertBinlogIterator(blobs, common.RowIDField, schemapb.DataType_Int64)
assert.Nil(t, err)
itr.Dispose()
assert.False(t, itr.HasNext())
_, err = itr.Next()
assert.Equal(t, ErrDisposed, err)
})
t.Run("not empty iterator", func(t *testing.T) {
blobs := generateTestData(t, 3)
itr, err := NewInsertBinlogIterator(blobs, common.RowIDField, schemapb.DataType_Int64)
assert.Nil(t, err)
for i := 1; i <= 3; i++ {
assert.True(t, itr.HasNext())
v, err := itr.Next()
assert.Nil(t, err)
value := v.(*Value)
f102 := make([]float32, 8)
for j := range f102 {
f102[j] = float32(i)
}
pk := &Int64PrimaryKey{
Value: int64(i),
}
expected := &Value{
int64(i),
pk,
int64(i),
false,
map[FieldID]interface{}{
common.TimeStampField: int64(i),
common.RowIDField: int64(i),
101: int32(i),
102: f102,
103: []byte{byte(i)},
},
}
assert.EqualValues(t, expected, value)
}
assert.False(t, itr.HasNext())
_, err = itr.Next()
assert.Equal(t, ErrNoMoreRecord, err)
})
}
func TestMergeIterator(t *testing.T) {
t.Run("empty iterators", func(t *testing.T) {
iterators := make([]Iterator, 0)
for i := 0; i < 3; i++ {
iterators = append(iterators, &InsertBinlogIterator{
data: &InsertData{},
})
}
itr := NewMergeIterator(iterators)
assert.False(t, itr.HasNext())
_, err := itr.Next()
assert.Equal(t, ErrNoMoreRecord, err)
})
t.Run("empty and non-empty iterators", func(t *testing.T) {
blobs := generateTestData(t, 3)
insertItr, err := NewInsertBinlogIterator(blobs, common.RowIDField, schemapb.DataType_Int64)
assert.Nil(t, err)
iterators := []Iterator{
&InsertBinlogIterator{data: &InsertData{}},
insertItr,
}
itr := NewMergeIterator(iterators)
for i := 1; i <= 3; i++ {
assert.True(t, itr.HasNext())
v, err := itr.Next()
assert.Nil(t, err)
value := v.(*Value)
f102 := make([]float32, 8)
for j := range f102 {
f102[j] = float32(i)
}
pk := &Int64PrimaryKey{
Value: int64(i),
}
expected := &Value{
int64(i),
pk,
int64(i),
false,
map[FieldID]interface{}{
common.TimeStampField: int64(i),
common.RowIDField: int64(i),
101: int32(i),
102: f102,
103: []byte{byte(i)},
},
}
assert.EqualValues(t, expected, value)
}
assert.False(t, itr.HasNext())
_, err = itr.Next()
assert.Equal(t, ErrNoMoreRecord, err)
})
t.Run("non-empty iterators", func(t *testing.T) {
blobs := generateTestData(t, 3)
itr1, err := NewInsertBinlogIterator(blobs, common.RowIDField, schemapb.DataType_Int64)
assert.Nil(t, err)
itr2, err := NewInsertBinlogIterator(blobs, common.RowIDField, schemapb.DataType_Int64)
assert.Nil(t, err)
iterators := []Iterator{itr1, itr2}
itr := NewMergeIterator(iterators)
for i := 1; i <= 3; i++ {
f102 := make([]float32, 8)
for j := range f102 {
f102[j] = float32(i)
}
pk := &Int64PrimaryKey{
Value: int64(i),
}
expected := &Value{
int64(i),
pk,
int64(i),
false,
map[FieldID]interface{}{
common.TimeStampField: int64(i),
common.RowIDField: int64(i),
101: int32(i),
102: f102,
103: []byte{byte(i)},
},
}
for j := 0; j < 2; j++ {
assert.True(t, itr.HasNext())
v, err := itr.Next()
assert.Nil(t, err)
value := v.(*Value)
assert.EqualValues(t, expected, value)
}
}
assert.False(t, itr.HasNext())
_, err = itr.Next()
assert.Equal(t, ErrNoMoreRecord, err)
})
t.Run("test dispose", func(t *testing.T) {
blobs := generateTestData(t, 3)
itr1, err := NewInsertBinlogIterator(blobs, common.RowIDField, schemapb.DataType_Int64)
assert.Nil(t, err)
itr := NewMergeIterator([]Iterator{itr1})
itr.Dispose()
assert.False(t, itr.HasNext())
_, err = itr.Next()
assert.Equal(t, ErrDisposed, err)
})
}