mirror of https://github.com/milvus-io/milvus.git
285 lines
8.8 KiB
Go
285 lines
8.8 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"
|
|
"math/rand"
|
|
"testing"
|
|
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
"github.com/milvus-io/milvus/client/v2/entity"
|
|
)
|
|
|
|
type VectorSuite struct {
|
|
suite.Suite
|
|
}
|
|
|
|
func (s *VectorSuite) TestBasic() {
|
|
s.Run("float_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 3
|
|
dim := rand.Intn(10) + 2
|
|
data := make([][]float32, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim, func(i int) float32 {
|
|
return rand.Float32()
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnFloatVector(name, dim, data)
|
|
s.Equal(entity.FieldTypeFloatVector, column.Type())
|
|
s.Equal(name, column.Name())
|
|
s.Equal(lo.Map(data, func(row []float32, _ int) entity.FloatVector { return entity.FloatVector(row) }), column.Data())
|
|
s.Equal(dim, column.Dim())
|
|
|
|
fd := column.FieldData()
|
|
s.Equal(name, fd.GetFieldName())
|
|
s.Equal(lo.Flatten(data), fd.GetVectors().GetFloatVector().GetData())
|
|
|
|
result, err := FieldDataColumn(fd, 0, -1)
|
|
s.NoError(err)
|
|
parsed, ok := result.(*ColumnFloatVector)
|
|
if s.True(ok) {
|
|
s.Equal(entity.FieldTypeFloatVector, parsed.Type())
|
|
s.Equal(name, parsed.Name())
|
|
s.Equal(lo.Map(data, func(row []float32, _ int) entity.FloatVector { return entity.FloatVector(row) }), parsed.Data())
|
|
s.Equal(dim, parsed.Dim())
|
|
}
|
|
})
|
|
|
|
s.Run("binary_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 3
|
|
dim := (rand.Intn(10) + 1) * 8
|
|
data := make([][]byte, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim/8, func(i int) byte {
|
|
return byte(rand.Intn(math.MaxUint8))
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnBinaryVector(name, dim, data)
|
|
s.Equal(entity.FieldTypeBinaryVector, column.Type())
|
|
s.Equal(name, column.Name())
|
|
s.Equal(lo.Map(data, func(row []byte, _ int) entity.BinaryVector { return entity.BinaryVector(row) }), column.Data())
|
|
s.Equal(dim, column.Dim())
|
|
|
|
fd := column.FieldData()
|
|
s.Equal(name, fd.GetFieldName())
|
|
s.Equal(lo.Flatten(data), fd.GetVectors().GetBinaryVector())
|
|
|
|
result, err := FieldDataColumn(fd, 0, -1)
|
|
s.NoError(err)
|
|
parsed, ok := result.(*ColumnBinaryVector)
|
|
if s.True(ok) {
|
|
s.Equal(entity.FieldTypeBinaryVector, parsed.Type())
|
|
s.Equal(name, parsed.Name())
|
|
s.Equal(lo.Map(data, func(row []byte, _ int) entity.BinaryVector { return entity.BinaryVector(row) }), parsed.Data())
|
|
s.Equal(dim, parsed.Dim())
|
|
}
|
|
})
|
|
|
|
s.Run("fp16_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 3
|
|
dim := rand.Intn(10) + 1
|
|
data := make([][]byte, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim*2, func(i int) byte {
|
|
return byte(rand.Intn(math.MaxUint8))
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnFloat16Vector(name, dim, data)
|
|
s.Equal(entity.FieldTypeFloat16Vector, column.Type())
|
|
s.Equal(name, column.Name())
|
|
s.Equal(lo.Map(data, func(row []byte, _ int) entity.Float16Vector { return entity.Float16Vector(row) }), column.Data())
|
|
s.Equal(dim, column.Dim())
|
|
|
|
fd := column.FieldData()
|
|
s.Equal(name, fd.GetFieldName())
|
|
s.Equal(lo.Flatten(data), fd.GetVectors().GetFloat16Vector())
|
|
|
|
result, err := FieldDataColumn(fd, 0, -1)
|
|
s.NoError(err)
|
|
parsed, ok := result.(*ColumnFloat16Vector)
|
|
if s.True(ok) {
|
|
s.Equal(entity.FieldTypeFloat16Vector, parsed.Type())
|
|
s.Equal(name, parsed.Name())
|
|
s.Equal(lo.Map(data, func(row []byte, _ int) entity.Float16Vector { return entity.Float16Vector(row) }), parsed.Data())
|
|
s.Equal(dim, parsed.Dim())
|
|
}
|
|
})
|
|
|
|
s.Run("bf16_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 3
|
|
dim := rand.Intn(10) + 1
|
|
data := make([][]byte, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim*2, func(i int) byte {
|
|
return byte(rand.Intn(math.MaxUint8))
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnBFloat16Vector(name, dim, data)
|
|
s.Equal(entity.FieldTypeBFloat16Vector, column.Type())
|
|
s.Equal(name, column.Name())
|
|
s.Equal(lo.Map(data, func(row []byte, _ int) entity.BFloat16Vector { return entity.BFloat16Vector(row) }), column.Data())
|
|
s.Equal(dim, column.Dim())
|
|
|
|
fd := column.FieldData()
|
|
s.Equal(name, fd.GetFieldName())
|
|
s.Equal(lo.Flatten(data), fd.GetVectors().GetBfloat16Vector())
|
|
|
|
result, err := FieldDataColumn(fd, 0, -1)
|
|
s.NoError(err)
|
|
parsed, ok := result.(*ColumnBFloat16Vector)
|
|
if s.True(ok) {
|
|
s.Equal(entity.FieldTypeBFloat16Vector, parsed.Type())
|
|
s.Equal(name, parsed.Name())
|
|
s.Equal(lo.Map(data, func(row []byte, _ int) entity.BFloat16Vector { return entity.BFloat16Vector(row) }), parsed.Data())
|
|
s.Equal(dim, parsed.Dim())
|
|
}
|
|
})
|
|
|
|
s.Run("fp32 <-> fp16/bf16 vector conversion", func() {
|
|
dim := 3
|
|
data := [][]float32{{0.1, 0.2, 0.3}, {0.4, 0.5, 0.6}, {0.7, 0.8, 0.9}, {1.0, 1.1, 1.2}}
|
|
|
|
fp16Vector := NewColumnFloat16VectorFromFp32Vector("fp16_vector", dim, data[:2])
|
|
fp16Vector.AppendValue(data[2])
|
|
fp16Vector.AppendValue(data[3])
|
|
for i, vec := range fp16Vector.Data() {
|
|
fp32Vector := vec.ToFloat32Vector()
|
|
s.Equal(dim, len(fp32Vector))
|
|
for j := 0; j < dim; j++ {
|
|
s.InDelta(data[i][j], fp32Vector[j], 7e-3)
|
|
}
|
|
}
|
|
|
|
bf16Vector := NewColumnBFloat16VectorFromFp32Vector("bf16_vector", dim, data[:2])
|
|
bf16Vector.AppendValue(data[2])
|
|
bf16Vector.AppendValue(data[3])
|
|
for i, vec := range bf16Vector.Data() {
|
|
fp32Vector := vec.ToFloat32Vector()
|
|
s.Equal(dim, len(fp32Vector))
|
|
for j := 0; j < dim; j++ {
|
|
s.InDelta(data[i][j], fp32Vector[j], 7e-3)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
func (s *VectorSuite) TestSlice() {
|
|
s.Run("float_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 100
|
|
dim := rand.Intn(10) + 2
|
|
data := make([][]float32, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim, func(i int) float32 {
|
|
return rand.Float32()
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnFloatVector(name, dim, data)
|
|
|
|
l := rand.Intn(n)
|
|
sliced := column.Slice(0, l)
|
|
slicedColumn, ok := sliced.(*ColumnFloatVector)
|
|
if s.True(ok) {
|
|
s.Equal(dim, slicedColumn.Dim())
|
|
s.Equal(lo.Map(data[:l], func(row []float32, _ int) entity.FloatVector { return entity.FloatVector(row) }), slicedColumn.Data())
|
|
}
|
|
})
|
|
|
|
s.Run("binary_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 100
|
|
dim := (rand.Intn(10) + 1) * 8
|
|
data := make([][]byte, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim/8, func(i int) byte {
|
|
return byte(rand.Intn(math.MaxUint8))
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnBinaryVector(name, dim, data)
|
|
|
|
l := rand.Intn(n)
|
|
sliced := column.Slice(0, l)
|
|
slicedColumn, ok := sliced.(*ColumnBinaryVector)
|
|
if s.True(ok) {
|
|
s.Equal(dim, slicedColumn.Dim())
|
|
s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.BinaryVector { return entity.BinaryVector(row) }), slicedColumn.Data())
|
|
}
|
|
})
|
|
|
|
s.Run("fp16_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 3
|
|
dim := rand.Intn(10) + 1
|
|
data := make([][]byte, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim*2, func(i int) byte {
|
|
return byte(rand.Intn(math.MaxUint8))
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnFloat16Vector(name, dim, data)
|
|
|
|
l := rand.Intn(n)
|
|
sliced := column.Slice(0, l)
|
|
slicedColumn, ok := sliced.(*ColumnFloat16Vector)
|
|
if s.True(ok) {
|
|
s.Equal(dim, slicedColumn.Dim())
|
|
s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.Float16Vector { return entity.Float16Vector(row) }), slicedColumn.Data())
|
|
}
|
|
})
|
|
|
|
s.Run("bf16_vector", func() {
|
|
name := fmt.Sprintf("field_%d", rand.Intn(1000))
|
|
n := 3
|
|
dim := rand.Intn(10) + 1
|
|
data := make([][]byte, 0, n)
|
|
for i := 0; i < n; i++ {
|
|
row := lo.RepeatBy(dim*2, func(i int) byte {
|
|
return byte(rand.Intn(math.MaxUint8))
|
|
})
|
|
data = append(data, row)
|
|
}
|
|
column := NewColumnBFloat16Vector(name, dim, data)
|
|
|
|
l := rand.Intn(n)
|
|
sliced := column.Slice(0, l)
|
|
slicedColumn, ok := sliced.(*ColumnBFloat16Vector)
|
|
if s.True(ok) {
|
|
s.Equal(dim, slicedColumn.Dim())
|
|
s.Equal(lo.Map(data[:l], func(row []byte, _ int) entity.BFloat16Vector { return entity.BFloat16Vector(row) }), slicedColumn.Data())
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestVectors(t *testing.T) {
|
|
suite.Run(t, new(VectorSuite))
|
|
}
|