Optimize segcore Reduce (#18902)

Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>
Signed-off-by: Nyudong.cai <yudong.cai@zilliz.com>

internal/core/src/segcore/Reduce.h

@@ -61,15 +61,26 @@ class ReduceHelper {
     FilterInvalidSearchResult(SearchResult* search_result);
 
     void
-    ReduceResultData(int slice_index);
+    FillPrimaryKey();
+
+    void
+    RefreshSearchResult();
+
+    void
+    FillEntryData();
+
+    int64_t
+    ReduceSearchResultForOneNQ(int64_t qi, int64_t topk, int64_t& result_offset);
+
+    void
+    ReduceResultData();
 
     std::vector<char>
-    GetSearchResultDataSlice(int slice_index_, int64_t result_count);
+    GetSearchResultDataSlice(int slice_index_);
 
  private:
     std::vector<int64_t> slice_topKs_;
     std::vector<int64_t> slice_nqs_;
-    int64_t unify_topK_;
     int64_t total_nq_;
     int64_t num_segments_;
     int64_t num_slices_;
@@ -77,10 +88,10 @@ class ReduceHelper {
     milvus::query::Plan* plan_;
     std::vector<SearchResult*>& search_results_;
 
-    //
-    std::vector<int32_t> nq_slice_offsets_;
-    std::vector<std::vector<int64_t>> final_search_records_;
-    std::vector<std::vector<int64_t>> final_real_topKs_;
+    std::vector<int64_t> slice_nqs_prefix_sum_;
+
+    // dim0: num_segments_; dim1: total_nq_; dim2: offset
+    std::vector<std::vector<std::vector<int64_t>>> final_search_records_;
 
     // output
     std::unique_ptr<SearchResultDataBlobs> search_result_data_blobs_;

internal/core/unittest/test_c_api.cpp

@@ -1347,25 +1347,16 @@ testReduceSearchWithExpr(int N, int topK, int num_queries) {
         auto suc = search_result_data.ParseFromArray(search_result_data_blobs->blobs[i].data(),
                                                      search_result_data_blobs->blobs[i].size());
         assert(suc);
-
-        assert(suc);
         assert(search_result_data.num_queries() == slice_nqs[i]);
         assert(search_result_data.top_k() == slice_topKs[i]);
-        assert(search_result_data.scores().size() == slice_topKs[i] * slice_nqs[i]);
-        assert(search_result_data.ids().int_id().data_size() == slice_topKs[i] * slice_nqs[i]);
+        assert(search_result_data.scores().size() == search_result_data.topks().at(0) * slice_nqs[i]);
+        assert(search_result_data.ids().int_id().data_size() == search_result_data.topks().at(0) * slice_nqs[i]);
 
-        // check topKs
+        // check real topks
         assert(search_result_data.topks().size() == slice_nqs[i]);
-        for (int j = 0; j < search_result_data.topks().size(); j++) {
-            assert(search_result_data.topks().at(j) == slice_topKs[i]);
+        for (auto real_topk : search_result_data.topks()) {
+            assert(real_topk <= slice_topKs[i]);
         }
-
-        // assert(search_result_data.scores().size() == slice_topKs[i] * slice_nqs[i]);
-        // assert(search_result_data.ids().int_id().data_size() == slice_topKs[i] * slice_nqs[i]);
-        // assert(search_result_data.top_k() == topK);
-        // assert(search_result_data.num_queries() == req_sizes[i]);
-        // assert(search_result_data.scores().size() == topK * req_sizes[i]);
-        // assert(search_result_data.ids().int_id().data_size() == topK * req_sizes[i]);
     }
 
     DeleteSearchResultDataBlobs(cSearchResultData);
@@ -1378,6 +1369,8 @@ testReduceSearchWithExpr(int N, int topK, int num_queries) {
 }
 
 TEST(CApiTest, ReduceSearchWithExpr) {
+    testReduceSearchWithExpr(2, 1, 1);
+    testReduceSearchWithExpr(2, 10, 10);
     testReduceSearchWithExpr(100, 1, 1);
     testReduceSearchWithExpr(100, 10, 10);
     testReduceSearchWithExpr(10000, 1, 1);