Fix retrieve_collection bugs (#5575)

Signed-off-by: fishpenguin <kun.yu@zilliz.com>

internal/proxynode/task.go

@@ -1591,7 +1591,15 @@ func (rt *RetrieveTask) PostExecute(ctx context.Context) error {
 		}
 
 		availableQueryNodeNum := 0
-		for _, partialRetrieveResult := range retrieveResult {
+		rt.result = &milvuspb.RetrieveResults{
+			Status: &commonpb.Status{
+				ErrorCode: 0,
+			},
+			Ids:        &schemapb.IDs{},
+			FieldsData: make([]*schemapb.FieldData, 0),
+		}
+		for idx, partialRetrieveResult := range retrieveResult {
+			log.Debug("Index-" + strconv.Itoa(idx))
 			if partialRetrieveResult.Ids == nil {
 				reason += "ids is nil\n"
 				continue
@@ -1604,11 +1612,28 @@ func (rt *RetrieveTask) PostExecute(ctx context.Context) error {
 				}
 
 				if !intOk {
-					rt.result.Ids.IdField.(*schemapb.IDs_IntId).IntId.Data = append(rt.result.Ids.IdField.(*schemapb.IDs_IntId).IntId.Data, intIds.IntId.Data...)
+					if idsStr, ok := rt.result.Ids.IdField.(*schemapb.IDs_StrId); ok {
+						idsStr.StrId.Data = append(idsStr.StrId.Data, strIds.StrId.Data...)
+					} else {
+						rt.result.Ids.IdField = &schemapb.IDs_StrId{
+							StrId: &schemapb.StringArray{
+								Data: strIds.StrId.Data,
+							},
+						}
+					}
 				} else {
-					rt.result.Ids.IdField.(*schemapb.IDs_StrId).StrId.Data = append(rt.result.Ids.IdField.(*schemapb.IDs_StrId).StrId.Data, strIds.StrId.Data...)
-				}
+					if idsInt, ok := rt.result.Ids.IdField.(*schemapb.IDs_IntId); ok {
+						idsInt.IntId.Data = append(idsInt.IntId.Data, intIds.IntId.Data...)
+					} else {
+						rt.result.Ids.IdField = &schemapb.IDs_IntId{
+							IntId: &schemapb.LongArray{
+								Data: intIds.IntId.Data,
+							},
+						}
+					}
 
+				}
+				rt.result.FieldsData = append(rt.result.FieldsData, partialRetrieveResult.FieldsData...)
 			}
 			availableQueryNodeNum++
 		}
@@ -1624,13 +1649,6 @@ func (rt *RetrieveTask) PostExecute(ctx context.Context) error {
 			}
 			return nil
 		}
-
-		rt.result = &milvuspb.RetrieveResults{
-			Status: &commonpb.Status{
-				ErrorCode: commonpb.ErrorCode_Success,
-				Reason:    reason,
-			},
-		}
 	}
 
 	log.Info("Retrieve PostExecute done.",

internal/querynode/retrieve_collection.go

@@ -208,6 +208,7 @@ func (rc *retrieveCollection) doUnsolvedMsgRetrieve() {
 
 				if rm.EndTs() <= serviceTime {
 					retrieveMsg = append(retrieveMsg, rm)
+					continue
 				}
 				rc.addToUnsolvedMsg(rm)
 			}