Repository: milvus-io/milvus-sdk-go
Pull request #502 (2023-06-26 09:57:28)

Add highlevel API `Get`
/kind improvement

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/502#" 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/v2.2.x/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":[]} -->
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> Merging [#502](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (be03757) into [v2.2.x](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/commit/7589432cfe2b5fa5d5cb8aabb2ac0f7526211193?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (7589432) will **decrease** coverage by `0.12%`.
> The diff coverage is `44.44%`.

[![Impacted file tree graph](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502/graphs/tree.svg?width=650&height=150&src=pr&token=WJK43PV390&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           v2.2.x     #502      +/-   ##
==========================================
- Coverage   80.46%   80.34%   -0.12%     
==========================================
  Files          35       35              
  Lines        2820     2829       +9     
==========================================
+ Hits         2269     2273       +4     
- Misses        438      443       +5     
  Partials      113      113              
```


| [Impacted Files](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [client/client.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudC5nbw==) | `31.70% <ø> (ø)` | |
| [client/options.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L29wdGlvbnMuZ28=) | `61.76% <0.00%> (-3.87%)` | :arrow_down: |
| [client/data.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/502?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2RhdGEuZ28=) | `69.01% <80.00%> (+0.15%)` | :arrow_up: |

yah01: /lgtm

Patch:
diff --git a/client/client.go b/client/client.go
index 1d29592..2cdfe30 100644
--- a/client/client.go
+++ b/client/client.go
@@ -135,6 +135,8 @@ type Client interface {
 	QueryByPks(ctx context.Context, collectionName string, partitionNames []string, ids entity.Column, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error)
 	// Query performs query records with boolean expression.
 	Query(ctx context.Context, collectionName string, partitionNames []string, expr string, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error)
+	// Get grabs the inserted entities using the primary key from the Collection.
+	Get(ctx context.Context, collectionName string, ids entity.Column, opts ...GetOption) (ResultSet, error)
 
 	// CalcDistance calculate the distance between vectors specified by ids or provided
 	CalcDistance(ctx context.Context, collName string, partitions []string,
diff --git a/client/data.go b/client/data.go
index 08e3169..403fee6 100644
--- a/client/data.go
+++ b/client/data.go
@@ -155,6 +155,30 @@ func PKs2Expr(backName string, ids entity.Column) string {
 	return expr
 }
 
+// Get grabs the inserted entities using the primary key from the Collection.
+func (c *GrpcClient) Get(ctx context.Context, collectionName string, ids entity.Column, opts ...GetOption) (ResultSet, error) {
+	if c.Service == nil {
+		return nil, ErrClientNotReady
+	}
+
+	o := &getOption{}
+	for _, opt := range opts {
+		opt(o)
+	}
+
+	if len(o.outputFields) == 0 {
+		coll, err := c.DescribeCollection(ctx, collectionName)
+		if err != nil {
+			return nil, err
+		}
+		for _, f := range coll.Schema.Fields {
+			o.outputFields = append(o.outputFields, f.Name)
+		}
+	}
+
+	return c.QueryByPks(ctx, collectionName, o.partitionNames, ids, o.outputFields)
+}
+
 // QueryByPks query record by specified primary key(s)
 func (c *GrpcClient) QueryByPks(ctx context.Context, collectionName string, partitionNames []string, ids entity.Column, outputFields []string, opts ...SearchQueryOptionFunc) (ResultSet, error) {
 	if c.Service == nil {
@@ -259,6 +283,15 @@ func getPKField(schema *entity.Schema) *entity.Field {
 	return nil
 }
 
+func getVectorField(schema *entity.Schema) *entity.Field {
+	for _, f := range schema.Fields {
+		if f.DataType == entity.FieldTypeFloatVector || f.DataType == entity.FieldTypeBinaryVector {
+			return f
+		}
+	}
+	return nil
+}
+
 func prepareSearchRequest(collName string, partitions []string,
 	expr string, outputFields []string, vectors []entity.Vector, vectorField string,
 	metricType entity.MetricType, topK int, sp entity.SearchParam, opt *SearchQueryOption) (*server.SearchRequest, error) {
diff --git a/client/options.go b/client/options.go
index 0209f33..cdb66e5 100644
--- a/client/options.go
+++ b/client/options.go
@@ -214,3 +214,22 @@ func IsBackup() BulkInsertOption {
 		req.Options = entity.MapKvPairs(optionMap)
 	}
 }
+
+type getOption struct {
+	partitionNames []string
+	outputFields   []string
+}
+
+type GetOption func(o *getOption)
+
+func GetWithPartitions(partionNames ...string) GetOption {
+	return func(o *getOption) {
+		o.partitionNames = partionNames
+	}
+}
+
+func GetWithOutputFields(outputFields ...string) GetOption {
+	return func(o *getOption) {
+		o.outputFields = outputFields
+	}
+}


Test patch:
diff --git a/client/data_test.go b/client/data_test.go
index 71f3374..245ae1f 100644
--- a/client/data_test.go
+++ b/client/data_test.go
@@ -401,6 +401,32 @@ func (s *QuerySuite) SetupSuite() {
 		WithField(entity.NewField().WithName("vector").WithDataType(entity.FieldTypeFloatVector).WithDim(testVectorDim))
 }
 
+func (s *QuerySuite) GetFail() {
+	c := s.client
+	ctx, cancel := context.WithCancel((context.Background()))
+	defer cancel()
+
+	idCol := entity.NewColumnInt64("ID", []int64{1})
+	s.Run("service_not_ready", func() {
+		_, err := (&GrpcClient{}).Get(ctx, testCollectionName, idCol)
+		s.Error(err)
+		s.ErrorIs(err, ErrClientNotReady)
+	})
+
+	s.Run("ids_len_0", func() {
+		_, err := c.Get(ctx, testCollectionName, entity.NewColumnInt64("ID", []int64{}), GetWithOutputFields("ID"))
+		s.Error(err)
+	})
+
+	s.Run("describe_failed", func() {
+		defer s.resetMock()
+		s.setupDescribeCollectionError(common.ErrorCode_Success, errors.New("mock error"))
+
+		_, err := c.Get(ctx, testCollectionName, idCol)
+		s.Error(err)
+	})
+}
+
 func (s *QuerySuite) TestQueryByPksFail() {
 	c := s.client
 	ctx, cancel := context.WithCancel(context.Background())