wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
685 lines
24 KiB
Go
685 lines
24 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/mlog"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/merr"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
func TestInsertDataSuite(t *testing.T) {
|
|
suite.Run(t, new(InsertDataSuite))
|
|
}
|
|
|
|
func TestArrayFieldDataSuite(t *testing.T) {
|
|
suite.Run(t, new(ArrayFieldDataSuite))
|
|
}
|
|
|
|
type InsertDataSuite struct {
|
|
suite.Suite
|
|
|
|
schema *schemapb.CollectionSchema
|
|
|
|
iDataOneRow *InsertData
|
|
iDataTwoRows *InsertData
|
|
iDataEmpty *InsertData
|
|
}
|
|
|
|
func (s *InsertDataSuite) SetupSuite() {
|
|
s.schema = genTestCollectionMeta().Schema
|
|
}
|
|
|
|
func (s *InsertDataSuite) TestInsertData() {
|
|
s.Run("nil schema", func() {
|
|
idata, err := NewInsertData(nil)
|
|
s.Error(err)
|
|
s.Nil(idata)
|
|
})
|
|
|
|
s.Run("nullable field schema", func() {
|
|
tests := []struct {
|
|
description string
|
|
dataType schemapb.DataType
|
|
typeParams []*commonpb.KeyValuePair
|
|
nullable bool
|
|
}{
|
|
{"nullable bool field", schemapb.DataType_Bool, nil, true},
|
|
{"nullable int8 field", schemapb.DataType_Int8, nil, true},
|
|
{"nullable int16 field", schemapb.DataType_Int16, nil, true},
|
|
{"nullable int32 field", schemapb.DataType_Int32, nil, true},
|
|
{"nullable int64 field", schemapb.DataType_Int64, nil, true},
|
|
{"nullable float field", schemapb.DataType_Float, nil, true},
|
|
{"nullable double field", schemapb.DataType_Double, nil, true},
|
|
{"nullable json field", schemapb.DataType_JSON, nil, true},
|
|
{"nullable array field", schemapb.DataType_Array, nil, true},
|
|
{"nullable string/varchar field", schemapb.DataType_String, nil, true},
|
|
{"nullable binary vector field", schemapb.DataType_BinaryVector, []*commonpb.KeyValuePair{{Key: "dim", Value: "8"}}, true},
|
|
{"nullable float vector field", schemapb.DataType_FloatVector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, true},
|
|
{"nullable float16 vector field", schemapb.DataType_Float16Vector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, true},
|
|
{"nullable bfloat16 vector field", schemapb.DataType_BFloat16Vector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, true},
|
|
{"nullable sparse float vector field", schemapb.DataType_SparseFloatVector, nil, true},
|
|
{"nullable int8 vector field", schemapb.DataType_Int8Vector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, true},
|
|
{"non-nullable binary vector field", schemapb.DataType_BinaryVector, []*commonpb.KeyValuePair{{Key: "dim", Value: "8"}}, false},
|
|
{"non-nullable float vector field", schemapb.DataType_FloatVector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, false},
|
|
{"non-nullable float16 vector field", schemapb.DataType_Float16Vector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, false},
|
|
{"non-nullable bfloat16 vector field", schemapb.DataType_BFloat16Vector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, false},
|
|
{"non-nullable sparse float vector field", schemapb.DataType_SparseFloatVector, nil, false},
|
|
{"non-nullable int8 vector field", schemapb.DataType_Int8Vector, []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}, false},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
DataType: test.dataType,
|
|
Nullable: test.nullable,
|
|
TypeParams: test.typeParams,
|
|
},
|
|
},
|
|
}
|
|
_, err := NewInsertData(schema)
|
|
s.Nil(err)
|
|
})
|
|
}
|
|
})
|
|
|
|
s.Run("invalid schema", func() {
|
|
tests := []struct {
|
|
description string
|
|
invalidType schemapb.DataType
|
|
}{
|
|
{"binary vector without dim", schemapb.DataType_BinaryVector},
|
|
{"float vector without dim", schemapb.DataType_FloatVector},
|
|
{"float16 vector without dim", schemapb.DataType_Float16Vector},
|
|
{"bfloat16 vector without dim", schemapb.DataType_BFloat16Vector},
|
|
{"int8 vector without dim", schemapb.DataType_Int8Vector},
|
|
}
|
|
|
|
for _, test := range tests {
|
|
s.Run(test.description, func() {
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
DataType: test.invalidType,
|
|
},
|
|
},
|
|
}
|
|
idata, err := NewInsertData(schema)
|
|
s.Error(err)
|
|
s.Nil(idata)
|
|
})
|
|
}
|
|
})
|
|
|
|
s.Run("empty iData", func() {
|
|
idata := &InsertData{}
|
|
s.True(idata.IsEmpty())
|
|
s.Equal(0, idata.GetRowNum())
|
|
s.Equal(0, idata.GetMemorySize())
|
|
|
|
err := idata.Append(map[FieldID]interface{}{1: struct{}{}})
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("init by New", func() {
|
|
s.True(s.iDataEmpty.IsEmpty())
|
|
s.Equal(0, s.iDataEmpty.GetRowNum())
|
|
s.Equal(161+1, s.iDataEmpty.GetMemorySize())
|
|
|
|
s.False(s.iDataOneRow.IsEmpty())
|
|
s.Equal(1, s.iDataOneRow.GetRowNum())
|
|
s.Equal(535+1, s.iDataOneRow.GetMemorySize())
|
|
|
|
s.False(s.iDataTwoRows.IsEmpty())
|
|
s.Equal(2, s.iDataTwoRows.GetRowNum())
|
|
s.Equal(734+1, s.iDataTwoRows.GetMemorySize())
|
|
|
|
for _, field := range s.iDataTwoRows.Data {
|
|
s.Equal(2, field.RowNum())
|
|
|
|
err := field.AppendRow(struct{}{})
|
|
mlog.Warn(context.TODO(), "error", mlog.Err(err))
|
|
s.ErrorIs(err, merr.ErrParameterInvalid)
|
|
}
|
|
})
|
|
}
|
|
|
|
func (s *InsertDataSuite) TestMemorySize() {
|
|
s.Equal(s.iDataEmpty.Data[RowIDField].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[TimestampField].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[BoolField].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[Int8Field].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[Int16Field].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[Int32Field].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[Int64Field].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[FloatField].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[DoubleField].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[StringField].GetMemorySize(), 1)
|
|
s.Equal(s.iDataEmpty.Data[ArrayField].GetMemorySize(), 1)
|
|
// +9 bytes: Nullable(1) + L2PMapping.GetMemorySize()(8)
|
|
s.Equal(s.iDataEmpty.Data[BinaryVectorField].GetMemorySize(), 4+9)
|
|
s.Equal(s.iDataEmpty.Data[FloatVectorField].GetMemorySize(), 4+9)
|
|
s.Equal(s.iDataEmpty.Data[Float16VectorField].GetMemorySize(), 4+9)
|
|
s.Equal(s.iDataEmpty.Data[BFloat16VectorField].GetMemorySize(), 4+9)
|
|
s.Equal(s.iDataEmpty.Data[SparseFloatVectorField].GetMemorySize(), 0+9)
|
|
s.Equal(s.iDataEmpty.Data[Int8VectorField].GetMemorySize(), 4+9)
|
|
s.Equal(s.iDataEmpty.Data[StructSubInt32Field].GetMemorySize(), 1)
|
|
// +1 byte: Nullable flag (VectorArrayFieldData has no L2PMapping under Plan B)
|
|
s.Equal(s.iDataEmpty.Data[StructSubFloatVectorField].GetMemorySize(), 0+1)
|
|
|
|
s.Equal(s.iDataOneRow.Data[RowIDField].GetMemorySize(), 9)
|
|
s.Equal(s.iDataOneRow.Data[TimestampField].GetMemorySize(), 9)
|
|
s.Equal(s.iDataOneRow.Data[BoolField].GetMemorySize(), 2)
|
|
s.Equal(s.iDataOneRow.Data[Int8Field].GetMemorySize(), 2)
|
|
s.Equal(s.iDataOneRow.Data[Int16Field].GetMemorySize(), 3)
|
|
s.Equal(s.iDataOneRow.Data[Int32Field].GetMemorySize(), 5)
|
|
s.Equal(s.iDataOneRow.Data[Int64Field].GetMemorySize(), 9)
|
|
s.Equal(s.iDataOneRow.Data[FloatField].GetMemorySize(), 5)
|
|
s.Equal(s.iDataOneRow.Data[DoubleField].GetMemorySize(), 9)
|
|
s.Equal(s.iDataOneRow.Data[StringField].GetMemorySize(), 20)
|
|
s.Equal(s.iDataOneRow.Data[JSONField].GetMemorySize(), len([]byte(`{"batch":1}`))+16+1)
|
|
s.Equal(s.iDataOneRow.Data[ArrayField].GetMemorySize(), 3*4+1)
|
|
// +9 bytes: Nullable(1) + L2PMapping.GetMemorySize()(8)
|
|
s.Equal(s.iDataOneRow.Data[BinaryVectorField].GetMemorySize(), 5+9)
|
|
s.Equal(s.iDataOneRow.Data[FloatVectorField].GetMemorySize(), 20+9)
|
|
s.Equal(s.iDataOneRow.Data[Float16VectorField].GetMemorySize(), 12+9)
|
|
s.Equal(s.iDataOneRow.Data[BFloat16VectorField].GetMemorySize(), 12+9)
|
|
s.Equal(s.iDataOneRow.Data[SparseFloatVectorField].GetMemorySize(), 28+9)
|
|
s.Equal(s.iDataOneRow.Data[Int8VectorField].GetMemorySize(), 8+9)
|
|
s.Equal(s.iDataOneRow.Data[StructSubInt32Field].GetMemorySize(), 3*4+1)
|
|
s.Equal(s.iDataOneRow.Data[StructSubFloatVectorField].GetMemorySize(), 3*4*2+4+1)
|
|
|
|
s.Equal(s.iDataTwoRows.Data[RowIDField].GetMemorySize(), 17)
|
|
s.Equal(s.iDataTwoRows.Data[TimestampField].GetMemorySize(), 17)
|
|
s.Equal(s.iDataTwoRows.Data[BoolField].GetMemorySize(), 3)
|
|
s.Equal(s.iDataTwoRows.Data[Int8Field].GetMemorySize(), 3)
|
|
s.Equal(s.iDataTwoRows.Data[Int16Field].GetMemorySize(), 5)
|
|
s.Equal(s.iDataTwoRows.Data[Int32Field].GetMemorySize(), 9)
|
|
s.Equal(s.iDataTwoRows.Data[Int64Field].GetMemorySize(), 17)
|
|
s.Equal(s.iDataTwoRows.Data[FloatField].GetMemorySize(), 9)
|
|
s.Equal(s.iDataTwoRows.Data[DoubleField].GetMemorySize(), 17)
|
|
s.Equal(s.iDataTwoRows.Data[StringField].GetMemorySize(), 39)
|
|
s.Equal(s.iDataTwoRows.Data[ArrayField].GetMemorySize(), 25)
|
|
// +9 bytes: Nullable(1) + L2PMapping.GetMemorySize()(8)
|
|
s.Equal(s.iDataTwoRows.Data[BinaryVectorField].GetMemorySize(), 6+9)
|
|
s.Equal(s.iDataTwoRows.Data[FloatVectorField].GetMemorySize(), 36+9)
|
|
s.Equal(s.iDataTwoRows.Data[Float16VectorField].GetMemorySize(), 20+9)
|
|
s.Equal(s.iDataTwoRows.Data[BFloat16VectorField].GetMemorySize(), 20+9)
|
|
s.Equal(s.iDataTwoRows.Data[SparseFloatVectorField].GetMemorySize(), 54+9)
|
|
s.Equal(s.iDataTwoRows.Data[Int8VectorField].GetMemorySize(), 12+9)
|
|
s.Equal(s.iDataTwoRows.Data[StructSubInt32Field].GetMemorySize(), 3*4+2*4+1)
|
|
s.Equal(s.iDataTwoRows.Data[StructSubFloatVectorField].GetMemorySize(), 3*4*2+4+2*4*2+4+1)
|
|
}
|
|
|
|
func (s *InsertDataSuite) TestGetRowSize() {
|
|
s.Equal(s.iDataOneRow.Data[RowIDField].GetRowSize(0), 8)
|
|
s.Equal(s.iDataOneRow.Data[TimestampField].GetRowSize(0), 8)
|
|
s.Equal(s.iDataOneRow.Data[BoolField].GetRowSize(0), 1)
|
|
s.Equal(s.iDataOneRow.Data[Int8Field].GetRowSize(0), 1)
|
|
s.Equal(s.iDataOneRow.Data[Int16Field].GetRowSize(0), 2)
|
|
s.Equal(s.iDataOneRow.Data[Int32Field].GetRowSize(0), 4)
|
|
s.Equal(s.iDataOneRow.Data[Int64Field].GetRowSize(0), 8)
|
|
s.Equal(s.iDataOneRow.Data[FloatField].GetRowSize(0), 4)
|
|
s.Equal(s.iDataOneRow.Data[DoubleField].GetRowSize(0), 8)
|
|
s.Equal(s.iDataOneRow.Data[StringField].GetRowSize(0), 19)
|
|
s.Equal(s.iDataOneRow.Data[JSONField].GetRowSize(0), len([]byte(`{"batch":1}`))+16)
|
|
s.Equal(s.iDataOneRow.Data[ArrayField].GetRowSize(0), 3*4)
|
|
s.Equal(s.iDataOneRow.Data[BinaryVectorField].GetRowSize(0), 1)
|
|
s.Equal(s.iDataOneRow.Data[FloatVectorField].GetRowSize(0), 16)
|
|
s.Equal(s.iDataOneRow.Data[Float16VectorField].GetRowSize(0), 8)
|
|
s.Equal(s.iDataOneRow.Data[BFloat16VectorField].GetRowSize(0), 8)
|
|
s.Equal(s.iDataOneRow.Data[SparseFloatVectorField].GetRowSize(0), 24)
|
|
s.Equal(s.iDataOneRow.Data[Int8VectorField].GetRowSize(0), 4)
|
|
s.Equal(s.iDataOneRow.Data[StructSubInt32Field].GetRowSize(0), 3*4)
|
|
s.Equal(s.iDataOneRow.Data[StructSubFloatVectorField].GetRowSize(0), 3*4*2+4)
|
|
}
|
|
|
|
func GetFields(schema *schemapb.CollectionSchema) []*schemapb.FieldSchema {
|
|
ret := make([]*schemapb.FieldSchema, 0, 100)
|
|
ret = append(ret, schema.GetFields()...)
|
|
for _, structField := range schema.GetStructArrayFields() {
|
|
ret = append(ret, structField.GetFields()...)
|
|
}
|
|
return ret
|
|
}
|
|
|
|
func (s *InsertDataSuite) TestGetDataType() {
|
|
for _, field := range GetFields(s.schema) {
|
|
fieldData, ok := s.iDataOneRow.Data[field.GetFieldID()]
|
|
s.True(ok)
|
|
s.Equal(field.GetDataType(), fieldData.GetDataType())
|
|
}
|
|
}
|
|
|
|
func (s *InsertDataSuite) TestGetNullable() {
|
|
for _, field := range GetFields(s.schema) {
|
|
fieldData, ok := s.iDataOneRow.Data[field.GetFieldID()]
|
|
s.True(ok)
|
|
s.Equal(field.GetNullable(), fieldData.GetNullable())
|
|
}
|
|
}
|
|
|
|
func (s *InsertDataSuite) SetupTest() {
|
|
var err error
|
|
s.iDataEmpty, err = NewInsertData(s.schema)
|
|
s.Require().NoError(err)
|
|
s.True(s.iDataEmpty.IsEmpty())
|
|
s.Equal(0, s.iDataEmpty.GetRowNum())
|
|
s.Equal(161+1, s.iDataEmpty.GetMemorySize())
|
|
|
|
row1 := map[FieldID]interface{}{
|
|
RowIDField: int64(3),
|
|
TimestampField: int64(3),
|
|
BoolField: true,
|
|
Int8Field: int8(3),
|
|
Int16Field: int16(3),
|
|
Int32Field: int32(3),
|
|
Int64Field: int64(3),
|
|
FloatField: float32(3),
|
|
DoubleField: float64(3),
|
|
StringField: "str",
|
|
BinaryVectorField: []byte{0},
|
|
FloatVectorField: []float32{4, 5, 6, 7},
|
|
Float16VectorField: []byte{0, 0, 0, 0, 255, 255, 255, 255},
|
|
BFloat16VectorField: []byte{0, 0, 0, 0, 255, 255, 255, 255},
|
|
SparseFloatVectorField: typeutil.CreateSparseFloatRow([]uint32{0, 1, 2}, []float32{4, 5, 6}),
|
|
Int8VectorField: []int8{-4, -5, 6, 7},
|
|
NullableFloatVectorField: []float32{1.0, 2.0, 3.0, 4.0},
|
|
NullableBinaryVectorField: []byte{1},
|
|
NullableFloat16VectorField: []byte{1, 2, 3, 4, 5, 6, 7, 8},
|
|
NullableBFloat16VectorField: []byte{1, 2, 3, 4, 5, 6, 7, 8},
|
|
NullableInt8VectorField: []int8{1, 2, 3, 4},
|
|
NullableSparseFloatVectorField: typeutil.CreateSparseFloatRow([]uint32{0, 1, 2}, []float32{4, 5, 6}),
|
|
ArrayField: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
|
},
|
|
},
|
|
JSONField: []byte(`{"batch":3}`),
|
|
StructSubInt32Field: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
|
},
|
|
},
|
|
StructSubFloatVectorField: &schemapb.VectorField{
|
|
Dim: 2,
|
|
Data: &schemapb.VectorField_FloatVector{
|
|
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6}},
|
|
},
|
|
},
|
|
}
|
|
|
|
s.iDataOneRow, err = NewInsertData(s.schema)
|
|
s.Require().NoError(err)
|
|
err = s.iDataOneRow.Append(row1)
|
|
s.Require().NoError(err)
|
|
|
|
for fID, field := range s.iDataOneRow.Data {
|
|
s.Equal(row1[fID], field.GetRow(0))
|
|
}
|
|
|
|
row2 := map[FieldID]interface{}{
|
|
RowIDField: int64(1),
|
|
TimestampField: int64(1),
|
|
BoolField: false,
|
|
Int8Field: int8(1),
|
|
Int16Field: int16(1),
|
|
Int32Field: int32(1),
|
|
Int64Field: int64(1),
|
|
FloatField: float32(1),
|
|
DoubleField: float64(1),
|
|
StringField: string("str"),
|
|
BinaryVectorField: []byte{0},
|
|
FloatVectorField: []float32{4, 5, 6, 7},
|
|
Float16VectorField: []byte{1, 2, 3, 4, 5, 6, 7, 8},
|
|
BFloat16VectorField: []byte{1, 2, 3, 4, 5, 6, 7, 8},
|
|
SparseFloatVectorField: typeutil.CreateSparseFloatRow([]uint32{2, 3, 4}, []float32{4, 5, 6}),
|
|
Int8VectorField: []int8{-128, -5, 6, 127},
|
|
NullableFloatVectorField: nil,
|
|
NullableBinaryVectorField: nil,
|
|
NullableFloat16VectorField: nil,
|
|
NullableBFloat16VectorField: nil,
|
|
NullableInt8VectorField: nil,
|
|
NullableSparseFloatVectorField: nil,
|
|
ArrayField: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{1, 2, 3}},
|
|
},
|
|
},
|
|
JSONField: []byte(`{"batch":1}`),
|
|
StructSubInt32Field: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{1, 2}},
|
|
},
|
|
},
|
|
StructSubFloatVectorField: &schemapb.VectorField{
|
|
Dim: 2,
|
|
Data: &schemapb.VectorField_FloatVector{
|
|
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}},
|
|
},
|
|
},
|
|
}
|
|
|
|
s.iDataTwoRows, err = NewInsertData(s.schema)
|
|
s.Require().NoError(err)
|
|
err = s.iDataTwoRows.Append(row1)
|
|
s.Require().NoError(err)
|
|
err = s.iDataTwoRows.Append(row2)
|
|
s.Require().NoError(err)
|
|
}
|
|
|
|
type ArrayFieldDataSuite struct {
|
|
suite.Suite
|
|
}
|
|
|
|
func (s *ArrayFieldDataSuite) TestArrayFieldData() {
|
|
fieldID2Type := map[int64]schemapb.DataType{
|
|
ArrayField + 1: schemapb.DataType_Bool,
|
|
ArrayField + 2: schemapb.DataType_Int8,
|
|
ArrayField + 3: schemapb.DataType_Int16,
|
|
ArrayField + 4: schemapb.DataType_Int32,
|
|
ArrayField + 5: schemapb.DataType_Int64,
|
|
ArrayField + 6: schemapb.DataType_Float,
|
|
ArrayField + 7: schemapb.DataType_Double,
|
|
ArrayField + 8: schemapb.DataType_VarChar,
|
|
}
|
|
|
|
schema := &schemapb.CollectionSchema{
|
|
Fields: []*schemapb.FieldSchema{
|
|
{
|
|
FieldID: RowIDField,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: TimestampField,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
{
|
|
FieldID: Int64Field,
|
|
IsPrimaryKey: true,
|
|
DataType: schemapb.DataType_Int64,
|
|
},
|
|
},
|
|
}
|
|
|
|
for fieldID, elementType := range fieldID2Type {
|
|
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
|
|
FieldID: fieldID,
|
|
DataType: schemapb.DataType_Array,
|
|
ElementType: elementType,
|
|
})
|
|
}
|
|
|
|
insertData, err := NewInsertData(schema)
|
|
s.NoError(err)
|
|
|
|
s.Equal(0, insertData.GetRowNum())
|
|
s.Equal(11, insertData.GetMemorySize())
|
|
s.True(insertData.IsEmpty())
|
|
|
|
fieldIDToData := map[int64]interface{}{
|
|
RowIDField: int64(1),
|
|
TimestampField: int64(2),
|
|
Int64Field: int64(3),
|
|
ArrayField + 1: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_BoolData{
|
|
BoolData: &schemapb.BoolArray{Data: []bool{true, false}},
|
|
},
|
|
},
|
|
ArrayField + 2: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{0, 0}},
|
|
},
|
|
},
|
|
ArrayField + 3: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{1, 1}},
|
|
},
|
|
},
|
|
ArrayField + 4: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_IntData{
|
|
IntData: &schemapb.IntArray{Data: []int32{2, 2}},
|
|
},
|
|
},
|
|
ArrayField + 5: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_LongData{
|
|
LongData: &schemapb.LongArray{Data: []int64{3, 3}},
|
|
},
|
|
},
|
|
ArrayField + 6: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_FloatData{
|
|
FloatData: &schemapb.FloatArray{Data: []float32{4, 4}},
|
|
},
|
|
},
|
|
ArrayField + 7: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_DoubleData{
|
|
DoubleData: &schemapb.DoubleArray{Data: []float64{5, 5}},
|
|
},
|
|
},
|
|
ArrayField + 8: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_StringData{
|
|
StringData: &schemapb.StringArray{Data: []string{"6", "6"}},
|
|
},
|
|
},
|
|
}
|
|
|
|
err = insertData.Append(fieldIDToData)
|
|
s.NoError(err)
|
|
s.Equal(1, insertData.GetRowNum())
|
|
s.Equal(126, insertData.GetMemorySize())
|
|
s.False(insertData.IsEmpty())
|
|
s.Equal(115, insertData.GetRowSize(0))
|
|
}
|
|
|
|
func makeFloatVec(dim int, vals ...float32) *schemapb.VectorField {
|
|
return &schemapb.VectorField{
|
|
Dim: int64(dim),
|
|
Data: &schemapb.VectorField_FloatVector{
|
|
FloatVector: &schemapb.FloatArray{Data: vals},
|
|
},
|
|
}
|
|
}
|
|
|
|
func TestNullableVectorAppendRowsRejectsNonCompactData(t *testing.T) {
|
|
validData := []bool{true, false, true}
|
|
sparseRows := &SparseFloatVectorFieldData{
|
|
SparseFloatArray: schemapb.SparseFloatArray{
|
|
Dim: 8,
|
|
Contents: [][]byte{
|
|
typeutil.CreateSparseFloatRow([]uint32{1}, []float32{1.0}),
|
|
typeutil.CreateSparseFloatRow([]uint32{2}, []float32{2.0}),
|
|
typeutil.CreateSparseFloatRow([]uint32{3}, []float32{3.0}),
|
|
},
|
|
},
|
|
}
|
|
|
|
tests := []struct {
|
|
name string
|
|
fieldData FieldData
|
|
rows any
|
|
}{
|
|
{
|
|
name: "binary vector",
|
|
fieldData: &BinaryVectorFieldData{Dim: 8, Nullable: true},
|
|
rows: []byte{0x01, 0x02, 0x03},
|
|
},
|
|
{
|
|
name: "float vector",
|
|
fieldData: &FloatVectorFieldData{Dim: 2, Nullable: true},
|
|
rows: []float32{1, 2, 3, 4, 5, 6},
|
|
},
|
|
{
|
|
name: "float16 vector",
|
|
fieldData: &Float16VectorFieldData{Dim: 2, Nullable: true},
|
|
rows: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
|
|
},
|
|
{
|
|
name: "bfloat16 vector",
|
|
fieldData: &BFloat16VectorFieldData{Dim: 2, Nullable: true},
|
|
rows: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
|
|
},
|
|
{
|
|
name: "sparse vector",
|
|
fieldData: &SparseFloatVectorFieldData{Nullable: true},
|
|
rows: sparseRows,
|
|
},
|
|
{
|
|
name: "int8 vector",
|
|
fieldData: &Int8VectorFieldData{Dim: 2, Nullable: true},
|
|
rows: []int8{1, 2, 3, 4, 5, 6},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
err := tt.fieldData.AppendRows(tt.rows, validData)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "compact")
|
|
assert.Equal(t, 0, tt.fieldData.RowNum())
|
|
assert.Empty(t, nullableVectorValidDataForTest(tt.fieldData))
|
|
})
|
|
}
|
|
}
|
|
|
|
func nullableVectorValidDataForTest(fieldData FieldData) []bool {
|
|
switch data := fieldData.(type) {
|
|
case *BinaryVectorFieldData:
|
|
return data.ValidData
|
|
case *FloatVectorFieldData:
|
|
return data.ValidData
|
|
case *Float16VectorFieldData:
|
|
return data.ValidData
|
|
case *BFloat16VectorFieldData:
|
|
return data.ValidData
|
|
case *SparseFloatVectorFieldData:
|
|
return data.ValidData
|
|
case *Int8VectorFieldData:
|
|
return data.ValidData
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func TestVectorArrayFieldData_NullableAppendAndGetRow(t *testing.T) {
|
|
fd := &VectorArrayFieldData{
|
|
Dim: 4,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Data: make([]*schemapb.VectorField, 0),
|
|
ValidData: make([]bool, 0),
|
|
Nullable: true,
|
|
}
|
|
|
|
vec0 := makeFloatVec(4, 1, 2, 3, 4)
|
|
vec2 := makeFloatVec(4, 5, 6, 7, 8)
|
|
|
|
require.NoError(t, fd.AppendRow(vec0))
|
|
require.NoError(t, fd.AppendRow(nil))
|
|
require.NoError(t, fd.AppendRow(vec2))
|
|
|
|
assert.Equal(t, 3, fd.RowNum(), "RowNum should count all rows including null placeholders")
|
|
assert.Equal(t, 3, len(fd.Data), "Plan B: Data is dense; null rows hold empty placeholders")
|
|
assert.Equal(t, []bool{true, false, true}, fd.GetValidData())
|
|
assert.True(t, fd.GetNullable())
|
|
|
|
assert.Equal(t, vec0, fd.GetRow(0))
|
|
assert.Nil(t, fd.GetRow(1), "null row should return nil despite placeholder in Data")
|
|
assert.Equal(t, vec2, fd.GetRow(2))
|
|
}
|
|
|
|
func TestVectorArrayFieldData_NonNullable(t *testing.T) {
|
|
fd := &VectorArrayFieldData{
|
|
Dim: 4,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Data: make([]*schemapb.VectorField, 0),
|
|
Nullable: false,
|
|
}
|
|
|
|
vec := makeFloatVec(4, 1, 2, 3, 4)
|
|
require.NoError(t, fd.AppendRow(vec))
|
|
|
|
assert.Equal(t, 1, fd.RowNum())
|
|
assert.False(t, fd.GetNullable())
|
|
assert.Nil(t, fd.GetValidData())
|
|
assert.Equal(t, vec, fd.GetRow(0))
|
|
}
|
|
|
|
func TestVectorArrayFieldData_AppendValidDataRows(t *testing.T) {
|
|
fd := &VectorArrayFieldData{
|
|
Dim: 4,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Data: []*schemapb.VectorField{
|
|
makeFloatVec(4, 1, 2, 3, 4),
|
|
makeFloatVec(4, 5, 6, 7, 8),
|
|
},
|
|
Nullable: true,
|
|
}
|
|
|
|
err := fd.AppendValidDataRows([]bool{true, false, true})
|
|
require.NoError(t, err)
|
|
|
|
assert.Equal(t, []bool{true, false, true}, fd.GetValidData())
|
|
|
|
assert.NoError(t, fd.AppendValidDataRows(nil))
|
|
assert.Error(t, fd.AppendValidDataRows("bad"))
|
|
}
|
|
|
|
func TestVectorArrayFieldData_GetMemorySize(t *testing.T) {
|
|
fd := &VectorArrayFieldData{
|
|
Dim: 4,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Data: []*schemapb.VectorField{makeFloatVec(4, 1, 2, 3, 4)},
|
|
ValidData: []bool{true, false},
|
|
Nullable: true,
|
|
}
|
|
|
|
assert.Greater(t, fd.GetMemorySize(), 0)
|
|
}
|
|
|
|
func TestVectorArrayFieldData_AllNull(t *testing.T) {
|
|
fd := &VectorArrayFieldData{
|
|
Dim: 4,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Data: make([]*schemapb.VectorField, 0),
|
|
ValidData: make([]bool, 0),
|
|
Nullable: true,
|
|
}
|
|
|
|
for i := 0; i < 3; i++ {
|
|
require.NoError(t, fd.AppendRow(nil))
|
|
}
|
|
|
|
assert.Equal(t, 3, fd.RowNum())
|
|
assert.Equal(t, 3, len(fd.Data), "Plan B: Data is dense, null rows hold placeholders")
|
|
for i := 0; i < 3; i++ {
|
|
assert.Nil(t, fd.GetRow(i))
|
|
}
|
|
}
|
|
|
|
func TestNewFieldData_NullableArrayOfVector(t *testing.T) {
|
|
schema := &schemapb.FieldSchema{
|
|
FieldID: 100,
|
|
Name: "vec_arr",
|
|
DataType: schemapb.DataType_ArrayOfVector,
|
|
ElementType: schemapb.DataType_FloatVector,
|
|
Nullable: true,
|
|
TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}},
|
|
}
|
|
fd, err := NewFieldData(schemapb.DataType_ArrayOfVector, schema, 10)
|
|
require.NoError(t, err)
|
|
|
|
vafd, ok := fd.(*VectorArrayFieldData)
|
|
require.True(t, ok)
|
|
assert.True(t, vafd.Nullable)
|
|
assert.NotNil(t, vafd.ValidData)
|
|
assert.Equal(t, int64(4), vafd.Dim)
|
|
}
|