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 Cpp (push) Waiting to run
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
320 lines
10 KiB
Go
320 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 column
|
|
|
|
import (
|
|
"fmt"
|
|
"math/rand"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/client/v2/entity"
|
|
)
|
|
|
|
type StructArraySuite struct {
|
|
suite.Suite
|
|
}
|
|
|
|
func (s *StructArraySuite) TestBasic() {
|
|
name := fmt.Sprintf("struct_array_%d", rand.Intn(100))
|
|
|
|
// Each row holds a variable-length array of struct elements; sub-columns are *Array types.
|
|
intRows := [][]int32{{1, 2}, {3}, {4, 5, 6}}
|
|
floatRows := [][]float32{{1.1, 2.2}, {3.3}, {4.4, 5.5, 6.6}}
|
|
strRows := [][]string{{"a", "b"}, {"c"}, {"d", "e", "f"}}
|
|
|
|
int32Col := NewColumnInt32Array("int_field", intRows)
|
|
floatCol := NewColumnFloatArray("float_field", floatRows)
|
|
varcharCol := NewColumnVarCharArray("varchar_field", strRows)
|
|
|
|
column := NewColumnStructArray(name, []Column{int32Col, floatCol, varcharCol})
|
|
|
|
s.Equal(name, column.Name())
|
|
s.Equal(entity.FieldTypeArray, column.Type())
|
|
s.EqualValues(3, column.Len())
|
|
|
|
fd := column.FieldData()
|
|
s.Equal(schemapb.DataType_ArrayOfStruct, fd.GetType())
|
|
s.Equal(name, fd.GetFieldName())
|
|
|
|
structArrays := fd.GetStructArrays()
|
|
s.NotNil(structArrays)
|
|
s.Equal(3, len(structArrays.GetFields()))
|
|
|
|
// Sub-fields must be Array (not flat scalars) to match server-side schema for struct sub-fields.
|
|
for _, sub := range structArrays.GetFields() {
|
|
s.Equal(schemapb.DataType_Array, sub.GetType())
|
|
}
|
|
|
|
val, err := column.Get(0)
|
|
s.NoError(err)
|
|
m, ok := val.(map[string]any)
|
|
s.True(ok)
|
|
s.Equal([]int32{1, 2}, m["int_field"])
|
|
s.Equal([]float32{1.1, 2.2}, m["float_field"])
|
|
s.Equal([]string{"a", "b"}, m["varchar_field"])
|
|
|
|
val, err = column.Get(2)
|
|
s.NoError(err)
|
|
m = val.(map[string]any)
|
|
s.Equal([]int32{4, 5, 6}, m["int_field"])
|
|
}
|
|
|
|
func (s *StructArraySuite) TestVectorSubField() {
|
|
dim := 4
|
|
rows := [][][]float32{
|
|
{{0.1, 0.2, 0.3, 0.4}, {0.5, 0.6, 0.7, 0.8}},
|
|
{{1.1, 1.2, 1.3, 1.4}},
|
|
}
|
|
idRows := [][]int64{{10, 20}, {30}}
|
|
|
|
idCol := NewColumnInt64Array("id", idRows)
|
|
embCol := NewColumnFloatVectorArray("emb", dim, rows)
|
|
column := NewColumnStructArray("clips", []Column{idCol, embCol})
|
|
|
|
fd := column.FieldData()
|
|
s.Equal(schemapb.DataType_ArrayOfStruct, fd.GetType())
|
|
s.Equal(2, len(fd.GetStructArrays().GetFields()))
|
|
|
|
embFD := fd.GetStructArrays().GetFields()[1]
|
|
s.Equal(schemapb.DataType_ArrayOfVector, embFD.GetType())
|
|
va := embFD.GetVectors().GetVectorArray()
|
|
s.NotNil(va)
|
|
s.EqualValues(dim, va.GetDim())
|
|
s.Equal(schemapb.DataType_FloatVector, va.GetElementType())
|
|
s.Equal(2, len(va.GetData()))
|
|
s.EqualValues(2*dim, len(va.GetData()[0].GetFloatVector().GetData()))
|
|
s.EqualValues(1*dim, len(va.GetData()[1].GetFloatVector().GetData()))
|
|
}
|
|
|
|
func (s *StructArraySuite) TestSlice() {
|
|
intRows := [][]int64{{10}, {20, 21}, {30, 31, 32}, {40}, {50, 51}}
|
|
boolRows := [][]bool{{true}, {false, true}, {true, false, true}, {false}, {true, false}}
|
|
|
|
int64Col := NewColumnInt64Array("id", intRows)
|
|
boolCol := NewColumnBoolArray("flag", boolRows)
|
|
column := NewColumnStructArray("struct_array_slice", []Column{int64Col, boolCol})
|
|
|
|
sliced := column.Slice(1, 4)
|
|
s.NotNil(sliced)
|
|
s.EqualValues(3, sliced.Len())
|
|
|
|
val, err := sliced.Get(0)
|
|
s.NoError(err)
|
|
m := val.(map[string]any)
|
|
s.Equal([]int64{20, 21}, m["id"])
|
|
s.Equal([]bool{false, true}, m["flag"])
|
|
}
|
|
|
|
func (s *StructArraySuite) TestAppendValue() {
|
|
intCol := NewColumnInt32Array("a", nil)
|
|
strCol := NewColumnVarCharArray("b", nil)
|
|
column := NewColumnStructArray("rows", []Column{intCol, strCol})
|
|
|
|
s.NoError(column.AppendValue(map[string]any{"a": []int32{1, 2}, "b": []string{"x", "y"}}))
|
|
s.NoError(column.AppendValue(map[string]any{"a": []int32{3}, "b": []string{"z"}}))
|
|
s.EqualValues(2, column.Len())
|
|
|
|
// missing sub-field
|
|
err := column.AppendValue(map[string]any{"a": []int32{4}})
|
|
s.Error(err)
|
|
|
|
// wrong shape (scalar instead of array)
|
|
err = column.AppendValue(map[string]any{"a": int32(1), "b": []string{"q"}})
|
|
s.Error(err)
|
|
}
|
|
|
|
func (s *StructArraySuite) TestParseStructArrayData() {
|
|
int32FieldData := &schemapb.FieldData{
|
|
Type: schemapb.DataType_Array,
|
|
FieldName: "age",
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_ArrayData{
|
|
ArrayData: &schemapb.ArrayArray{
|
|
ElementType: schemapb.DataType_Int32,
|
|
Data: []*schemapb.ScalarField{
|
|
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{10, 11}}}},
|
|
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{20}}}},
|
|
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{30, 31, 32}}}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
varcharFieldData := &schemapb.FieldData{
|
|
Type: schemapb.DataType_Array,
|
|
FieldName: "name",
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_ArrayData{
|
|
ArrayData: &schemapb.ArrayArray{
|
|
ElementType: schemapb.DataType_VarChar,
|
|
Data: []*schemapb.ScalarField{
|
|
{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"alice", "ann"}}}},
|
|
{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"bob"}}}},
|
|
{Data: &schemapb.ScalarField_StringData{StringData: &schemapb.StringArray{Data: []string{"c1", "c2", "c3"}}}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
structArrayField := &schemapb.StructArrayField{
|
|
Fields: []*schemapb.FieldData{int32FieldData, varcharFieldData},
|
|
}
|
|
|
|
col, err := parseStructArrayData("person", structArrayField, 0, -1)
|
|
s.NoError(err)
|
|
s.NotNil(col)
|
|
s.Equal("person", col.Name())
|
|
s.Equal(entity.FieldTypeArray, col.Type())
|
|
|
|
val, err := col.Get(0)
|
|
s.NoError(err)
|
|
m := val.(map[string]any)
|
|
s.Equal([]int32{10, 11}, m["age"])
|
|
s.Equal([]string{"alice", "ann"}, m["name"])
|
|
|
|
val, err = col.Get(1)
|
|
s.NoError(err)
|
|
m = val.(map[string]any)
|
|
s.Equal([]int32{20}, m["age"])
|
|
s.Equal([]string{"bob"}, m["name"])
|
|
}
|
|
|
|
func (s *StructArraySuite) TestParseTopLevelArrayOfStruct() {
|
|
// Verify FieldDataColumn dispatches DataType_ArrayOfStruct to parseStructArrayData.
|
|
int32Sub := &schemapb.FieldData{
|
|
Type: schemapb.DataType_Array,
|
|
FieldName: "x",
|
|
Field: &schemapb.FieldData_Scalars{
|
|
Scalars: &schemapb.ScalarField{
|
|
Data: &schemapb.ScalarField_ArrayData{
|
|
ArrayData: &schemapb.ArrayArray{
|
|
ElementType: schemapb.DataType_Int32,
|
|
Data: []*schemapb.ScalarField{
|
|
{Data: &schemapb.ScalarField_IntData{IntData: &schemapb.IntArray{Data: []int32{1, 2}}}},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
top := &schemapb.FieldData{
|
|
Type: schemapb.DataType_ArrayOfStruct,
|
|
FieldName: "wrap",
|
|
Field: &schemapb.FieldData_StructArrays{
|
|
StructArrays: &schemapb.StructArrayField{Fields: []*schemapb.FieldData{int32Sub}},
|
|
},
|
|
}
|
|
col, err := FieldDataColumn(top, 0, -1)
|
|
s.NoError(err)
|
|
s.Equal("wrap", col.Name())
|
|
val, err := col.Get(0)
|
|
s.NoError(err)
|
|
s.Equal([]int32{1, 2}, val.(map[string]any)["x"])
|
|
}
|
|
|
|
func (s *StructArraySuite) TestParseVectorArrayDataErrors() {
|
|
mkFD := func(elemType schemapb.DataType, dim int64, rows []*schemapb.VectorField) *schemapb.FieldData {
|
|
return &schemapb.FieldData{
|
|
Type: schemapb.DataType_ArrayOfVector,
|
|
FieldName: "emb",
|
|
Field: &schemapb.FieldData_Vectors{
|
|
Vectors: &schemapb.VectorField{
|
|
Dim: dim,
|
|
Data: &schemapb.VectorField_VectorArray{
|
|
VectorArray: &schemapb.VectorArray{
|
|
Dim: dim, ElementType: elemType, Data: rows,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
s.Run("unknown dim rejected", func() {
|
|
fd := mkFD(schemapb.DataType_FloatVector, 0, nil)
|
|
_, err := FieldDataColumn(fd, 0, -1)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("payload not a multiple of dim", func() {
|
|
// dim=4 but row has 5 floats -> must error, not silently truncate.
|
|
row := &schemapb.VectorField{
|
|
Dim: 4,
|
|
Data: &schemapb.VectorField_FloatVector{FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5}}},
|
|
}
|
|
fd := mkFD(schemapb.DataType_FloatVector, 4, []*schemapb.VectorField{row})
|
|
_, err := FieldDataColumn(fd, 0, -1)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("binary dim not multiple of 8", func() {
|
|
row := &schemapb.VectorField{
|
|
Dim: 4,
|
|
Data: &schemapb.VectorField_BinaryVector{BinaryVector: []byte{0}},
|
|
}
|
|
fd := mkFD(schemapb.DataType_BinaryVector, 4, []*schemapb.VectorField{row})
|
|
_, err := FieldDataColumn(fd, 0, -1)
|
|
s.Error(err)
|
|
})
|
|
|
|
s.Run("nil row rejected", func() {
|
|
fd := mkFD(schemapb.DataType_FloatVector, 4, []*schemapb.VectorField{nil})
|
|
_, err := FieldDataColumn(fd, 0, -1)
|
|
s.Error(err)
|
|
})
|
|
}
|
|
|
|
func (s *StructArraySuite) TestAppendValueRollback() {
|
|
intCol := NewColumnInt32Array("a", nil)
|
|
strCol := NewColumnVarCharArray("b", nil)
|
|
col := NewColumnStructArray("rows", []Column{intCol, strCol}).(*columnStructArray)
|
|
|
|
// Seed with one good row so both sub-columns are at length 1.
|
|
s.NoError(col.AppendValue(map[string]any{"a": []int32{1}, "b": []string{"x"}}))
|
|
s.EqualValues(1, col.Len())
|
|
|
|
// Second row: sub-field "a" accepts the []int32, but "b" gets wrong type —
|
|
// rollback must restore sub-column "a" to length 1 so the struct stays in lock-step.
|
|
err := col.AppendValue(map[string]any{"a": []int32{2}, "b": 42})
|
|
s.Error(err)
|
|
s.EqualValues(1, col.fields[0].Len(), "sub-field 'a' must be rolled back")
|
|
s.EqualValues(1, col.fields[1].Len(), "sub-field 'b' must not have been appended")
|
|
s.EqualValues(1, col.Len(), "struct array length stays consistent after rollback")
|
|
}
|
|
|
|
func (s *StructArraySuite) TestLenMismatchPanics() {
|
|
// Manually drift sub-column lengths to simulate a prior corruption and verify Len reports it.
|
|
intCol := NewColumnInt32Array("a", [][]int32{{1}, {2}})
|
|
strCol := NewColumnVarCharArray("b", [][]string{{"x"}})
|
|
col := NewColumnStructArray("rows", []Column{intCol, strCol})
|
|
|
|
s.Panics(func() { _ = col.Len() })
|
|
}
|
|
|
|
func TestStructArray(t *testing.T) {
|
|
suite.Run(t, new(StructArraySuite))
|
|
}
|