Repository: milvus-io/milvus-sdk-go
Pull request #152 (2021-07-19 11:26:55)

Add param check for Search api
Add param check for `Search` API

Fix: #133 

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/152#" title="Author self-approved">congqixia</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
sunby: /lgtm
/approve
sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/152#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/152#issuecomment-882476253" title="Approved">sunby</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia,sunby]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/client/client_grpc_data.go b/client/client_grpc_data.go
index dc3c797..e47e072 100644
--- a/client/client_grpc_data.go
+++ b/client/client_grpc_data.go
@@ -175,6 +175,54 @@ func (c *grpcClient) Search(ctx context.Context, collName string, partitions []s
 	if c.service == nil {
 		return []SearchResult{}, ErrClientNotReady
 	}
+	// 1. check all input params
+	if err := c.checkCollectionExists(ctx, collName); err != nil {
+		return nil, err
+	}
+	for _, partition := range partitions {
+		err := c.checkPartitionExists(ctx, collName, partition)
+		if err != nil {
+			return nil, err
+		}
+	}
+	// TODO maybe add expr analysis?
+	coll, err := c.DescribeCollection(ctx, collName)
+	if err != nil {
+		return nil, err
+	}
+	mNameField := make(map[string]*entity.Field)
+	for _, field := range coll.Schema.Fields {
+		mNameField[field.Name] = field
+	}
+	for _, outField := range outputFields {
+		_, has := mNameField[outField]
+		if !has {
+			return nil, fmt.Errorf("field %s does not exist in collection %s", outField, collName)
+		}
+	}
+	vfDef, has := mNameField[vectorField]
+	if !has {
+		return nil, fmt.Errorf("vector field %s does not exist in collection %s", vectorField, collName)
+	}
+	dimStr := vfDef.TypeParams["dim"]
+	for _, vector := range vectors {
+		if fmt.Sprintf("%d", vector.Dim()) != dimStr {
+			return nil, fmt.Errorf("vector %s has dim of %s while found search vector with dim %d", vectorField,
+				dimStr, vector.Dim())
+		}
+	}
+	switch vfDef.DataType {
+	case entity.FieldTypeFloatVector:
+		if metricType != entity.IP && metricType != entity.L2 {
+			return nil, fmt.Errorf("Float vector does not support metric type %s", metricType)
+		}
+	case entity.FieldTypeBinaryVector:
+		if metricType == entity.IP || metricType == entity.L2 {
+			return nil, fmt.Errorf("Binary vector does not support metric type %s", metricType)
+		}
+	}
+
+	// 2. Request milvus service
 	bs, _ := json.Marshal(sp.Params())
 	searchParams := entity.MapKvPairs(map[string]string{
 		"anns_field":  vectorField,
@@ -200,6 +248,7 @@ func (c *grpcClient) Search(ctx context.Context, collName string, partitions []s
 	if err := handleRespStatus(resp.GetStatus()); err != nil {
 		return []SearchResult{}, err
 	}
+	// 3. parse result into result
 	results := resp.GetResults()
 	offset := 0
 	sr := make([]SearchResult, 0, results.GetNumQueries())


Test patch:
diff --git a/client/client_grpc_data_test.go b/client/client_grpc_data_test.go
index 6b0da00..3f66835 100644
--- a/client/client_grpc_data_test.go
+++ b/client/client_grpc_data_test.go
@@ -179,16 +179,117 @@ func TestGrpcSearch(t *testing.T) {
 	ctx := context.Background()
 
 	c := testClient(ctx, t)
-
 	vectors := generateFloatVector(4096, testVectorDim)
-	sp, err := entity.NewIndexFlatSearchParam(10)
-	assert.Nil(t, err)
-	assert.NotNil(t, sp)
-	results, err := c.Search(ctx, testCollectionName, []string{}, "int64 > 0", []string{"int64"}, []entity.Vector{entity.FloatVector(vectors[0])},
-		testVectorField, entity.L2, 10, sp)
 
-	assert.Nil(t, err)
-	assert.NotNil(t, results)
+	t.Run("search fail due to meta error", func(t *testing.T) {
+		sp, err := entity.NewIndexFlatSearchParam(10)
+		assert.Nil(t, err)
+
+		// collection name
+		mock.delInjection(mHasCollection)
+		r, err := c.Search(ctx, testCollectionName, []string{}, "", []string{}, []entity.Vector{}, "vector",
+			entity.L2, 5, sp)
+		assert.Nil(t, r)
+		assert.NotNil(t, err)
+
+		// partition
+		mock.setInjection(mHasCollection, hasCollectionDefault)
+		r, err = c.Search(ctx, testCollectionName, []string{"_non_exist"}, "", []string{}, []entity.Vector{}, "vector",
+			entity.L2, 5, sp)
+		assert.Nil(t, r)
+		assert.NotNil(t, err)
+
+		// output field
+		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, 0, testCollectionName, defaultSchema()))
+		r, err = c.Search(ctx, testCollectionName, []string{}, "", []string{"extra"}, []entity.Vector{}, "vector",
+			entity.L2, 5, sp)
+		assert.Nil(t, r)
+		assert.NotNil(t, err)
+
+		// vector field
+		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, 0, testCollectionName, defaultSchema()))
+		r, err = c.Search(ctx, testCollectionName, []string{}, "", []string{"int64"}, []entity.Vector{}, "no_vector",
+			entity.L2, 5, sp)
+		assert.Nil(t, r)
+		assert.NotNil(t, err)
+
+		// vector dim
+		badVectors := generateFloatVector(1, testVectorDim*2)
+		r, err = c.Search(ctx, testCollectionName, []string{}, "", []string{"int64"}, []entity.Vector{entity.FloatVector(badVectors[0])}, "vector",
+			entity.L2, 5, sp)
+		assert.Nil(t, r)
+		assert.NotNil(t, err)
+
+		// metric type
+		r, err = c.Search(ctx, testCollectionName, []string{}, "", []string{"int64"}, []entity.Vector{entity.FloatVector(vectors[0])}, "vector",
+			entity.HAMMING, 5, sp)
+		assert.Nil(t, r)
+		assert.NotNil(t, err)
+
+	})
+
+	t.Run("ok search", func(t *testing.T) {
+		mock.setInjection(mHasCollection, hasCollectionDefault)
+		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, 0, testCollectionName, defaultSchema()))
+
+		expr := `int64 > 0`
+
+		mock.setInjection(mSearch, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.SearchRequest)
+			resp := &server.SearchResults{}
+			if !ok {
+				s, err := badRequestStatus()
+				resp.Status = s
+				return resp, err
+			}
+			assert.Equal(t, testCollectionName, req.GetCollectionName())
+			assert.Equal(t, expr, req.GetDsl())
+			assert.Equal(t, common.DslType_BoolExprV1, req.GetDslType())
+			assert.ElementsMatch(t, []string{"int64"}, req.GetOutputFields())
+
+			resp.Results = &schema.SearchResultData{
+				NumQueries: 1,
+				TopK:       10,
+				FieldsData: []*schema.FieldData{
+					{
+						Type:      schema.DataType_Int64,
+						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},
+									},
+								},
+							},
+						},
+					},
+				},
+				Ids: &schema.IDs{
+					IdField: &schema.IDs_IntId{
+						IntId: &schema.LongArray{
+							Data: []int64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},
+						},
+					},
+				},
+				Scores: make([]float32, 10),
+				Topks:  []int64{10},
+			}
+
+			s, err := successStatus()
+			resp.Status = s
+			return resp, err
+		})
+
+		sp, err := entity.NewIndexFlatSearchParam(10)
+		assert.Nil(t, err)
+		assert.NotNil(t, sp)
+		results, err := c.Search(ctx, testCollectionName, []string{}, expr, []string{"int64"}, []entity.Vector{entity.FloatVector(vectors[0])},
+			testVectorField, entity.L2, 10, sp)
+
+		assert.Nil(t, err)
+		assert.NotNil(t, results)
+	})
 }
 
 func TestGrpcCalcDistanceWithIDs(t *testing.T) {
@@ -396,3 +497,13 @@ func generateFloatVector(num, dim int) [][]float32 {
 	}
 	return r
 }
+
+func generateBinaryVector(num, dim int) [][]byte {
+	r := make([][]byte, 0, num)
+	for i := 0; i < num; i++ {
+		v := make([]byte, 0, dim/8)
+		rand.Read(v)
+		r = append(r, v)
+	}
+	return r
+}