Add topks in searchResult (milvus-io#581)

Signed-off-by: fishpenguin <[email protected]>

Author: fishpenguin

grpc-proto/gen/schema_pb2.py

@@ -120,4 +120,5 @@ message SearchResultData {
   repeated FieldData fields_data = 3;
   repeated float scores = 4;
   IDs ids = 5;
+  repeated int64 topks = 6;
 }

grpc-proto/schema.proto

@@ -355,10 +355,11 @@ def _pack(self, raw):
         self._nq = raw.results.num_queries
         self._topk = raw.results.top_k
         self._hits = []
+        offset = 0
         for i in range(self._nq):
             hit = schema_pb2.SearchResultData()
-            start_pos = i * self._topk
-            end_pos = (i + 1) * self._topk
+            start_pos = offset
+            end_pos = offset + raw.results.topks[i]
             hit.scores.append(raw.results.scores[start_pos: end_pos])
             hit.ids.append(raw.results.ids.int_id.data[start_pos: end_pos])
             for field_data in raw.result.fields_data:
@@ -385,6 +386,7 @@ def _pack(self, raw):
                         field_data.vectors.float_data.data[start_pos * dim: end_pos * dim])
                 hit.fields_data.append(field)
             self._hits.append(hit)
+            offset += raw.results.topks[i]
 
     def get__item(self, item):
         return Hits(self._hits[item], self._auto_id)
@@ -411,10 +413,11 @@ def _pack(self, raw_list):
             nq = raw.results.num_queries
             self._nq += nq
             self._topk = raw.results.top_k
+            offset = 0
             for i in range(nq):
                 hit = schema_pb2.SearchResultData()
-                start_pos = i * self._topk
-                end_pos = (i + 1) * self._topk
+                start_pos = offset
+                end_pos = offset + raw.results.topks[i]
                 hit.scores.extend(raw.results.scores[start_pos: end_pos])
                 hit.ids.int_id.data.extend(raw.results.ids.int_id.data[start_pos: end_pos])
                 for field_data in raw.results.fields_data:
@@ -441,6 +444,7 @@ def _pack(self, raw_list):
                                                                start_pos * dim: end_pos * dim])
                     hit.fields_data.append(field)
                 self._hits.append(hit)
+                offset += raw.result.topks[i]
 
     def get__item(self, item):
         return Hits(self._hits[item], self._auto_id)