Skip to content

Commit

Permalink
test: add cases for go milvusclient query (#33884)
Browse files Browse the repository at this point in the history
issue: #33419

---------

Signed-off-by: ThreadDao <yufen.zong@zilliz.com>
  • Loading branch information
ThreadDao committed Jun 26, 2024
1 parent cc54ab2 commit a77e295
Show file tree
Hide file tree
Showing 10 changed files with 756 additions and 46 deletions.
19 changes: 19 additions & 0 deletions client/column/columns.go
Original file line number Diff line number Diff line change
Expand Up @@ -310,6 +310,25 @@ func FieldDataColumn(fd *schemapb.FieldData, begin, end int) (Column, error) {
vector = append(vector, v)
}
return NewColumnBFloat16Vector(fd.GetFieldName(), dim, vector), nil
case schemapb.DataType_SparseFloatVector:
sparseVectors := fd.GetVectors().GetSparseFloatVector()
if sparseVectors == nil {
return nil, errFieldDataTypeNotMatch
}
data := sparseVectors.Contents
if end < 0 {
end = len(data)
}
data = data[begin:end]
vectors := make([]entity.SparseEmbedding, 0, len(data))
for _, bs := range data {
vector, err := entity.DeserializeSliceSparseEmbedding(bs)
if err != nil {
return nil, err
}
vectors = append(vectors, vector)
}
return NewColumnSparseVectors(fd.GetFieldName(), vectors), nil
default:
return nil, fmt.Errorf("unsupported data type %s", fd.GetType())
}
Expand Down
2 changes: 1 addition & 1 deletion client/entity/sparse.go
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ func (e sliceSparseEmbedding) Swap(i, j int) {
e.values[i], e.values[j] = e.values[j], e.values[i]
}

func deserializeSliceSparceEmbedding(bs []byte) (sliceSparseEmbedding, error) {
func DeserializeSliceSparseEmbedding(bs []byte) (sliceSparseEmbedding, error) {
length := len(bs)
if length%8 != 0 {
return sliceSparseEmbedding{}, errors.New("not valid sparse embedding bytes")
Expand Down
2 changes: 1 addition & 1 deletion client/entity/sparse_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ func TestSliceSparseEmbedding(t *testing.T) {
assert.EqualValues(t, length, se.Len())

bs := se.Serialize()
nv, err := deserializeSliceSparceEmbedding(bs)
nv, err := DeserializeSliceSparseEmbedding(bs)
require.NoError(t, err)

for i := 0; i < length; i++ {
Expand Down
104 changes: 83 additions & 21 deletions tests/go_client/common/response_checker.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package common

import (
"fmt"
"reflect"
"strings"
"testing"

Expand Down Expand Up @@ -32,7 +33,7 @@ func CheckErr(t *testing.T, actualErr error, expErrNil bool, expErrorMsg ...stri
}
}
if !contains {
t.FailNow()
t.Fatalf("CheckErr failed, actualErr doesn contains any expErrorMsg, please check test cases!")
}
}
}
Expand Down Expand Up @@ -60,31 +61,90 @@ func EqualColumn(t *testing.T, columnA column.Column, columnB column.Column) {
case entity.FieldTypeVarChar:
require.ElementsMatch(t, columnA.(*column.ColumnVarChar).Data(), columnB.(*column.ColumnVarChar).Data())
case entity.FieldTypeJSON:
log.Debug("columnA", zap.Any("data", columnA.(*column.ColumnJSONBytes).Data()))
log.Debug("columnB", zap.Any("data", columnB.(*column.ColumnJSONBytes).Data()))
require.ElementsMatch(t, columnA.(*column.ColumnJSONBytes).Data(), columnB.(*column.ColumnJSONBytes).Data())
log.Debug("data", zap.String("name", columnA.Name()), zap.Any("type", columnA.Type()), zap.Any("data", columnA.FieldData()))
log.Debug("data", zap.String("name", columnB.Name()), zap.Any("type", columnB.Type()), zap.Any("data", columnB.FieldData()))
require.Equal(t, reflect.TypeOf(columnA), reflect.TypeOf(columnB))
switch columnA.(type) {
case *column.ColumnDynamic:
require.ElementsMatch(t, columnA.(*column.ColumnDynamic).Data(), columnB.(*column.ColumnDynamic).Data())
case *column.ColumnJSONBytes:
require.ElementsMatch(t, columnA.(*column.ColumnJSONBytes).Data(), columnB.(*column.ColumnJSONBytes).Data())
}
case entity.FieldTypeFloatVector:
require.ElementsMatch(t, columnA.(*column.ColumnFloatVector).Data(), columnB.(*column.ColumnFloatVector).Data())
case entity.FieldTypeBinaryVector:
require.ElementsMatch(t, columnA.(*column.ColumnBinaryVector).Data(), columnB.(*column.ColumnBinaryVector).Data())
case entity.FieldTypeFloat16Vector:
require.ElementsMatch(t, columnA.(*column.ColumnFloat16Vector).Data(), columnB.(*column.ColumnFloat16Vector).Data())
case entity.FieldTypeBFloat16Vector:
require.ElementsMatch(t, columnA.(*column.ColumnBFloat16Vector).Data(), columnB.(*column.ColumnBFloat16Vector).Data())
case entity.FieldTypeSparseVector:
require.ElementsMatch(t, columnA.(*column.ColumnSparseFloatVector).Data(), columnB.(*column.ColumnSparseFloatVector).Data())
case entity.FieldTypeArray:
log.Info("TODO support column element type")
EqualArrayColumn(t, columnA, columnB)
default:
log.Info("Support column type is:", zap.Any("FieldType", []entity.FieldType{
entity.FieldTypeBool,
entity.FieldTypeInt8, entity.FieldTypeInt16, entity.FieldTypeInt32, entity.FieldTypeInt64,
entity.FieldTypeFloat, entity.FieldTypeDouble, entity.FieldTypeString, entity.FieldTypeVarChar,
entity.FieldTypeArray, entity.FieldTypeFloatVector, entity.FieldTypeBinaryVector,
entity.FieldTypeInt8, entity.FieldTypeInt16, entity.FieldTypeInt32,
entity.FieldTypeInt64, entity.FieldTypeFloat, entity.FieldTypeDouble, entity.FieldTypeString,
entity.FieldTypeVarChar, entity.FieldTypeArray, entity.FieldTypeFloatVector, entity.FieldTypeBinaryVector,
}))
}
}

// EqualColumn assert field data is equal of two columns
func EqualArrayColumn(t *testing.T, columnA column.Column, columnB column.Column) {
require.Equal(t, columnA.Name(), columnB.Name())
require.IsType(t, columnA.Type(), entity.FieldTypeArray)
require.IsType(t, columnB.Type(), entity.FieldTypeArray)
switch columnA.(type) {
case *column.ColumnBoolArray:
require.ElementsMatch(t, columnA.(*column.ColumnBoolArray).Data(), columnB.(*column.ColumnBoolArray).Data())
case *column.ColumnInt8Array:
require.ElementsMatch(t, columnA.(*column.ColumnInt8Array).Data(), columnB.(*column.ColumnInt8Array).Data())
case *column.ColumnInt16Array:
require.ElementsMatch(t, columnA.(*column.ColumnInt16Array).Data(), columnB.(*column.ColumnInt16Array).Data())
case *column.ColumnInt32Array:
require.ElementsMatch(t, columnA.(*column.ColumnInt32Array).Data(), columnB.(*column.ColumnInt32Array).Data())
case *column.ColumnInt64Array:
require.ElementsMatch(t, columnA.(*column.ColumnInt64Array).Data(), columnB.(*column.ColumnInt64Array).Data())
case *column.ColumnFloatArray:
require.ElementsMatch(t, columnA.(*column.ColumnFloatArray).Data(), columnB.(*column.ColumnFloatArray).Data())
case *column.ColumnDoubleArray:
require.ElementsMatch(t, columnA.(*column.ColumnDoubleArray).Data(), columnB.(*column.ColumnDoubleArray).Data())
case *column.ColumnVarCharArray:
require.ElementsMatch(t, columnA.(*column.ColumnVarCharArray).Data(), columnB.(*column.ColumnVarCharArray).Data())
default:
log.Info("Support array element type is:", zap.Any("FieldType", []entity.FieldType{
entity.FieldTypeBool, entity.FieldTypeInt8, entity.FieldTypeInt16,
entity.FieldTypeInt32, entity.FieldTypeInt64, entity.FieldTypeFloat, entity.FieldTypeDouble, entity.FieldTypeVarChar,
}))
}
}

// CheckInsertResult check insert result, ids len (insert count), ids data (pks, but no auto ids)
func CheckInsertResult(t *testing.T, expIds column.Column, insertRes clientv2.InsertResult) {
require.Equal(t, expIds.Len(), insertRes.IDs.Len())
require.Equal(t, expIds.Len(), int(insertRes.InsertCount))
actualIds := insertRes.IDs
switch expIds.Type() {
// pk field support int64 and varchar type
case entity.FieldTypeInt64:
require.ElementsMatch(t, actualIds.(*column.ColumnInt64).Data(), expIds.(*column.ColumnInt64).Data())
case entity.FieldTypeVarChar:
require.ElementsMatch(t, actualIds.(*column.ColumnVarChar).Data(), expIds.(*column.ColumnVarChar).Data())
default:
log.Info("The primary field only support ", zap.Any("type", []entity.FieldType{entity.FieldTypeInt64, entity.FieldTypeVarChar}))
}
}

// CheckOutputFields check query output fields
func CheckOutputFields(t *testing.T, expFields []string, actualColumns []column.Column) {
actualFields := make([]string, 0)
for _, actualColumn := range actualColumns {
actualFields = append(actualFields, actualColumn.Name())
}
log.Debug("CheckOutputFields", zap.Any("expFields", expFields), zap.Any("actualFields", actualFields))
require.ElementsMatchf(t, expFields, actualFields, fmt.Sprintf("Expected search output fields: %v, actual: %v", expFields, actualFields))
}

Expand All @@ -97,18 +157,20 @@ func CheckSearchResult(t *testing.T, actualSearchResults []clientv2.ResultSet, e
}
}

// CheckInsertResult check insert result, ids len (insert count), ids data (pks, but no auto ids)
func CheckInsertResult(t *testing.T, expIds column.Column, insertRes clientv2.InsertResult) {
require.Equal(t, expIds.Len(), insertRes.IDs.Len())
require.Equal(t, expIds.Len(), int(insertRes.InsertCount))
actualIds := insertRes.IDs
switch expIds.Type() {
// pk field support int64 and varchar type
case entity.FieldTypeInt64:
require.ElementsMatch(t, actualIds.(*column.ColumnInt64).Data(), expIds.(*column.ColumnInt64).Data())
case entity.FieldTypeVarChar:
require.ElementsMatch(t, actualIds.(*column.ColumnVarChar).Data(), expIds.(*column.ColumnVarChar).Data())
default:
log.Info("The primary field only support ", zap.Any("type", []entity.FieldType{entity.FieldTypeInt64, entity.FieldTypeVarChar}))
// CheckQueryResult check query result, column name, type and field
func CheckQueryResult(t *testing.T, expColumns []column.Column, actualColumns []column.Column) {
require.Equal(t, len(actualColumns), len(expColumns),
"The len of actual columns %d should greater or equal to the expected columns %d", len(actualColumns), len(expColumns))
for _, expColumn := range expColumns {
exist := false
for _, actualColumn := range actualColumns {
if expColumn.Name() == actualColumn.Name() {
exist = true
EqualColumn(t, expColumn, actualColumn)
}
}
if !exist {
log.Error("CheckQueryResult actualColumns no column", zap.String("name", expColumn.Name()))
}
}
}
8 changes: 4 additions & 4 deletions tests/go_client/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,16 @@ go 1.21
toolchain go1.21.10

require (
github.com/milvus-io/milvus/client/v2 v2.0.0-20240612020955-8cb350598c9a
github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5
github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3
github.com/quasilyte/go-ruleguard/dsl v0.3.22
github.com/stretchr/testify v1.9.0
github.com/x448/float16 v0.8.4
go.uber.org/zap v1.27.0
google.golang.org/grpc v1.64.0
)

//replace github.com/milvus-io/milvus/client/v2 v2.0.0-20240521081339-017fd7bc25de => ../../../milvus/client
//replace github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5 => ../../../milvus/client

require (
github.com/beorn7/perks v1.0.1 // indirect
Expand Down Expand Up @@ -55,7 +56,7 @@ require (
github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
github.com/magiconair/properties v1.8.5 // indirect
github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240430035521-259ae1d10016 // indirect
github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3 // indirect
github.com/mitchellh/mapstructure v1.4.1 // indirect
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
github.com/modern-go/reflect2 v1.0.2 // indirect
Expand All @@ -69,7 +70,6 @@ require (
github.com/prometheus/client_model v0.3.0 // indirect
github.com/prometheus/common v0.42.0 // indirect
github.com/prometheus/procfs v0.9.0 // indirect
github.com/quasilyte/go-ruleguard/dsl v0.3.22 // indirect
github.com/rogpeppe/go-internal v1.10.0 // indirect
github.com/samber/lo v1.27.0 // indirect
github.com/shirou/gopsutil/v3 v3.22.9 // indirect
Expand Down
10 changes: 4 additions & 6 deletions tests/go_client/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -401,12 +401,10 @@ github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfr
github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240430035521-259ae1d10016 h1:8WV4maXLeGEyJCCYIc1DmZ18H+VFAjMrwXJg5iI2nX4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240430035521-259ae1d10016/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240521081339-017fd7bc25de h1:sZ1K+F+1RMlWzuAFlhF7m68e9+G6UlRDoL7Q4Bmsc08=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240521081339-017fd7bc25de/go.mod h1:FU4rii9fMaMy1CGWgjROyKo5SCaJqcxZeIhdhuuEhrQ=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240612020955-8cb350598c9a h1:iyXy8ee4pDutzO40VA9xem1O1h/o5gI5x86ZVqVzto4=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240612020955-8cb350598c9a/go.mod h1:thfuEkUztRRmQ+qu4hCoO/6uxDJoUVNNx4vHqx9yh5I=
github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3 h1:KUSaWVePVlHMIluAXf2qmNffI1CMlGFLLiP+4iy9014=
github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5 h1:KNE4Smy6HxIpoJHLpds5BI2ZyvhrE5FcLhUcsjxAxAk=
github.com/milvus-io/milvus/client/v2 v2.0.0-20240621033600-e653ad27e2d5/go.mod h1:thfuEkUztRRmQ+qu4hCoO/6uxDJoUVNNx4vHqx9yh5I=
github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3 h1:ZBpRWhBa7FTFxW4YYVv9AUESoW1Xyb3KNXTzTqfkZmw=
github.com/milvus-io/milvus/pkg v0.0.2-0.20240317152703-17b4938985f3/go.mod h1:jQ2BUZny1COsgv1Qbcv8dmbppW+V9J/c4YQZNb3EOm8=
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
Expand Down
11 changes: 6 additions & 5 deletions tests/go_client/testcases/collection_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ func TestCreateArrayCollections(t *testing.T) {

// test create collection with partition key not supported field type
func TestCreateCollectionPartitionKey(t *testing.T) {
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout*2)
mc := createDefaultMilvusClient(ctx, t)

int64Field := entity.NewField().WithName(common.DefaultInt64FieldName).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)
Expand Down Expand Up @@ -421,7 +421,7 @@ func TestCreateCollectionWithValidName(t *testing.T) {
func TestCreateCollectionWithInvalidFieldName(t *testing.T) {
t.Parallel()
// connect
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout*2)
mc := createDefaultMilvusClient(ctx, t)

// create collection with invalid field name
Expand All @@ -436,7 +436,8 @@ func TestCreateCollectionWithInvalidFieldName(t *testing.T) {
common.CheckErr(t, err, false, "field name should not be empty",
"The first character of a field name must be an underscore or letter",
"Field name cannot only contain numbers, letters, and underscores",
"The length of a field name must be less than 255 characters")
"The length of a field name must be less than 255 characters",
"Field name can only contain numbers, letters, and underscores")
}
}

Expand Down Expand Up @@ -515,7 +516,7 @@ func TestCreateCollectionInvalidFields(t *testing.T) {

// create autoID or not collection with non-int64 and non-varchar field
func TestCreateCollectionInvalidAutoPkField(t *testing.T) {
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout*2)
mc := createDefaultMilvusClient(ctx, t)
t.Parallel()
// create collection with autoID true or not
Expand Down Expand Up @@ -557,7 +558,7 @@ func TestCreateCollectionDuplicateField(t *testing.T) {

// test create collection with partition key not supported field type
func TestCreateCollectionInvalidPartitionKeyType(t *testing.T) {
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout*2)
mc := createDefaultMilvusClient(ctx, t)

int64Field := entity.NewField().WithName(common.DefaultInt64FieldName).WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)
Expand Down
19 changes: 19 additions & 0 deletions tests/go_client/testcases/helper/data_helper.go
Original file line number Diff line number Diff line change
Expand Up @@ -394,3 +394,22 @@ func GenDynamicColumnData(start int, nb int) []column.Column {
}
return data
}

func MergeColumnsToDynamic(nb int, columns []column.Column, columnName string) *column.ColumnJSONBytes {
values := make([][]byte, 0, nb)
for i := 0; i < nb; i++ {
m := make(map[string]interface{})
for _, c := range columns {
// range guaranteed
m[c.Name()], _ = c.Get(i)
}
bs, err := json.Marshal(&m)
if err != nil {
log.Fatal("MergeColumnsToDynamic failed:", zap.Error(err))
}
values = append(values, bs)
}
jsonColumn := column.NewColumnJSONBytes(columnName, values).WithIsDynamic(true)

return jsonColumn
}
15 changes: 7 additions & 8 deletions tests/go_client/testcases/helper/helper.go
Original file line number Diff line number Diff line change
Expand Up @@ -136,16 +136,15 @@ func (chainTask *CollectionPrepare) InsertData(ctx context.Context, t *testing.T
}
fields := ip.Schema.Fields
insertOpt := clientv2.NewColumnBasedInsertOption(ip.Schema.CollectionName)
if ip.Schema.EnableDynamicField {
insertOpt.WithColumns(GenDynamicColumnData(option.start, ip.Nb)...)
}
for _, field := range fields {
if field.IsDynamic {
insertOpt.WithColumns(GenDynamicColumnData(option.start, ip.Nb)...)
} else {
if field.DataType == entity.FieldTypeArray {
option.TWithElementType(field.ElementType)
}
column := GenColumnData(ip.Nb, field.DataType, *option)
insertOpt.WithColumns(column)
if field.DataType == entity.FieldTypeArray {
option.TWithElementType(field.ElementType)
}
column := GenColumnData(ip.Nb, field.DataType, *option)
insertOpt.WithColumns(column)
}
if ip.PartitionName != "" {
insertOpt.WithPartition(ip.PartitionName)
Expand Down
Loading

0 comments on commit a77e295

Please sign in to comment.