Files
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 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
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

388 lines
11 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 (
"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 VectorArraySuite struct {
suite.Suite
}
func (s *VectorArraySuite) TestFloatVectorArrayBasic() {
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}},
}
col := NewColumnFloatVectorArray("emb", dim, rows)
s.Equal("emb", col.Name())
s.Equal(entity.FieldTypeArray, col.Type())
s.Equal(entity.FieldTypeFloatVector, col.ElementType())
s.Equal(dim, col.Dim())
s.Equal(2, col.Len())
s.Equal(2, col.ValidCount())
s.False(col.Nullable())
v, err := col.Get(0)
s.NoError(err)
row0, ok := v.([]entity.FloatVector)
s.Require().True(ok)
s.Equal(2, len(row0))
// Out-of-range Get.
_, err = col.Get(-1)
s.Error(err)
_, err = col.Get(99)
s.Error(err)
// Scalar conversions are unsupported on vector array columns.
_, err = col.GetAsInt64(0)
s.Error(err)
_, err = col.GetAsString(0)
s.Error(err)
_, err = col.GetAsDouble(0)
s.Error(err)
_, err = col.GetAsBool(0)
s.Error(err)
// Nullable helpers are no-ops / zero-valued for vector arrays.
isNull, err := col.IsNull(0)
s.NoError(err)
s.False(isNull)
s.Error(col.AppendNull())
col.SetNullable(true)
s.False(col.Nullable())
s.NoError(col.ValidateNullable())
col.CompactNullableValues()
// AppendValue via both canonical shapes.
s.NoError(col.AppendValue([]entity.FloatVector{entity.FloatVector([]float32{2.1, 2.2, 2.3, 2.4})}))
s.NoError(col.AppendValue([][]float32{{3.1, 3.2, 3.3, 3.4}}))
s.Equal(4, col.Len())
// AppendValue rejects bad shapes.
s.Error(col.AppendValue([]int{1, 2, 3}))
// FieldData round-trip.
fd := col.FieldData()
s.Equal(schemapb.DataType_ArrayOfVector, fd.GetType())
s.Equal("emb", fd.GetFieldName())
va := fd.GetVectors().GetVectorArray()
s.Require().NotNil(va)
s.EqualValues(dim, va.GetDim())
s.Equal(schemapb.DataType_FloatVector, va.GetElementType())
s.Equal(4, len(va.GetData()))
}
func (s *VectorArraySuite) TestFloat16VectorArrayBasic() {
dim := 4
byteRow := make([]byte, dim*2)
rows := [][][]byte{{byteRow, byteRow}, {byteRow}}
col := NewColumnFloat16VectorArray("emb", dim, rows)
s.Equal(entity.FieldTypeFloat16Vector, col.ElementType())
s.Equal(2, col.Len())
// AppendValue variants.
s.NoError(col.AppendValue([]entity.Float16Vector{entity.Float16Vector(byteRow)}))
s.NoError(col.AppendValue([][]byte{byteRow}))
s.Error(col.AppendValue(123))
s.Equal(4, col.Len())
fd := col.FieldData()
s.Equal(schemapb.DataType_ArrayOfVector, fd.GetType())
s.Equal(schemapb.DataType_Float16Vector, fd.GetVectors().GetVectorArray().GetElementType())
}
func (s *VectorArraySuite) TestBFloat16VectorArrayBasic() {
dim := 4
byteRow := make([]byte, dim*2)
rows := [][][]byte{{byteRow}}
col := NewColumnBFloat16VectorArray("emb", dim, rows)
s.Equal(entity.FieldTypeBFloat16Vector, col.ElementType())
s.Equal(1, col.Len())
s.NoError(col.AppendValue([]entity.BFloat16Vector{entity.BFloat16Vector(byteRow)}))
s.NoError(col.AppendValue([][]byte{byteRow}))
s.Error(col.AppendValue("bad"))
s.Equal(3, col.Len())
fd := col.FieldData()
s.Equal(schemapb.DataType_BFloat16Vector, fd.GetVectors().GetVectorArray().GetElementType())
}
func (s *VectorArraySuite) TestBinaryVectorArrayBasic() {
dim := 8 // binary dim is bits; 1 byte per vector.
byteRow := make([]byte, dim/8)
rows := [][][]byte{{byteRow, byteRow}}
col := NewColumnBinaryVectorArray("emb", dim, rows)
s.Equal(entity.FieldTypeBinaryVector, col.ElementType())
s.Equal(1, col.Len())
s.NoError(col.AppendValue([]entity.BinaryVector{entity.BinaryVector(byteRow)}))
s.NoError(col.AppendValue([][]byte{byteRow}))
s.Error(col.AppendValue(42))
s.Equal(3, col.Len())
fd := col.FieldData()
s.Equal(schemapb.DataType_BinaryVector, fd.GetVectors().GetVectorArray().GetElementType())
}
func (s *VectorArraySuite) TestInt8VectorArrayBasic() {
dim := 4
rows := [][][]int8{{{1, 2, 3, 4}, {5, 6, 7, 8}}, {{9, 10, 11, 12}}}
col := NewColumnInt8VectorArray("emb", dim, rows)
s.Equal(entity.FieldTypeInt8Vector, col.ElementType())
s.Equal(2, col.Len())
s.NoError(col.AppendValue([]entity.Int8Vector{entity.Int8Vector([]int8{1, 2, 3, 4})}))
s.NoError(col.AppendValue([][]int8{{5, 6, 7, 8}}))
s.Error(col.AppendValue("bad"))
s.Equal(4, col.Len())
fd := col.FieldData()
s.Equal(schemapb.DataType_Int8Vector, fd.GetVectors().GetVectorArray().GetElementType())
}
func (s *VectorArraySuite) TestBaseAppendValueRejectsWrongType() {
dim := 4
col := NewColumnFloat16VectorArray("emb", dim, nil)
// The columnVectorArrayBase AppendValue path (via embedded struct) rejects mismatched types.
s.Error(col.columnVectorArrayBase.AppendValue([]int{1, 2}))
// The positive path is exercised in TestFloat16VectorArrayBasic.
}
func (s *VectorArraySuite) TestParseVectorArrayDataFloatSuccess() {
dim := 4
row := &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4, 5, 6, 7, 8}},
},
}
fd := &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: "emb",
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Dim: int64(dim),
ElementType: schemapb.DataType_FloatVector,
Data: []*schemapb.VectorField{row, row},
},
},
},
},
}
col, err := FieldDataColumn(fd, 0, -1)
s.NoError(err)
s.Equal(2, col.Len())
fv, ok := col.(*ColumnFloatVectorArray)
s.Require().True(ok)
s.Equal(dim, fv.Dim())
}
func (s *VectorArraySuite) TestParseVectorArrayDataByteTypes() {
dim := 4
// Build a single payload holding 2 inner vectors of `dim*2` bytes each.
fp16Row := &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Float16Vector{Float16Vector: make([]byte, dim*2*2)},
}
bf16Row := &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Bfloat16Vector{Bfloat16Vector: make([]byte, dim*2*2)},
}
int8Row := &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_Int8Vector{Int8Vector: make([]byte, dim*2)},
}
// Binary: dim=16 bits -> 2 bytes per inner vector.
binRow := &schemapb.VectorField{
Dim: 16,
Data: &schemapb.VectorField_BinaryVector{BinaryVector: make([]byte, 2*2)},
}
cases := []struct {
name string
elem schemapb.DataType
innerFD *schemapb.VectorField
arrDim int64
wantCol Column
}{
{"float16", schemapb.DataType_Float16Vector, fp16Row, int64(dim), (*ColumnFloat16VectorArray)(nil)},
{"bfloat16", schemapb.DataType_BFloat16Vector, bf16Row, int64(dim), (*ColumnBFloat16VectorArray)(nil)},
{"int8", schemapb.DataType_Int8Vector, int8Row, int64(dim), (*ColumnInt8VectorArray)(nil)},
{"binary", schemapb.DataType_BinaryVector, binRow, 16, (*ColumnBinaryVectorArray)(nil)},
}
for _, c := range cases {
s.Run(c.name, func() {
fd := &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: "emb",
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: c.arrDim,
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Dim: c.arrDim,
ElementType: c.elem,
Data: []*schemapb.VectorField{c.innerFD},
},
},
},
},
}
col, err := FieldDataColumn(fd, 0, -1)
s.NoError(err)
s.Equal(1, col.Len())
s.IsType(c.wantCol, col)
})
}
}
func (s *VectorArraySuite) TestParseVectorArrayDataUnsupportedElement() {
fd := &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: "emb",
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Dim: 4,
ElementType: schemapb.DataType_SparseFloatVector, // not supported inside ArrayOfVector
Data: nil,
},
},
},
},
}
_, err := FieldDataColumn(fd, 0, -1)
s.Error(err)
}
func (s *VectorArraySuite) TestParseVectorArrayDataBeginEndClamping() {
dim := 2
row := &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2}},
},
}
fd := &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: "emb",
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: int64(dim),
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Dim: int64(dim),
ElementType: schemapb.DataType_FloatVector,
Data: []*schemapb.VectorField{row, row, row},
},
},
},
},
}
// negative begin gets clamped to 0; end > len clamped to len.
col, err := FieldDataColumn(fd, -5, 10)
s.NoError(err)
s.Equal(3, col.Len())
// begin > end collapses to empty.
col2, err := FieldDataColumn(fd, 10, 1)
s.NoError(err)
s.Equal(0, col2.Len())
}
func (s *VectorArraySuite) TestParseVectorArrayDataFallbackDim() {
// Outer Dim=0 should fall back to the first non-zero inner VectorField.Dim.
row := &schemapb.VectorField{
Dim: 4,
Data: &schemapb.VectorField_FloatVector{
FloatVector: &schemapb.FloatArray{Data: []float32{1, 2, 3, 4}},
},
}
fd := &schemapb.FieldData{
Type: schemapb.DataType_ArrayOfVector,
FieldName: "emb",
Field: &schemapb.FieldData_Vectors{
Vectors: &schemapb.VectorField{
Dim: 0,
Data: &schemapb.VectorField_VectorArray{
VectorArray: &schemapb.VectorArray{
Dim: 0,
ElementType: schemapb.DataType_FloatVector,
Data: []*schemapb.VectorField{row},
},
},
},
},
}
col, err := FieldDataColumn(fd, 0, -1)
s.NoError(err)
fv, ok := col.(*ColumnFloatVectorArray)
s.Require().True(ok)
s.Equal(4, fv.Dim())
}
func (s *VectorArraySuite) TestSlice() {
dim := 2
rows := [][][]float32{
{{1, 2}},
{{3, 4}, {5, 6}},
{{7, 8}},
{{9, 10}, {11, 12}},
}
col := NewColumnFloatVectorArray("emb", dim, rows)
sliced := col.Slice(1, 3)
s.Equal(2, sliced.Len())
// end clamped to len.
sliced2 := col.Slice(2, 99)
s.Equal(2, sliced2.Len())
// end == -1 means to the end.
sliced3 := col.Slice(1, -1)
s.Equal(3, sliced3.Len())
// start > end is clamped to end.
sliced4 := col.Slice(3, 1)
s.Equal(0, sliced4.Len())
}
func TestVectorArray(t *testing.T) {
suite.Run(t, new(VectorArraySuite))
}