From 9b549e0fa4db8b7627f663e334b547a929e498d9 Mon Sep 17 00:00:00 2001
From: yukun <kun.yu@zilliz.com>
Date: Thu, 24 Jun 2021 21:08:14 +0800
Subject: [PATCH] Add primaryKey to output fields (#6094)

Signed-off-by: fishpenguin <kun.yu@zilliz.com>
---
 internal/proxy/task.go | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/internal/proxy/task.go b/internal/proxy/task.go
index 78e3ed372b..afcc43fc21 100644
--- a/internal/proxy/task.go
+++ b/internal/proxy/task.go
@@ -2026,11 +2026,11 @@ func (rt *RetrieveTask) PreExecute(ctx context.Context) error {
 		return errors.New(errMsg)
 	}
 	rt.Ids = rt.retrieve.Ids
+	schema, err := globalMetaCache.GetCollectionSchema(ctx, collectionName)
+	if err != nil {
+		return err
+	}
 	if len(rt.retrieve.OutputFields) == 0 {
-		schema, err := globalMetaCache.GetCollectionSchema(ctx, collectionName)
-		if err != nil {
-			return err
-		}
 		for _, field := range schema.Fields {
 			if field.FieldID >= 100 {
 				rt.OutputFields = append(rt.OutputFields, field.Name)
@@ -2038,6 +2038,20 @@ func (rt *RetrieveTask) PreExecute(ctx context.Context) error {
 		}
 	} else {
 		rt.OutputFields = rt.retrieve.OutputFields
+		for _, field := range schema.Fields {
+			if field.IsPrimaryKey {
+				containPrimaryKey := false
+				for _, reqFields := range rt.retrieve.OutputFields {
+					if reqFields == field.Name {
+						containPrimaryKey = true
+					}
+				}
+				if !containPrimaryKey {
+					rt.OutputFields = append(rt.OutputFields, field.Name)
+				}
+				break
+			}
+		}
 	}
 
 	travelTimestamp := rt.retrieve.TravelTimestamp