diff --git a/client/client_grpc_data.go b/client/client_grpc_data.go index 0d8856aa..ad4f71e2 100644 --- a/client/client_grpc_data.go +++ b/client/client_grpc_data.go @@ -391,6 +391,15 @@ func (c *grpcClient) QueryByPks(ctx context.Context, collectionName string, part fieldsData := resp.GetFieldsData() columns := make([]entity.Column, 0, len(fieldsData)) for _, fieldData := range resp.GetFieldsData() { + if fieldData.GetType() == schema.DataType_FloatVector || + fieldData.GetType() == schema.DataType_BinaryVector { + column, err := entity.FieldDataVector(fieldData) + if err != nil { + return nil, err + } + columns = append(columns, column) + continue + } column, err := entity.FieldDataColumn(fieldData, 0, -1) if err != nil { return nil, err diff --git a/client/client_grpc_data_test.go b/client/client_grpc_data_test.go index 0f33a84d..a4949885 100644 --- a/client/client_grpc_data_test.go +++ b/client/client_grpc_data_test.go @@ -410,16 +410,31 @@ func TestGrpcQueryByPks(t *testing.T) { }, }, }, + { + Type: schema.DataType_FloatVector, + FieldName: testVectorField, + Field: &schema.FieldData_Vectors{ + Vectors: &schema.VectorField{ + Dim: 1, + Data: &schema.VectorField_FloatVector{ + FloatVector: &schema.FloatArray{ + Data: []float32{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0}, + }, + }, + }, + }, + }, } return resp, err }) defer mock.delInjection(mQuery) - columns, err := c.QueryByPks(ctx, testCollectionName, []string{partName}, entity.NewColumnInt64(testPrimaryField, []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), []string{"int64"}) + columns, err := c.QueryByPks(ctx, testCollectionName, []string{partName}, entity.NewColumnInt64(testPrimaryField, []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}), []string{"int64", testVectorField}) assert.NoError(t, err) - assert.Equal(t, 1, len(columns)) + assert.Equal(t, 2, len(columns)) assert.Equal(t, entity.FieldTypeInt64, columns[0].Type()) + assert.Equal(t, entity.FieldTypeFloatVector, columns[1].Type()) assert.Equal(t, 10, columns[0].Len()) colInt64, ok := columns[0].(*entity.ColumnInt64) @@ -519,6 +534,36 @@ func TestGrpcQueryByPks(t *testing.T) { }) _, err = c.QueryByPks(ctx, testCollectionName, []string{}, entity.NewColumnInt64(testPrimaryField, []int64{1}), []string{"*"}) assert.Error(t, err) + + mock.setInjection(mQuery, func(_ context.Context, raw proto.Message) (proto.Message, error) { + _, ok := raw.(*server.QueryRequest) + if !ok { + t.FailNow() + } + + resp := &server.QueryResults{} + s, err := successStatus() + resp.Status = s + resp.FieldsData = []*schema.FieldData{ + { + Type: schema.DataType_FloatVector, + FieldName: "int64", + Field: &schema.FieldData_Scalars{ + Scalars: &schema.ScalarField{ + Data: &schema.ScalarField_LongData{ + LongData: &schema.LongArray{ + Data: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, + }, + }, + }, + }, + }, + } + + return resp, err + }) + _, err = c.QueryByPks(ctx, testCollectionName, []string{}, entity.NewColumnInt64(testPrimaryField, []int64{1}), []string{"*"}) + assert.Error(t, err) }) } diff --git a/entity/columns.go b/entity/columns.go index 83f2d75d..72302700 100644 --- a/entity/columns.go +++ b/entity/columns.go @@ -191,7 +191,43 @@ func FieldDataColumn(fd *schema.FieldData, begin, end int) (Column, error) { return NewColumnString(fd.GetFieldName(), data.StringData.GetData()[begin:]), nil } return NewColumnString(fd.GetFieldName(), data.StringData.GetData()[begin:end]), nil + default: + return nil, errors.New("unsupported data type") + } +} +// FieldDataColumn converts schema.FieldData to vector Column +func FieldDataVector(fd *schema.FieldData) (Column, error) { + switch fd.GetType() { + case schema.DataType_FloatVector: + vectors := fd.GetVectors() + data := vectors.GetFloatVector().GetData() + if data == nil { + return nil, errFieldDataTypeNotMatch + } + dim := int(vectors.GetDim()) + vector := make([][]float32, 0, len(data)/dim) // shall not have remanunt + for i := 0; i < len(data)/dim; i++ { + v := make([]float32, dim) + copy(v, data[i*dim:(i+1)*dim]) + vector = append(vector, v) + } + return NewColumnFloatVector(fd.GetFieldName(), dim, vector), nil + case schema.DataType_BinaryVector: + vectors := fd.GetVectors() + data := vectors.GetBinaryVector() + if data == nil { + return nil, errFieldDataTypeNotMatch + } + dim := int(vectors.GetDim()) + blen := dim / 8 + vector := make([][]byte, 0, len(data)/blen) + for i := 0; i < len(data)/blen; i++ { + v := make([]byte, blen) + copy(v, data[i*blen:(i+1)*blen]) + vector = append(vector, v) + } + return NewColumnBinaryVector(fd.GetFieldName(), dim, vector), nil default: return nil, errors.New("unsupported data type") } diff --git a/entity/columns_vector_gen.go b/entity/columns_vector_gen.go index e8fbd74d..30441bda 100755 --- a/entity/columns_vector_gen.go +++ b/entity/columns_vector_gen.go @@ -1,5 +1,5 @@ // Code generated by go generate; DO NOT EDIT -// This file is generated by go genrated at 2021-07-28 17:22:17.531936577 +0800 CST m=+0.003667571 +// This file is generated by go genrated at 2022-01-26 15:50:36.315715425 +0800 CST m=+0.003539472 package entity diff --git a/entity/columns_vector_gen_test.go b/entity/columns_vector_gen_test.go index 38208ca8..8f652574 100755 --- a/entity/columns_vector_gen_test.go +++ b/entity/columns_vector_gen_test.go @@ -1,5 +1,5 @@ // Code generated by go generate; DO NOT EDIT -// This file is generated by go genrated at 2021-07-28 17:22:17.531936577 +0800 CST m=+0.003667571 +// This file is generated by go genrated at 2022-01-26 15:50:36.315715425 +0800 CST m=+0.003539472 package entity @@ -9,17 +9,25 @@ import ( "testing" "time" + "github.com/milvus-io/milvus-sdk-go/v2/internal/proto/schema" "github.com/stretchr/testify/assert" ) func TestColumnBinaryVector(t *testing.T) { rand.Seed(time.Now().UnixNano()) columnName := fmt.Sprintf("column_BinaryVector_%d", rand.Int()) - columnLen := 8 + rand.Intn(10) - dim := ([]int{8, 32, 64, 128})[rand.Intn(4)] + columnLen := 12 + rand.Intn(10) + dim := ([]int{64, 128, 256, 512})[rand.Intn(4)] - v := make([][]byte, columnLen) - column := NewColumnBinaryVector(columnName,dim, v) + v := make([][]byte,0, columnLen) + dlen := dim + dlen /= 8 + + for i := 0; i < columnLen; i++ { + entry := make([]byte, dlen) + v = append(v, entry) + } + column := NewColumnBinaryVector(columnName, dim, v) t.Run("test meta", func(t *testing.T) { ft := FieldTypeBinaryVector @@ -51,6 +59,19 @@ func TestColumnBinaryVector(t *testing.T) { fd := column.FieldData() assert.NotNil(t, fd) assert.Equal(t, fd.GetFieldName(), columnName) + + c, err := FieldDataVector(fd) + assert.NotNil(t, c) + assert.NoError(t, err) + }) + + t.Run("test column field data error", func(t *testing.T) { + fd := &schema.FieldData{ + Type: schema.DataType_BinaryVector, + FieldName: columnName, + } + _, err := FieldDataVector(fd) + assert.Error(t, err) }) } @@ -58,11 +79,18 @@ func TestColumnBinaryVector(t *testing.T) { func TestColumnFloatVector(t *testing.T) { rand.Seed(time.Now().UnixNano()) columnName := fmt.Sprintf("column_FloatVector_%d", rand.Int()) - columnLen := 8 + rand.Intn(10) - dim := ([]int{8, 32, 64, 128})[rand.Intn(4)] + columnLen := 12 + rand.Intn(10) + dim := ([]int{64, 128, 256, 512})[rand.Intn(4)] - v := make([][]float32, columnLen) - column := NewColumnFloatVector(columnName,dim, v) + v := make([][]float32,0, columnLen) + dlen := dim + + + for i := 0; i < columnLen; i++ { + entry := make([]float32, dlen) + v = append(v, entry) + } + column := NewColumnFloatVector(columnName, dim, v) t.Run("test meta", func(t *testing.T) { ft := FieldTypeFloatVector @@ -94,6 +122,19 @@ func TestColumnFloatVector(t *testing.T) { fd := column.FieldData() assert.NotNil(t, fd) assert.Equal(t, fd.GetFieldName(), columnName) + + c, err := FieldDataVector(fd) + assert.NotNil(t, c) + assert.NoError(t, err) + }) + + t.Run("test column field data error", func(t *testing.T) { + fd := &schema.FieldData{ + Type: schema.DataType_FloatVector, + FieldName: columnName, + } + _, err := FieldDataVector(fd) + assert.Error(t, err) }) } diff --git a/entity/gen/gen.go b/entity/gen/gen.go index 182747ab..5cdce426 100644 --- a/entity/gen/gen.go +++ b/entity/gen/gen.go @@ -344,17 +344,25 @@ import ( "testing" "time" + "github.com/milvus-io/milvus-sdk-go/v2/internal/proto/schema" "github.com/stretchr/testify/assert" ) {{range .Types}}{{with .}} func TestColumn{{.TypeName}}(t *testing.T) { rand.Seed(time.Now().UnixNano()) columnName := fmt.Sprintf("column_{{.TypeName}}_%d", rand.Int()) - columnLen := 8 + rand.Intn(10) - dim := ([]int{8, 32, 64, 128})[rand.Intn(4)] + columnLen := 12 + rand.Intn(10) + dim := ([]int{64, 128, 256, 512})[rand.Intn(4)] - v := make([]{{.TypeDef}}, columnLen) - column := NewColumn{{.TypeName}}(columnName,dim, v) + v := make([]{{.TypeDef}},0, columnLen) + dlen := dim + {{if eq .TypeName "BinaryVector" }}dlen /= 8{{end}} + + for i := 0; i < columnLen; i++ { + entry := make({{.TypeDef}}, dlen) + v = append(v, entry) + } + column := NewColumn{{.TypeName}}(columnName, dim, v) t.Run("test meta", func(t *testing.T) { ft := FieldType{{.TypeName}} @@ -386,6 +394,19 @@ func TestColumn{{.TypeName}}(t *testing.T) { fd := column.FieldData() assert.NotNil(t, fd) assert.Equal(t, fd.GetFieldName(), columnName) + + c, err := FieldDataVector(fd) + assert.NotNil(t, c) + assert.NoError(t, err) + }) + + t.Run("test column field data error", func(t *testing.T) { + fd := &schema.FieldData{ + Type: schema.DataType_{{.TypeName}}, + FieldName: columnName, + } + _, err := FieldDataVector(fd) + assert.Error(t, err) }) } diff --git a/entity/indexes_gen.go b/entity/indexes_gen.go index bd477a9d..57af4d70 100755 --- a/entity/indexes_gen.go +++ b/entity/indexes_gen.go @@ -1,5 +1,5 @@ // Code generated by go generate; DO NOT EDIT -// This file is generated by go generate at 2021-07-22 12:13:27.534211219 +0800 CST m=+0.006076838 +// This file is generated by go generate at 2022-01-26 15:50:36.70745261 +0800 CST m=+0.004592013 package entity @@ -47,7 +47,7 @@ func(i *IndexFlat) Params() map[string]string { } } -// NewIndexFlat create index with contruction parameters +// NewIndexFlat create index with construction parameters func NewIndexFlat(metricType MetricType, nlist int, ) (*IndexFlat, error) { @@ -103,7 +103,7 @@ func(i *IndexBinFlat) Params() map[string]string { } } -// NewIndexBinFlat create index with contruction parameters +// NewIndexBinFlat create index with construction parameters func NewIndexBinFlat(metricType MetricType, nlist int, ) (*IndexBinFlat, error) { @@ -159,7 +159,7 @@ func(i *IndexIvfFlat) Params() map[string]string { } } -// NewIndexIvfFlat create index with contruction parameters +// NewIndexIvfFlat create index with construction parameters func NewIndexIvfFlat(metricType MetricType, nlist int, ) (*IndexIvfFlat, error) { @@ -215,7 +215,7 @@ func(i *IndexBinIvfFlat) Params() map[string]string { } } -// NewIndexBinIvfFlat create index with contruction parameters +// NewIndexBinIvfFlat create index with construction parameters func NewIndexBinIvfFlat(metricType MetricType, nlist int, ) (*IndexBinIvfFlat, error) { @@ -271,7 +271,7 @@ func(i *IndexIvfSQ8) Params() map[string]string { } } -// NewIndexIvfSQ8 create index with contruction parameters +// NewIndexIvfSQ8 create index with construction parameters func NewIndexIvfSQ8(metricType MetricType, nlist int, ) (*IndexIvfSQ8, error) { @@ -327,7 +327,7 @@ func(i *IndexIvfSQ8H) Params() map[string]string { } } -// NewIndexIvfSQ8H create index with contruction parameters +// NewIndexIvfSQ8H create index with construction parameters func NewIndexIvfSQ8H(metricType MetricType, nlist int, ) (*IndexIvfSQ8H, error) { @@ -387,7 +387,7 @@ func(i *IndexIvfPQ) Params() map[string]string { } } -// NewIndexIvfPQ create index with contruction parameters +// NewIndexIvfPQ create index with construction parameters func NewIndexIvfPQ(metricType MetricType, nlist int, @@ -466,7 +466,7 @@ func(i *IndexRNSG) Params() map[string]string { } } -// NewIndexRNSG create index with contruction parameters +// NewIndexRNSG create index with construction parameters func NewIndexRNSG(metricType MetricType, out_degree int, @@ -557,7 +557,7 @@ func(i *IndexHNSW) Params() map[string]string { } } -// NewIndexHNSW create index with contruction parameters +// NewIndexHNSW create index with construction parameters func NewIndexHNSW(metricType MetricType, M int, @@ -626,7 +626,7 @@ func(i *IndexRHNSWFlat) Params() map[string]string { } } -// NewIndexRHNSWFlat create index with contruction parameters +// NewIndexRHNSWFlat create index with construction parameters func NewIndexRHNSWFlat(metricType MetricType, M int, @@ -697,7 +697,7 @@ func(i *IndexRHNSW_PQ) Params() map[string]string { } } -// NewIndexRHNSW_PQ create index with contruction parameters +// NewIndexRHNSW_PQ create index with construction parameters func NewIndexRHNSW_PQ(metricType MetricType, M int, @@ -772,7 +772,7 @@ func(i *IndexRHNSW_SQ) Params() map[string]string { } } -// NewIndexRHNSW_SQ create index with contruction parameters +// NewIndexRHNSW_SQ create index with construction parameters func NewIndexRHNSW_SQ(metricType MetricType, M int, @@ -843,7 +843,7 @@ func(i *IndexIvfHNSW) Params() map[string]string { } } -// NewIndexIvfHNSW create index with contruction parameters +// NewIndexIvfHNSW create index with construction parameters func NewIndexIvfHNSW(metricType MetricType, nlist int, @@ -921,7 +921,7 @@ func(i *IndexANNOY) Params() map[string]string { } } -// NewIndexANNOY create index with contruction parameters +// NewIndexANNOY create index with construction parameters func NewIndexANNOY(metricType MetricType, n_trees int, ) (*IndexANNOY, error) { @@ -981,7 +981,7 @@ func(i *IndexNGTPANNG) Params() map[string]string { } } -// NewIndexNGTPANNG create index with contruction parameters +// NewIndexNGTPANNG create index with construction parameters func NewIndexNGTPANNG(metricType MetricType, edge_size int, @@ -1063,7 +1063,7 @@ func(i *IndexNGTONNG) Params() map[string]string { } } -// NewIndexNGTONNG create index with contruction parameters +// NewIndexNGTONNG create index with construction parameters func NewIndexNGTONNG(metricType MetricType, edge_size int, diff --git a/entity/indexes_gen_test.go b/entity/indexes_gen_test.go index 4d85e304..b99a3852 100755 --- a/entity/indexes_gen_test.go +++ b/entity/indexes_gen_test.go @@ -1,5 +1,5 @@ // Code generated by go generate; DO NOT EDIT -// This file is generated by go generate at 2021-07-19 17:30:04.100488399 +0800 CST m=+0.005484811 +// This file is generated by go generate at 2022-01-26 15:50:36.70745261 +0800 CST m=+0.004592013 package entity diff --git a/entity/indexes_search_param_gen.go b/entity/indexes_search_param_gen.go index 317b69d2..6ac54355 100755 --- a/entity/indexes_search_param_gen.go +++ b/entity/indexes_search_param_gen.go @@ -1,5 +1,5 @@ // Code generated by go generate; DO NOT EDIT -// This file is generated by go generate at 2021-07-22 12:13:27.534211219 +0800 CST m=+0.006076838 +// This file is generated by go generate at 2022-01-26 15:50:36.70745261 +0800 CST m=+0.004592013 package entity diff --git a/entity/indexes_search_param_gen_test.go b/entity/indexes_search_param_gen_test.go index ec4ed5ad..37b3b7f6 100755 --- a/entity/indexes_search_param_gen_test.go +++ b/entity/indexes_search_param_gen_test.go @@ -1,5 +1,5 @@ // Code generated by go generate; DO NOT EDIT -// This file is generated by go generate at 2021-07-22 12:13:27.534211219 +0800 CST m=+0.006076838 +// This file is generated by go generate at 2022-01-26 15:50:36.70745261 +0800 CST m=+0.004592013 package entity