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) }