Repository: milvus-io/milvus-sdk-go
Pull request #154 (2021-07-20 07:10:04)

Implement CalcDistance with vector param feature
Fully implement `CalcDistance`
Corresponding test cases added

Fix: #153 

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/154#" 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/154#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/154#issuecomment-883183362" 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.go b/client/client.go
index 19dc62c..acfab87 100644
--- a/client/client.go
+++ b/client/client.go
@@ -90,9 +90,9 @@ type Client interface {
 	Search(ctx context.Context, collName string, partitions []string,
 		expr string, outputFields []string, vectors []entity.Vector, vectorField string, metricType entity.MetricType, topK int, sp entity.SearchParam) ([]SearchResult, error)
 
-	// CalcDistanceWithIDs calculate the distance between vectors specified by ids provided
-	CalcDistanceWithIDs(ctx context.Context, collName string, partitions []string, fieldName string,
-		metricType entity.MetricType, idsl, idsr entity.Column) (entity.Column, error)
+	// CalcDistance calculate the distance between vectors specified by ids or provided
+	CalcDistance(ctx context.Context, collName string, partitions []string,
+		metricType entity.MetricType, opLeft, opRight entity.Column) (entity.Column, error)
 }
 
 // SearchResult contains the result from Search api of client
diff --git a/client/client_grpc_data.go b/client/client_grpc_data.go
index e47e072..aa09591 100644
--- a/client/client_grpc_data.go
+++ b/client/client_grpc_data.go
@@ -341,91 +341,42 @@ func (c *grpcClient) GetQuerySegmentInfo(ctx context.Context, collName string) (
 	return segments, nil
 }
 
-// CalcDistanceWithIDs calculate distance between vectors
-func (c *grpcClient) CalcDistanceWithIDs(ctx context.Context, collName string, partitions []string, fieldName string,
-	metricType entity.MetricType, idsl, idsr entity.Column) (entity.Column, error) {
+func (c *grpcClient) CalcDistance(ctx context.Context, collName string, partitions []string,
+	metricType entity.MetricType, opLeft, opRight entity.Column) (entity.Column, error) {
 	if c.service == nil {
 		return nil, ErrClientNotReady
 	}
-
-	if idsl == nil || idsr == nil {
-		return nil, errors.New("ids cannot be nil")
+	if opLeft == nil || opRight == nil {
+		return nil, errors.New("operators cannot be nil")
 	}
 
-	if idsl.Len() != idsl.Len() {
-		return nil, errors.New("ids length not match")
+	// check meta
+	if err := c.checkCollectionExists(ctx, collName); err != nil {
+		return nil, err
 	}
-	if idsl.Type() != idsl.Type() {
-		return nil, errors.New("ids type not match")
+	for _, partition := range partitions {
+		if err := c.checkPartitionExists(ctx, collName, partition); err != nil {
+			return nil, err
+		}
 	}
-
-	// describe collection to check id provided is primary key
-	coll, err := c.DescribeCollection(ctx, collName)
-	if err != nil {
+	if err := c.checkCollField(ctx, collName, opLeft.Name()); err != nil {
 		return nil, err
 	}
-
-	if !(isCollectionPrimaryKey(coll, idsl) && isCollectionPrimaryKey(coll, idsr)) {
-		return nil, errors.New("column(s) passed is no the primary key of the collection")
-	}
-
-	colToIDs := func(col entity.Column) *server.VectorIDs {
-		switch col.Type() {
-		case entity.FieldTypeInt64:
-			int64Column, ok := col.(*entity.ColumnInt64)
-			if !ok {
-				return nil // server shall report error
-			}
-			return &server.VectorIDs{
-				CollectionName: collName,
-				PartitionNames: partitions,
-				FieldName:      fieldName,
-				IdArray: &schema.IDs{
-					IdField: &schema.IDs_IntId{
-						IntId: &schema.LongArray{
-							Data: int64Column.Data(),
-						},
-					},
-				},
-			}
-		case entity.FieldTypeString:
-			//NOT supported yet
-			stringColumn, ok := col.(*entity.ColumnString)
-			if !ok {
-				return nil // server shall report error
-			}
-			return &server.VectorIDs{
-				CollectionName: collName,
-				PartitionNames: partitions,
-				FieldName:      fieldName,
-				IdArray: &schema.IDs{
-					IdField: &schema.IDs_StrId{
-						StrId: &schema.StringArray{
-							Data: stringColumn.Data(),
-						},
-					},
-				},
-			}
-
-		}
-		return nil
+	if err := c.checkCollField(ctx, collName, opRight.Name()); err != nil {
+		return nil, err
 	}
 
 	req := &server.CalcDistanceRequest{
-		OpLeft: &server.VectorsArray{
-			Array: &server.VectorsArray_IdArray{
-				IdArray: colToIDs(idsl),
-			},
-		},
-		OpRight: &server.VectorsArray{
-			Array: &server.VectorsArray_IdArray{
-				IdArray: colToIDs(idsr),
-			},
-		},
+		OpLeft:  columnToVectorsArray(collName, partitions, opLeft),
+		OpRight: columnToVectorsArray(collName, partitions, opRight),
 		Params: entity.MapKvPairs(map[string]string{
 			"metric": string(metricType),
 		}),
 	}
+	if req.OpLeft == nil || req.OpRight == nil {
+		return nil, errors.New("invalid operator passed")
+	}
+
 	resp, err := c.service.CalcDistance(ctx, req)
 	if err != nil {
 		return nil, err
@@ -444,6 +395,85 @@ func (c *grpcClient) CalcDistanceWithIDs(ctx context.Context, collName string, p
 	return nil, errors.New("distance field not supported")
 }
 
+func columnToVectorsArray(collName string, partitions []string, column entity.Column) *server.VectorsArray {
+	result := &server.VectorsArray{}
+	switch column.Type() {
+	case entity.FieldTypeInt64: // int64 id
+		int64Column, ok := column.(*entity.ColumnInt64)
+		if !ok {
+			return nil // server shall report error
+		}
+		ids := &server.VectorIDs{
+			CollectionName: collName,
+			PartitionNames: partitions,
+			FieldName:      column.Name(), //TODO use field name or column name?
+			IdArray: &schema.IDs{
+				IdField: &schema.IDs_IntId{
+					IntId: &schema.LongArray{
+						Data: int64Column.Data(),
+					},
+				},
+			},
+		}
+		result.Array = &server.VectorsArray_IdArray{IdArray: ids}
+	case entity.FieldTypeString: // string id
+		stringColumn, ok := column.(*entity.ColumnString)
+		if !ok {
+			return nil
+		}
+		ids := &server.VectorIDs{
+			CollectionName: collName,
+			PartitionNames: partitions,
+			FieldName:      column.Name(),
+			IdArray: &schema.IDs{
+				IdField: &schema.IDs_StrId{
+					StrId: &schema.StringArray{
+						Data: stringColumn.Data(),
+					},
+				},
+			},
+		}
+		result.Array = &server.VectorsArray_IdArray{IdArray: ids}
+	case entity.FieldTypeFloatVector:
+		fvColumn, ok := column.(*entity.ColumnFloatVector)
+		if !ok {
+			return nil
+		}
+		fvdata := fvColumn.Data()
+		data := make([]float32, 0, fvColumn.Len()*fvColumn.Dim())
+		for _, row := range fvdata {
+			data = append(data, row...)
+		}
+		result.Array = &server.VectorsArray_DataArray{DataArray: &schema.VectorField{
+			Dim: int64(fvColumn.Dim()),
+			Data: &schema.VectorField_FloatVector{
+				FloatVector: &schema.FloatArray{
+					Data: data,
+				},
+			},
+		}}
+	case entity.FieldTypeBinaryVector:
+		bvColumn, ok := column.(*entity.ColumnBinaryVector)
+		if !ok {
+			return nil
+		}
+		bvdata := bvColumn.Data()
+		data := make([]byte, 0, bvColumn.Dim()*bvColumn.Len()/8)
+		for _, row := range bvdata {
+			data = append(data, row...)
+		}
+		result.Array = &server.VectorsArray_DataArray{DataArray: &schema.VectorField{
+			Dim: int64(bvColumn.Dim()),
+			Data: &schema.VectorField_BinaryVector{
+				BinaryVector: data,
+			},
+		}}
+	default:
+		return nil
+	}
+	return result
+}
+
 func vector2PlaceholderGroupBytes(vectors []entity.Vector) []byte {
 	phg := &server.PlaceholderGroup{
 		Placeholders: []*server.PlaceholderValue{
diff --git a/client/client_grpc_partition.go b/client/client_grpc_partition.go
index c7b8b3d..749088f 100644
--- a/client/client_grpc_partition.go
+++ b/client/client_grpc_partition.go
@@ -119,6 +119,9 @@ func (c *grpcClient) ShowPartitions(ctx context.Context, collName string) ([]*en
 	if err != nil {
 		return []*entity.Partition{}, err
 	}
+	if err := handleRespStatus(resp.GetStatus()); err != nil {
+		return []*entity.Partition{}, err
+	}
 	partitions := make([]*entity.Partition, 0, len(resp.GetPartitionIDs()))
 	for idx, partitionID := range resp.GetPartitionIDs() {
 		partitions = append(partitions, &entity.Partition{ID: partitionID, Name: resp.GetPartitionNames()[idx]})
diff --git a/entity/columns_binaryvector.go b/entity/columns_binaryvector.go
index 335e6a6..e1eccb2 100755
--- a/entity/columns_binaryvector.go
+++ b/entity/columns_binaryvector.go
@@ -1,5 +1,5 @@
 // Code generated by go generate; DO NOT EDIT
-// This file is generated by go genrated at 2021-07-19 17:30:03.688864552 +0800 CST m=+0.002944243
+// This file is generated by go genrated at 2021-07-20 13:48:08.384238299 +0800 CST m=+0.003190552
 
 package entity 
 
@@ -33,6 +33,11 @@ func (c *ColumnBinaryVector) Dim() int {
 	return c.dim
 }
 
+// Data returns column data
+func (c *ColumnBinaryVector) Data() [][]byte {
+	return c.values
+}
+
 // FieldData return column data mapped to schema.FieldData
 func (c *ColumnBinaryVector) FieldData() *schema.FieldData {
 	fd := &schema.FieldData{
diff --git a/entity/columns_floatvector.go b/entity/columns_floatvector.go
index 0960277..80ccd7b 100755
--- a/entity/columns_floatvector.go
+++ b/entity/columns_floatvector.go
@@ -1,5 +1,5 @@
 // Code generated by go generate; DO NOT EDIT
-// This file is generated by go genrated at 2021-07-19 17:30:03.688864552 +0800 CST m=+0.002944243
+// This file is generated by go genrated at 2021-07-20 13:48:08.384238299 +0800 CST m=+0.003190552
 
 package entity 
 
@@ -33,6 +33,11 @@ func (c *ColumnFloatVector) Dim() int {
 	return c.dim
 }
 
+// Data returns column data
+func (c *ColumnFloatVector) Data() [][]float32 {
+	return c.values
+}
+
 // FieldData return column data mapped to schema.FieldData
 func (c *ColumnFloatVector) FieldData() *schema.FieldData {
 	fd := &schema.FieldData{
diff --git a/entity/gen/gen.go b/entity/gen/gen.go
index 5981494..1fa8d91 100644
--- a/entity/gen/gen.go
+++ b/entity/gen/gen.go
@@ -136,6 +136,11 @@ func (c *Column{{.TypeName}}) Dim() int {
 	return c.dim
 }
 
+// Data returns column data
+func (c *Column{{.TypeName}}) Data() []{{.TypeDef}} {
+	return c.values
+}
+
 // FieldData return column data mapped to schema.FieldData
 func (c *Column{{.TypeName}}) FieldData() *schema.FieldData {
 	fd := &schema.FieldData{
@@ -308,6 +313,7 @@ func TestColumn{{.TypeName}}(t *testing.T) {
 		assert.Equal(t, FieldType{{.TypeName}}, column.Type())
 		assert.Equal(t, columnLen, column.Len())
 		assert.Equal(t, dim, column.Dim())
+		assert.Equal(t ,v, column.Data())
 	})
 
 	t.Run("test column field data", func(t *testing.T) {


Test patch:
diff --git a/client/client_grpc_collection_test.go b/client/client_grpc_collection_test.go
index 4983879..286575b 100644
--- a/client/client_grpc_collection_test.go
+++ b/client/client_grpc_collection_test.go
@@ -44,7 +44,7 @@ func TestGrpcClientNil(t *testing.T) {
 		mt := m.Type                                   // type of function
 		if m.Name == "Close" || m.Name == "Connect" || // skip connect & close
 			m.Name == "Search" || // type alias MetricType treated as string
-			m.Name == "CalcDistanceWithIDs" ||
+			m.Name == "CalcDistance" ||
 			m.Name == "Insert" { // complex methods with ...
 			continue
 		}
diff --git a/client/client_grpc_data_test.go b/client/client_grpc_data_test.go
index 3f66835..d2d3453 100644
--- a/client/client_grpc_data_test.go
+++ b/client/client_grpc_data_test.go
@@ -296,7 +296,7 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 	ctx := context.Background()
 	t.Run("bad client calls CalcDistance", func(t *testing.T) {
 		c := &grpcClient{}
-		r, err := c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.L2, nil, nil)
+		r, err := c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2, nil, nil)
 		assert.Nil(t, r)
 		assert.NotNil(t, err)
 		assert.EqualValues(t, ErrClientNotReady, err)
@@ -324,24 +324,24 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 		ctxDone, cancel := context.WithCancel(context.Background())
 		cancel()
 
-		r, err := c.CalcDistanceWithIDs(ctxDone, testCollectionName, []string{}, "vector", entity.L2,
+		r, err := c.CalcDistance(ctxDone, testCollectionName, []string{}, entity.L2,
 			entity.NewColumnInt64("int64", []int64{1}), entity.NewColumnInt64("int64", []int64{1}))
 		assert.Nil(t, r)
 		assert.NotNil(t, err)
 	})
 
 	t.Run("invalid ids call", func(t *testing.T) {
-		r, err := c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.L2,
+		r, err := c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
 			nil, nil)
 		assert.Nil(t, r)
 		assert.NotNil(t, err)
 
-		r, err = c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.L2,
+		r, err = c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
 			entity.NewColumnInt64("non-exists", []int64{1}), entity.NewColumnInt64("non-exists", []int64{1}))
 		assert.Nil(t, r)
 		assert.NotNil(t, err)
 
-		r, err = c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.L2,
+		r, err = c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
 			entity.NewColumnInt64("non-exists", []int64{1}), entity.NewColumnInt64("int64", []int64{1}))
 		assert.Nil(t, r)
 		assert.NotNil(t, err)
@@ -349,7 +349,7 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 	})
 
 	t.Run("valid calls", func(t *testing.T) {
-		mock.setInjection(mCalcDistanceIDs, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		mock.setInjection(mCalcDistance, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 			req, ok := raw.(*server.CalcDistanceRequest)
 			resp := &server.CalcDistanceResults{}
 			if !ok {
@@ -358,23 +358,38 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 				return resp, err
 			}
 			idsLeft := req.GetOpLeft().GetIdArray()
+			valuesLeft := req.GetOpLeft().GetDataArray()
 			idsRight := req.GetOpRight().GetIdArray()
-			assert.NotNil(t, idsLeft)
-			assert.NotNil(t, idsRight)
-
-			assert.Equal(t, testCollectionName, idsLeft.CollectionName)
-			assert.Equal(t, testCollectionName, idsRight.CollectionName)
+			valuesRight := req.GetOpRight().GetDataArray()
+			assert.True(t, idsLeft != nil || valuesLeft != nil)
+			assert.True(t, idsRight != nil || valuesRight != nil)
 
-			// only int ids supported for now TODO update string test cases
-			assert.NotNil(t, idsLeft.IdArray.GetIntId())
-			assert.NotNil(t, idsRight.IdArray.GetIntId())
+			if idsLeft != nil {
+				assert.Equal(t, testCollectionName, idsLeft.CollectionName)
+			}
+			if idsRight != nil {
+				assert.Equal(t, testCollectionName, idsRight.CollectionName)
+			}
 
 			// this injection returns float distance
-			dl := len(idsLeft.IdArray.GetIntId().Data)
+			dl := 0
+			if idsLeft != nil {
+				dl = len(idsLeft.IdArray.GetIntId().Data)
+			}
+			if valuesLeft != nil {
+				dl = len(valuesLeft.GetFloatVector().GetData()) / int(valuesLeft.Dim)
+			}
+			dr := 0
+			if idsRight != nil {
+				dr = len(idsRight.IdArray.GetIntId().Data)
+			}
+			if valuesRight != nil {
+				dr = len(valuesRight.GetFloatVector().GetData()) / int(valuesRight.Dim)
+			}
 
 			resp.Array = &server.CalcDistanceResults_FloatDist{
 				FloatDist: &schema.FloatArray{
-					Data: make([]float32, dl),
+					Data: make([]float32, dl*dr),
 				},
 			}
 
@@ -382,13 +397,24 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 			resp.Status = s
 			return resp, err
 		})
-		r, err := c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.L2,
-			entity.NewColumnInt64("int64", []int64{1}), entity.NewColumnInt64("int64", []int64{1}))
+		r, err := c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
+			entity.NewColumnInt64("vector", []int64{1}), entity.NewColumnInt64("vector", []int64{1}))
+		assert.Nil(t, err)
+		assert.NotNil(t, r)
+
+		vectors := generateFloatVector(5, testVectorDim)
+		r, err = c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
+			entity.NewColumnInt64("vector", []int64{1}), entity.NewColumnFloatVector("vector", testVectorDim, vectors))
+		assert.Nil(t, err)
+		assert.NotNil(t, r)
+
+		r, err = c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
+			entity.NewColumnFloatVector("vector", testVectorDim, vectors), entity.NewColumnInt64("vector", []int64{1}))
 		assert.Nil(t, err)
 		assert.NotNil(t, r)
 
 		// test IntDistance,
-		mock.setInjection(mCalcDistanceIDs, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		mock.setInjection(mCalcDistance, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 			req, ok := raw.(*server.CalcDistanceRequest)
 			resp := &server.CalcDistanceResults{}
 			if !ok {
@@ -397,23 +423,38 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 				return resp, err
 			}
 			idsLeft := req.GetOpLeft().GetIdArray()
+			valuesLeft := req.GetOpLeft().GetDataArray()
 			idsRight := req.GetOpRight().GetIdArray()
-			assert.NotNil(t, idsLeft)
-			assert.NotNil(t, idsRight)
-
-			assert.Equal(t, testCollectionName, idsLeft.CollectionName)
-			assert.Equal(t, testCollectionName, idsRight.CollectionName)
+			valuesRight := req.GetOpRight().GetDataArray()
+			assert.True(t, idsLeft != nil || valuesLeft != nil)
+			assert.True(t, idsRight != nil || valuesRight != nil)
 
-			// only int ids supported for now TODO update string test cases
-			assert.NotNil(t, idsLeft.IdArray.GetIntId())
-			assert.NotNil(t, idsRight.IdArray.GetIntId())
+			if idsLeft != nil {
+				assert.Equal(t, testCollectionName, idsLeft.CollectionName)
+			}
+			if idsRight != nil {
+				assert.Equal(t, testCollectionName, idsRight.CollectionName)
+			}
 
 			// this injection returns float distance
-			dl := len(idsLeft.IdArray.GetIntId().Data)
+			dl := 0
+			if idsLeft != nil {
+				dl = len(idsLeft.IdArray.GetIntId().Data)
+			}
+			if valuesLeft != nil {
+				dl = len(valuesLeft.GetFloatVector().GetData()) / int(valuesLeft.Dim)
+			}
+			dr := 0
+			if idsRight != nil {
+				dr = len(idsRight.IdArray.GetIntId().Data)
+			}
+			if valuesRight != nil {
+				dr = len(valuesRight.GetFloatVector().GetData()) / int(valuesRight.Dim)
+			}
 
 			resp.Array = &server.CalcDistanceResults_IntDist{
 				IntDist: &schema.IntArray{
-					Data: make([]int32, dl),
+					Data: make([]int32, dl*dr),
 				},
 			}
 
@@ -421,8 +462,8 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 			resp.Status = s
 			return resp, err
 		})
-		r, err = c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.HAMMING,
-			entity.NewColumnInt64("int64", []int64{1}), entity.NewColumnInt64("int64", []int64{1}))
+		r, err = c.CalcDistance(ctx, testCollectionName, []string{}, entity.HAMMING,
+			entity.NewColumnInt64("vector", []int64{1}), entity.NewColumnInt64("vector", []int64{1}))
 		assert.Nil(t, err)
 		assert.NotNil(t, r)
 
@@ -446,7 +487,7 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 			resp.Status = s
 			return resp, err
 		})
-		mock.setInjection(mCalcDistanceIDs, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		mock.setInjection(mCalcDistance, func(_ context.Context, raw proto.Message) (proto.Message, error) {
 			req, ok := raw.(*server.CalcDistanceRequest)
 			resp := &server.CalcDistanceResults{}
 			if !ok {
@@ -479,13 +520,33 @@ func TestGrpcCalcDistanceWithIDs(t *testing.T) {
 			resp.Status = s
 			return resp, err
 		})
-		r, err = c.CalcDistanceWithIDs(ctx, testCollectionName, []string{}, "vector", entity.L2,
-			entity.NewColumnString("str", []string{"1"}), entity.NewColumnString("str", []string{"1"}))
+		r, err = c.CalcDistance(ctx, testCollectionName, []string{}, entity.L2,
+			entity.NewColumnString("vector", []string{"1"}), entity.NewColumnString("vector", []string{"1"}))
 		assert.Nil(t, err)
 		assert.NotNil(t, r)
 	})
 }
 
+func TestIsCollectionPrimaryKey(t *testing.T) {
+	t.Run("nil cases", func(t *testing.T) {
+		assert.False(t, isCollectionPrimaryKey(nil, nil))
+		assert.False(t, isCollectionPrimaryKey(&entity.Collection{}, entity.NewColumnInt64("id", []int64{})))
+	})
+
+	t.Run("check cases", func(t *testing.T) {
+		assert.False(t, isCollectionPrimaryKey(&entity.Collection{
+			Schema: defaultSchema(),
+		}, entity.NewColumnInt64("id", []int64{})))
+		assert.False(t, isCollectionPrimaryKey(&entity.Collection{
+			Schema: defaultSchema(),
+		}, entity.NewColumnInt32("int64", []int32{})))
+		assert.True(t, isCollectionPrimaryKey(&entity.Collection{
+			Schema: defaultSchema(),
+		}, entity.NewColumnInt64("int64", []int64{})))
+
+	})
+}
+
 func generateFloatVector(num, dim int) [][]float32 {
 	r := make([][]float32, 0, num)
 	for i := 0; i < num; i++ {
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
index 45123a1..7fd368a 100644
--- a/client/client_mock_test.go
+++ b/client/client_mock_test.go
@@ -123,10 +123,10 @@ const (
 	mGetIndexState         serviceMethod = 23
 	mGetIndexBuildProgress serviceMethod = 24
 
-	mInsert          serviceMethod = 30
-	mFlush           serviceMethod = 31
-	mSearch          serviceMethod = 32
-	mCalcDistanceIDs serviceMethod = 33
+	mInsert       serviceMethod = 30
+	mFlush        serviceMethod = 31
+	mSearch       serviceMethod = 32
+	mCalcDistance serviceMethod = 33
 
 	mGetPersistentSegmentInfo serviceMethod = 98
 	mGetQuerySegmentInfo      serviceMethod = 99
@@ -427,7 +427,7 @@ func (m *mockServer) RegisterLink(_ context.Context, _ *server.RegisterLinkReque
 }
 
 func (m *mockServer) CalcDistance(ctx context.Context, req *server.CalcDistanceRequest) (*server.CalcDistanceResults, error) {
-	f := m.getInjection(mCalcDistanceIDs)
+	f := m.getInjection(mCalcDistance)
 	if f != nil {
 		r, err := f(ctx, req)
 		return r.(*server.CalcDistanceResults), err
diff --git a/entity/columns_binaryvector_test.go b/entity/columns_binaryvector_test.go
index 54b10dc..73d58b2 100755
--- a/entity/columns_binaryvector_test.go
+++ b/entity/columns_binaryvector_test.go
@@ -1,5 +1,5 @@
 // Code generated by go generate; DO NOT EDIT
-// This file is generated by go genrated at 2021-07-19 17:30:03.688864552 +0800 CST m=+0.002944243
+// This file is generated by go genrated at 2021-07-20 13:48:08.384238299 +0800 CST m=+0.003190552
 
 package entity 
 
@@ -35,6 +35,7 @@ func TestColumnBinaryVector(t *testing.T) {
 		assert.Equal(t, FieldTypeBinaryVector, column.Type())
 		assert.Equal(t, columnLen, column.Len())
 		assert.Equal(t, dim, column.Dim())
+		assert.Equal(t ,v, column.Data())
 	})
 
 	t.Run("test column field data", func(t *testing.T) {
diff --git a/entity/columns_floatvector_test.go b/entity/columns_floatvector_test.go
index 3711b26..2585196 100755
--- a/entity/columns_floatvector_test.go
+++ b/entity/columns_floatvector_test.go
@@ -1,5 +1,5 @@
 // Code generated by go generate; DO NOT EDIT
-// This file is generated by go genrated at 2021-07-19 17:30:03.688864552 +0800 CST m=+0.002944243
+// This file is generated by go genrated at 2021-07-20 13:48:08.384238299 +0800 CST m=+0.003190552
 
 package entity 
 
@@ -35,6 +35,7 @@ func TestColumnFloatVector(t *testing.T) {
 		assert.Equal(t, FieldTypeFloatVector, column.Type())
 		assert.Equal(t, columnLen, column.Len())
 		assert.Equal(t, dim, column.Dim())
+		assert.Equal(t ,v, column.Data())
 	})
 
 	t.Run("test column field data", func(t *testing.T) {