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
388 lines
11 KiB
Go
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))
|
|
}
|