enhance: [cp 3.0]: improve query node search reduce performance (#51515)

pr: https://github.com/milvus-io/milvus/pull/51316
issue: https://github.com/milvus-io/milvus/issues/51315

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>

internal/core/src/segcore/plan_c.cpp

@@ -9,11 +9,12 @@
 // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 // or implied. See the License for the specific language governing permissions and limitations under the License
 
-#include <string.h>
+#include <algorithm>
 #include <exception>
 #include <memory>
 #include <optional>
 #include <string>
+#include <string.h>
 #include <vector>
 
 #include "NamedType/named_type_impl.hpp"
@@ -142,6 +143,12 @@ GetMetricType(CSearchPlan plan) {
     return strdup(metric_str.c_str());
 }
 
+bool
+HasTargetEntries(CSearchPlan plan) {
+    auto search_plan = static_cast<milvus::query::Plan*>(plan);
+    return !search_plan->target_entries_.empty();
+}
+
 void
 SetMetricType(CSearchPlan plan, const char* metric_type) {
     auto search_plan = static_cast<milvus::query::Plan*>(plan);

internal/core/src/segcore/plan_c.h

@@ -50,6 +50,9 @@ GetFieldID(CSearchPlan plan, int64_t* field_id);
 const char*
 GetMetricType(CSearchPlan plan);
 
+bool
+HasTargetEntries(CSearchPlan plan);
+
 void
 SetMetricType(CSearchPlan plan, const char* metric_type);
 

internal/core/src/segcore/search_result_export_c.cpp

@@ -948,6 +948,16 @@ FillOutputFieldsOrderedImpl(CSearchResult* search_results,
                 milvus::ThreadPoolPriority::MIDDLE);
             std::vector<std::future<void>> futures;
             futures.reserve(seg_results.size());
+            auto futures_guard = folly::makeGuard([&futures]() {
+                for (auto& f : futures) {
+                    if (f.valid()) {
+                        try {
+                            f.get();
+                        } catch (...) {
+                        }
+                    }
+                }
+            });
             for (auto& [seg_idx, seg_res] : seg_results) {
                 auto sr = static_cast<SearchResult*>(search_results[seg_idx]);
                 auto segment =
@@ -960,16 +970,6 @@ FillOutputFieldsOrderedImpl(CSearchResult* search_results,
                             plan, seg_res_ptr->temp_result, &op_ctx);
                     }));
             }
-            auto futures_guard = folly::makeGuard([&futures]() {
-                for (auto& f : futures) {
-                    if (f.valid()) {
-                        try {
-                            f.get();
-                        } catch (...) {
-                        }
-                    }
-                }
-            });
             for (auto& future : futures) {
                 future.get();
             }

internal/core/src/segcore/search_result_export_c_test.cpp

@@ -1150,7 +1150,9 @@ TEST(SearchResultExport, FillOutputFieldsOrdered_Basic) {
     auto plan_bytes = BuildSimpleVectorSearchPlan(vec_fid, /*topk=*/2);
     auto plan = milvus::query::CreateSearchPlanByExpr(
         schema, plan_bytes.data(), plan_bytes.size());
+    plan->target_entries_.push_back(pk_fid);
     plan->target_entries_.push_back(output_fid);
+    EXPECT_TRUE(HasTargetEntries(reinterpret_cast<CSearchPlan>(plan.get())));
 
     SearchResult sr;
     sr.segment_ = segment.get();
@@ -1175,13 +1177,33 @@ TEST(SearchResultExport, FillOutputFieldsOrdered_Basic) {
     milvus::proto::schema::SearchResultData result_data;
     ASSERT_TRUE(
         result_data.ParseFromArray(c_proto.proto_blob, c_proto.proto_size));
-    ASSERT_EQ(result_data.fields_data_size(), 1);
-    EXPECT_EQ(result_data.fields_data(0).field_id(), output_fid.get());
+    ASSERT_EQ(result_data.fields_data_size(), 2);
+    EXPECT_EQ(result_data.fields_data(0).field_id(), pk_fid.get());
     EXPECT_EQ(result_data.fields_data(0).scalars().long_data().data_size(), 2);
+    EXPECT_EQ(result_data.fields_data(1).field_id(), output_fid.get());
+    EXPECT_EQ(result_data.fields_data(1).scalars().long_data().data_size(), 2);
 
     free(const_cast<void*>(c_proto.proto_blob));
 }
 
+TEST(SearchResultExport, HasTargetEntries) {
+    using namespace milvus;
+
+    auto schema = std::make_shared<Schema>();
+    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
+    schema->set_primary_field_id(pk_fid);
+    auto vec_fid = schema->AddDebugField(
+        "fakevec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2);
+
+    auto plan_bytes = BuildSimpleVectorSearchPlan(vec_fid, /*topk=*/2);
+    auto plan = milvus::query::CreateSearchPlanByExpr(
+        schema, plan_bytes.data(), plan_bytes.size());
+    EXPECT_FALSE(HasTargetEntries(reinterpret_cast<CSearchPlan>(plan.get())));
+
+    plan->target_entries_.push_back(pk_fid);
+    EXPECT_TRUE(HasTargetEntries(reinterpret_cast<CSearchPlan>(plan.get())));
+}
+
 TEST(SearchResultExport,
      FillOutputFieldsOrdered_EmptyArrayOutputsPreserveElementType) {
     using namespace milvus;
@@ -1200,6 +1222,7 @@ TEST(SearchResultExport,
     auto plan_bytes = BuildSimpleVectorSearchPlan(vec_fid, /*topk=*/2);
     auto plan = milvus::query::CreateSearchPlanByExpr(
         schema, plan_bytes.data(), plan_bytes.size());
+    plan->target_entries_.push_back(pk_fid);
     plan->target_entries_.push_back(scalar_fid);
     plan->target_entries_.push_back(vec_fid);
     plan->target_entries_.push_back(array_fid);
@@ -1221,22 +1244,28 @@ TEST(SearchResultExport,
     milvus::proto::schema::SearchResultData result_data;
     ASSERT_TRUE(
         result_data.ParseFromArray(c_proto.proto_blob, c_proto.proto_size));
-    ASSERT_EQ(result_data.fields_data_size(), 4);
+    ASSERT_EQ(result_data.fields_data_size(), 5);
 
-    const auto& scalar_field = result_data.fields_data(0);
+    const auto& pk_field = result_data.fields_data(0);
+    EXPECT_EQ(pk_field.field_id(), pk_fid.get());
+    EXPECT_EQ(pk_field.type(), milvus::proto::schema::DataType::Int64);
+    EXPECT_TRUE(pk_field.has_scalars());
+    EXPECT_TRUE(pk_field.scalars().has_long_data());
+
+    const auto& scalar_field = result_data.fields_data(1);
     EXPECT_EQ(scalar_field.field_id(), scalar_fid.get());
     EXPECT_EQ(scalar_field.type(), milvus::proto::schema::DataType::Int64);
     EXPECT_TRUE(scalar_field.has_scalars());
     EXPECT_TRUE(scalar_field.scalars().has_long_data());
 
-    const auto& vector_field = result_data.fields_data(1);
+    const auto& vector_field = result_data.fields_data(2);
     EXPECT_EQ(vector_field.field_id(), vec_fid.get());
     EXPECT_EQ(vector_field.type(),
               milvus::proto::schema::DataType::FloatVector);
     EXPECT_TRUE(vector_field.has_vectors());
     EXPECT_TRUE(vector_field.vectors().has_float_vector());
 
-    const auto& array_field = result_data.fields_data(2);
+    const auto& array_field = result_data.fields_data(3);
     EXPECT_EQ(array_field.field_id(), array_fid.get());
     EXPECT_EQ(array_field.type(), milvus::proto::schema::DataType::Array);
     EXPECT_TRUE(array_field.has_scalars());
@@ -1244,7 +1273,7 @@ TEST(SearchResultExport,
     EXPECT_EQ(array_field.scalars().array_data().element_type(),
               milvus::proto::schema::DataType::Int64);
 
-    const auto& vector_array_field = result_data.fields_data(3);
+    const auto& vector_array_field = result_data.fields_data(4);
     EXPECT_EQ(vector_array_field.field_id(), vector_array_fid.get());
     EXPECT_EQ(vector_array_field.type(),
               milvus::proto::schema::DataType::ArrayOfVector);

internal/querynodev2/tasks/heap_merge_reduce.go

@@ -356,26 +356,6 @@ type stringElementDedupKey struct {
 	elementIndex int32
 }
 
-func int64DedupKey(e *mergeEntry) any {
-	if e.elementIdx == nil {
-		return e.idInt64Val()
-	}
-	return int64ElementDedupKey{
-		pk:           e.idInt64Val(),
-		elementIndex: e.elementIndexVal(),
-	}
-}
-
-func stringDedupKey(e *mergeEntry) any {
-	if e.elementIdx == nil {
-		return e.idStringVal()
-	}
-	return stringElementDedupKey{
-		pk:           e.idStringVal(),
-		elementIndex: e.elementIndexVal(),
-	}
-}
-
 // mergeChunkInt64Pk performs the k-way merge for one chunk with int64 PK.
 func mergeChunkInt64Pk(
 	pool memory.Allocator,
@@ -486,30 +466,54 @@ func mergeChunkStringPk(
 
 // mergeStandardInt64Pk performs standard k-way merge for one NQ (int64 PK).
 func mergeStandardInt64Pk(h *mergeHeapInt64Pk, topK int64) ([]int64, []float32, []segmentSource) {
-	dedupSet := make(map[any]struct{}, topK)
+	if h.Len() > 0 && (*h)[0].elementIdx != nil {
+		return mergeStandardInt64ElementPk(h, topK)
+	}
+
+	dedupSet := make(map[int64]struct{}, topK)
 	ids := make([]int64, 0, topK)
 	scores := make([]float32, 0, topK)
 	sources := make([]segmentSource, 0, topK)
 
 	for int64(len(ids)) < topK && h.Len() > 0 {
-		e := heap.Pop(h).(*mergeEntry)
+		e := (*h)[0]
 		pk := e.idInt64Val()
-		dedupKey := int64DedupKey(e)
-
-		if _, dup := dedupSet[dedupKey]; !dup {
+		if _, dup := dedupSet[pk]; !dup {
 			ids = append(ids, pk)
 			scores = append(scores, e.scoreVal())
-			dedupSet[dedupKey] = struct{}{}
+			dedupSet[pk] = struct{}{}
 			sources = append(sources, segmentSource{
 				InputIdx:    e.inputIdx,
 				SegOffset:   e.segOffsetVal(),
 				OriginalIdx: e.cursor,
 			})
 		}
+		h.advanceRoot()
+	}
+	return ids, scores, sources
+}
 
-		if e.advance() {
-			heap.Push(h, e)
+func mergeStandardInt64ElementPk(h *mergeHeapInt64Pk, topK int64) ([]int64, []float32, []segmentSource) {
+	dedupSet := make(map[int64ElementDedupKey]struct{}, topK)
+	ids := make([]int64, 0, topK)
+	scores := make([]float32, 0, topK)
+	sources := make([]segmentSource, 0, topK)
+
+	for int64(len(ids)) < topK && h.Len() > 0 {
+		e := (*h)[0]
+		pk := e.idInt64Val()
+		key := int64ElementDedupKey{pk: pk, elementIndex: e.elementIndexVal()}
+		if _, dup := dedupSet[key]; !dup {
+			ids = append(ids, pk)
+			scores = append(scores, e.scoreVal())
+			dedupSet[key] = struct{}{}
+			sources = append(sources, segmentSource{
+				InputIdx:    e.inputIdx,
+				SegOffset:   e.segOffsetVal(),
+				OriginalIdx: e.cursor,
+			})
 		}
+		h.advanceRoot()
 	}
 	return ids, scores, sources
 }
@@ -524,23 +528,22 @@ func mergeGroupByInt64Pk(
 		ids, scores, sources := mergeStandardInt64Pk(h, topK)
 		return ids, scores, sources, nil
 	}
-	totalLimit := topK * groupSize
-	dedupSet := make(map[any]struct{}, totalLimit)
-	counter := newCompositeGroupCounter(topK, groupSize)
+	if h.Len() > 0 && (*h)[0].elementIdx != nil {
+		return mergeGroupByInt64ElementPk(h, topK, groupSize, numGroupFields)
+	}
 
+	totalLimit := topK * groupSize
+	dedupSet := make(map[int64]struct{}, totalLimit)
+	counter := newCompositeGroupCounter(topK, groupSize)
 	ids := make([]int64, 0, totalLimit)
 	scores := make([]float32, 0, totalLimit)
 	sources := make([]segmentSource, 0, totalLimit)
 
 	for int64(len(ids)) < totalLimit && h.Len() > 0 {
-		e := heap.Pop(h).(*mergeEntry)
+		e := (*h)[0]
 		pk := e.idInt64Val()
-		dedupKey := int64DedupKey(e)
-
-		if _, dup := dedupSet[dedupKey]; dup {
-			if e.advance() {
-				heap.Push(h, e)
-			}
+		if _, dup := dedupSet[pk]; dup {
+			h.advanceRoot()
 			continue
 		}
 
@@ -549,15 +552,13 @@ func mergeGroupByInt64Pk(
 			return nil, nil, nil, err
 		}
 		if !counter.shouldAccept(values) {
-			if e.advance() {
-				heap.Push(h, e)
-			}
+			h.advanceRoot()
 			continue
 		}
 
 		ids = append(ids, pk)
 		scores = append(scores, e.scoreVal())
-		dedupSet[dedupKey] = struct{}{}
+		dedupSet[pk] = struct{}{}
 		sources = append(sources, segmentSource{
 			InputIdx:    e.inputIdx,
 			SegOffset:   e.segOffsetVal(),
@@ -566,40 +567,107 @@ func mergeGroupByInt64Pk(
 		if counter.allSaturated() {
 			break
 		}
+		h.advanceRoot()
+	}
+	return ids, scores, sources, nil
+}
 
-		if e.advance() {
-			heap.Push(h, e)
+func mergeGroupByInt64ElementPk(
+	h *mergeHeapInt64Pk,
+	topK, groupSize int64,
+	numGroupFields int,
+) ([]int64, []float32, []segmentSource, error) {
+	totalLimit := topK * groupSize
+	dedupSet := make(map[int64ElementDedupKey]struct{}, totalLimit)
+	counter := newCompositeGroupCounter(topK, groupSize)
+	ids := make([]int64, 0, totalLimit)
+	scores := make([]float32, 0, totalLimit)
+	sources := make([]segmentSource, 0, totalLimit)
+
+	for int64(len(ids)) < totalLimit && h.Len() > 0 {
+		e := (*h)[0]
+		pk := e.idInt64Val()
+		key := int64ElementDedupKey{pk: pk, elementIndex: e.elementIndexVal()}
+		if _, dup := dedupSet[key]; dup {
+			h.advanceRoot()
+			continue
 		}
+
+		values, err := extractCompositeGroupValues(e, numGroupFields)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		if !counter.shouldAccept(values) {
+			h.advanceRoot()
+			continue
+		}
+
+		ids = append(ids, pk)
+		scores = append(scores, e.scoreVal())
+		dedupSet[key] = struct{}{}
+		sources = append(sources, segmentSource{
+			InputIdx:    e.inputIdx,
+			SegOffset:   e.segOffsetVal(),
+			OriginalIdx: e.cursor,
+		})
+		if counter.allSaturated() {
+			break
+		}
+		h.advanceRoot()
 	}
 	return ids, scores, sources, nil
 }
 
 // mergeStandardStringPk performs standard k-way merge for one NQ (string PK).
 func mergeStandardStringPk(h *mergeHeapStringPk, topK int64) ([]string, []float32, []segmentSource) {
-	dedupSet := make(map[any]struct{}, topK)
+	if h.Len() > 0 && (*h)[0].elementIdx != nil {
+		return mergeStandardStringElementPk(h, topK)
+	}
+
+	dedupSet := make(map[string]struct{}, topK)
 	ids := make([]string, 0, topK)
 	scores := make([]float32, 0, topK)
 	sources := make([]segmentSource, 0, topK)
 
 	for int64(len(ids)) < topK && h.Len() > 0 {
-		e := heap.Pop(h).(*mergeEntry)
+		e := (*h)[0]
 		pk := e.idStringVal()
-		dedupKey := stringDedupKey(e)
-
-		if _, dup := dedupSet[dedupKey]; !dup {
+		if _, dup := dedupSet[pk]; !dup {
 			ids = append(ids, pk)
 			scores = append(scores, e.scoreVal())
-			dedupSet[dedupKey] = struct{}{}
+			dedupSet[pk] = struct{}{}
 			sources = append(sources, segmentSource{
 				InputIdx:    e.inputIdx,
 				SegOffset:   e.segOffsetVal(),
 				OriginalIdx: e.cursor,
 			})
 		}
+		h.advanceRoot()
+	}
+	return ids, scores, sources
+}
 
-		if e.advance() {
-			heap.Push(h, e)
+func mergeStandardStringElementPk(h *mergeHeapStringPk, topK int64) ([]string, []float32, []segmentSource) {
+	dedupSet := make(map[stringElementDedupKey]struct{}, topK)
+	ids := make([]string, 0, topK)
+	scores := make([]float32, 0, topK)
+	sources := make([]segmentSource, 0, topK)
+
+	for int64(len(ids)) < topK && h.Len() > 0 {
+		e := (*h)[0]
+		pk := e.idStringVal()
+		key := stringElementDedupKey{pk: pk, elementIndex: e.elementIndexVal()}
+		if _, dup := dedupSet[key]; !dup {
+			ids = append(ids, pk)
+			scores = append(scores, e.scoreVal())
+			dedupSet[key] = struct{}{}
+			sources = append(sources, segmentSource{
+				InputIdx:    e.inputIdx,
+				SegOffset:   e.segOffsetVal(),
+				OriginalIdx: e.cursor,
+			})
 		}
+		h.advanceRoot()
 	}
 	return ids, scores, sources
 }
@@ -614,23 +682,22 @@ func mergeGroupByStringPk(
 		ids, scores, sources := mergeStandardStringPk(h, topK)
 		return ids, scores, sources, nil
 	}
-	totalLimit := topK * groupSize
-	dedupSet := make(map[any]struct{}, totalLimit)
-	counter := newCompositeGroupCounter(topK, groupSize)
+	if h.Len() > 0 && (*h)[0].elementIdx != nil {
+		return mergeGroupByStringElementPk(h, topK, groupSize, numGroupFields)
+	}
 
+	totalLimit := topK * groupSize
+	dedupSet := make(map[string]struct{}, totalLimit)
+	counter := newCompositeGroupCounter(topK, groupSize)
 	ids := make([]string, 0, totalLimit)
 	scores := make([]float32, 0, totalLimit)
 	sources := make([]segmentSource, 0, totalLimit)
 
 	for int64(len(ids)) < totalLimit && h.Len() > 0 {
-		e := heap.Pop(h).(*mergeEntry)
+		e := (*h)[0]
 		pk := e.idStringVal()
-		dedupKey := stringDedupKey(e)
-
-		if _, dup := dedupSet[dedupKey]; dup {
-			if e.advance() {
-				heap.Push(h, e)
-			}
+		if _, dup := dedupSet[pk]; dup {
+			h.advanceRoot()
 			continue
 		}
 
@@ -639,15 +706,13 @@ func mergeGroupByStringPk(
 			return nil, nil, nil, err
 		}
 		if !counter.shouldAccept(values) {
-			if e.advance() {
-				heap.Push(h, e)
-			}
+			h.advanceRoot()
 			continue
 		}
 
 		ids = append(ids, pk)
 		scores = append(scores, e.scoreVal())
-		dedupSet[dedupKey] = struct{}{}
+		dedupSet[pk] = struct{}{}
 		sources = append(sources, segmentSource{
 			InputIdx:    e.inputIdx,
 			SegOffset:   e.segOffsetVal(),
@@ -656,10 +721,53 @@ func mergeGroupByStringPk(
 		if counter.allSaturated() {
 			break
 		}
+		h.advanceRoot()
+	}
+	return ids, scores, sources, nil
+}
 
-		if e.advance() {
-			heap.Push(h, e)
+func mergeGroupByStringElementPk(
+	h *mergeHeapStringPk,
+	topK, groupSize int64,
+	numGroupFields int,
+) ([]string, []float32, []segmentSource, error) {
+	totalLimit := topK * groupSize
+	dedupSet := make(map[stringElementDedupKey]struct{}, totalLimit)
+	counter := newCompositeGroupCounter(topK, groupSize)
+	ids := make([]string, 0, totalLimit)
+	scores := make([]float32, 0, totalLimit)
+	sources := make([]segmentSource, 0, totalLimit)
+
+	for int64(len(ids)) < totalLimit && h.Len() > 0 {
+		e := (*h)[0]
+		pk := e.idStringVal()
+		key := stringElementDedupKey{pk: pk, elementIndex: e.elementIndexVal()}
+		if _, dup := dedupSet[key]; dup {
+			h.advanceRoot()
+			continue
 		}
+
+		values, err := extractCompositeGroupValues(e, numGroupFields)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+		if !counter.shouldAccept(values) {
+			h.advanceRoot()
+			continue
+		}
+
+		ids = append(ids, pk)
+		scores = append(scores, e.scoreVal())
+		dedupSet[key] = struct{}{}
+		sources = append(sources, segmentSource{
+			InputIdx:    e.inputIdx,
+			SegOffset:   e.segOffsetVal(),
+			OriginalIdx: e.cursor,
+		})
+		if counter.allSaturated() {
+			break
+		}
+		h.advanceRoot()
 	}
 	return ids, scores, sources, nil
 }

internal/querynodev2/tasks/heap_merge_reduce_test.go

@@ -20,6 +20,7 @@ package tasks
 
 import (
 	"container/heap"
+	"fmt"
 	"testing"
 
 	"github.com/apache/arrow/go/v17/arrow"
@@ -261,6 +262,54 @@ func buildTestDFWithStringPKAndElementIndices(pool memory.Allocator, idsPerChunk
 	return builder.Build()
 }
 
+func buildTestDFWithStringPKElementIndicesAndGroupBy(
+	pool memory.Allocator,
+	idsPerChunk [][]string,
+	scoresPerChunk [][]float32,
+	elemIdxPerChunk [][]int32,
+	groupByPerChunk [][]int64,
+	groupByFieldName string,
+) *chain.DataFrame {
+	numChunks := len(idsPerChunk)
+	chunkSizes := make([]int64, numChunks)
+	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, elementIndicesCol, groupByFieldName}
+	collector := chain.NewChunkCollector(colNames, numChunks)
+
+	for i := 0; i < numChunks; i++ {
+		chunkSizes[i] = int64(len(idsPerChunk[i]))
+
+		idBuilder := array.NewStringBuilder(pool)
+		idBuilder.AppendValues(idsPerChunk[i], nil)
+		collector.Set(idFieldName, i, idBuilder.NewArray())
+		idBuilder.Release()
+
+		scoreBuilder := array.NewFloat32Builder(pool)
+		scoreBuilder.AppendValues(scoresPerChunk[i], nil)
+		collector.Set(scoreFieldName, i, scoreBuilder.NewArray())
+		scoreBuilder.Release()
+
+		elementBuilder := array.NewInt32Builder(pool)
+		elementBuilder.AppendValues(elemIdxPerChunk[i], nil)
+		collector.Set(elementIndicesCol, i, elementBuilder.NewArray())
+		elementBuilder.Release()
+
+		groupBuilder := array.NewInt64Builder(pool)
+		groupBuilder.AppendValues(groupByPerChunk[i], nil)
+		collector.Set(groupByFieldName, i, groupBuilder.NewArray())
+		groupBuilder.Release()
+
+		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
+	}
+
+	builder := chain.NewDataFrameBuilder()
+	builder.SetChunkSizes(chunkSizes)
+	for _, name := range colNames {
+		_ = builder.AddColumnFromChunks(name, collector.Consume(name))
+	}
+	collector.Release()
+	return builder.Build()
+}
+
 func buildTestDFWithSegOffset(pool memory.Allocator, idsPerChunk [][]int64, scoresPerChunk [][]float32, offsetsPerChunk [][]int64) *chain.DataFrame {
 	numChunks := len(idsPerChunk)
 	chunkSizes := make([]int64, numChunks)
@@ -1704,6 +1753,35 @@ func TestHeapMergeReduce_ElementIndicesDedup(t *testing.T) {
 	assert.Equal(t, int32(0), eiCol.Value(3))
 }
 
+func TestHeapMergeReduce_GroupByStringElementIndicesDedup(t *testing.T) {
+	pool := memory.NewGoAllocator()
+
+	df0 := buildTestDFWithStringPKElementIndicesAndGroupBy(pool,
+		[][]string{{"a", "a", "b"}},
+		[][]float32{{0.99, 0.97, 0.96}},
+		[][]int32{{0, 1, 0}},
+		[][]int64{{1, 1, 2}},
+		groupByCol)
+	defer df0.Release()
+
+	df1 := buildTestDFWithStringPKElementIndicesAndGroupBy(pool,
+		[][]string{{"a", "c"}},
+		[][]float32{{0.98, 0.95}},
+		[][]int32{{1, 0}},
+		[][]int64{{1, 3}},
+		groupByCol)
+	defer df1.Release()
+
+	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 2, &groupByOptions{GroupSize: 2})
+	require.NoError(t, err)
+	defer result.DF.Release()
+
+	ids := result.DF.Column(idFieldName).Chunk(0).(*array.String)
+	elements := result.DF.Column(elementIndicesCol).Chunk(0).(*array.Int32)
+	assert.Equal(t, []string{"a", "a", "b"}, []string{ids.Value(0), ids.Value(1), ids.Value(2)})
+	assert.Equal(t, []int32{0, 1, 0}, []int32{elements.Value(0), elements.Value(1), elements.Value(2)})
+}
+
 func TestHeapMergeReduce_GroupByElementIndicesDedup(t *testing.T) {
 	pool := memory.NewGoAllocator()
 
@@ -1788,3 +1866,354 @@ func TestHeapMergeReduce_NoElementIndices(t *testing.T) {
 	assert.False(t, result.DF.HasColumn(elementIndicesCol),
 		"no $element_indices column should exist when inputs lack it")
 }
+
+func TestMergeHeapAdvanceRoot(t *testing.T) {
+	pool := memory.NewGoAllocator()
+
+	t.Run("Int64", func(t *testing.T) {
+		dfs := []*chain.DataFrame{
+			buildTestDF(pool, [][]int64{{1, 4}}, [][]float32{{0.9, 0.5}}),
+			buildTestDF(pool, [][]int64{{2}}, [][]float32{{0.8}}),
+			buildTestDF(pool, [][]int64{{3, 5}}, [][]float32{{0.7, 0.4}}),
+		}
+		defer func() {
+			for _, df := range dfs {
+				df.Release()
+			}
+		}()
+
+		cols := resolveInputCols(dfs, nil, false)
+		entries, err := buildMergeEntries(cols, 0, false, false)
+		require.NoError(t, err)
+		h := mergeHeapInt64Pk(entries)
+		heap.Init(&h)
+
+		var got []int64
+		for h.Len() > 0 {
+			got = append(got, h[0].idInt64Val())
+			h.advanceRoot()
+		}
+		assert.Equal(t, []int64{1, 2, 3, 4, 5}, got)
+	})
+
+	t.Run("String", func(t *testing.T) {
+		dfs := []*chain.DataFrame{
+			buildTestDFWithStringPK(pool, [][]string{{"a", "d"}}, [][]float32{{0.9, 0.5}}),
+			buildTestDFWithStringPK(pool, [][]string{{"b"}}, [][]float32{{0.8}}),
+			buildTestDFWithStringPK(pool, [][]string{{"c", "e"}}, [][]float32{{0.7, 0.4}}),
+		}
+		defer func() {
+			for _, df := range dfs {
+				df.Release()
+			}
+		}()
+
+		cols := resolveInputCols(dfs, nil, false)
+		entries, err := buildMergeEntries(cols, 0, false, true)
+		require.NoError(t, err)
+		h := mergeHeapStringPk(entries)
+		heap.Init(&h)
+
+		var got []string
+		for h.Len() > 0 {
+			got = append(got, h[0].idStringVal())
+			h.advanceRoot()
+		}
+		assert.Equal(t, []string{"a", "b", "c", "d", "e"}, got)
+	})
+}
+
+func TestMergeHeapAdvanceRoot_NearEpsilonLegacyParity(t *testing.T) {
+	const (
+		high float32 = 0.0999999717
+		mid  float32 = 0.0999999642
+		low  float32 = 0.0999999568
+	)
+
+	pool := memory.NewGoAllocator()
+
+	t.Run("Int64", func(t *testing.T) {
+		dfs := []*chain.DataFrame{
+			buildTestDF(pool, [][]int64{{1, 7}}, [][]float32{{0.2, low}}),
+			buildTestDF(pool, [][]int64{{7}}, [][]float32{{mid}}),
+			buildTestDF(pool, [][]int64{{7}}, [][]float32{{high}}),
+		}
+		defer func() {
+			for _, df := range dfs {
+				df.Release()
+			}
+		}()
+
+		cols := resolveInputCols(dfs, nil, false)
+		entries, err := buildMergeEntries(cols, 0, false, false)
+		require.NoError(t, err)
+
+		legacyHeap := &mergeHeapInt64Pk{}
+		optimizedHeap := &mergeHeapInt64Pk{}
+		for _, entry := range entries {
+			legacyCopy := *entry
+			optimizedCopy := *entry
+			heap.Push(legacyHeap, &legacyCopy)
+			heap.Push(optimizedHeap, &optimizedCopy)
+		}
+
+		wantIDs, wantScores, wantSources := legacyMergeStandardInt64Pk(legacyHeap, 2)
+		gotIDs, gotScores, gotSources := mergeStandardInt64Pk(optimizedHeap, 2)
+
+		assert.Equal(t, []int64{1, 7}, wantIDs)
+		assert.Equal(t, []float32{0.2, high}, wantScores)
+		assert.Equal(t, wantIDs, gotIDs)
+		assert.Equal(t, wantScores, gotScores)
+		assert.Equal(t, wantSources, gotSources)
+	})
+
+	t.Run("String", func(t *testing.T) {
+		dfs := []*chain.DataFrame{
+			buildTestDFWithStringPK(pool, [][]string{{"first", "duplicate"}}, [][]float32{{0.2, low}}),
+			buildTestDFWithStringPK(pool, [][]string{{"duplicate"}}, [][]float32{{mid}}),
+			buildTestDFWithStringPK(pool, [][]string{{"duplicate"}}, [][]float32{{high}}),
+		}
+		defer func() {
+			for _, df := range dfs {
+				df.Release()
+			}
+		}()
+
+		cols := resolveInputCols(dfs, nil, false)
+		entries, err := buildMergeEntries(cols, 0, false, true)
+		require.NoError(t, err)
+
+		legacyHeap := &mergeHeapStringPk{}
+		optimizedHeap := &mergeHeapStringPk{}
+		for _, entry := range entries {
+			legacyCopy := *entry
+			optimizedCopy := *entry
+			heap.Push(legacyHeap, &legacyCopy)
+			heap.Push(optimizedHeap, &optimizedCopy)
+		}
+
+		wantIDs, wantScores, wantSources := legacyMergeStandardStringPk(legacyHeap, 2)
+		gotIDs, gotScores, gotSources := mergeStandardStringPk(optimizedHeap, 2)
+
+		assert.Equal(t, []string{"first", "duplicate"}, wantIDs)
+		assert.Equal(t, []float32{0.2, high}, wantScores)
+		assert.Equal(t, wantIDs, gotIDs)
+		assert.Equal(t, wantScores, gotScores)
+		assert.Equal(t, wantSources, gotSources)
+	})
+}
+
+func legacyMergeStandardInt64Pk(h *mergeHeapInt64Pk, topK int64) ([]int64, []float32, []segmentSource) {
+	dedupSet := make(map[any]struct{}, topK)
+	ids := make([]int64, 0, topK)
+	scores := make([]float32, 0, topK)
+	sources := make([]segmentSource, 0, topK)
+	for int64(len(ids)) < topK && h.Len() > 0 {
+		e := heap.Pop(h).(*mergeEntry)
+		pk := e.idInt64Val()
+		if _, dup := dedupSet[pk]; !dup {
+			ids = append(ids, pk)
+			scores = append(scores, e.scoreVal())
+			dedupSet[pk] = struct{}{}
+			sources = append(sources, segmentSource{
+				InputIdx:    e.inputIdx,
+				SegOffset:   e.segOffsetVal(),
+				OriginalIdx: e.cursor,
+			})
+		}
+		if e.advance() {
+			heap.Push(h, e)
+		}
+	}
+	return ids, scores, sources
+}
+
+func legacyMergeStandardStringPk(h *mergeHeapStringPk, topK int64) ([]string, []float32, []segmentSource) {
+	dedupSet := make(map[any]struct{}, topK)
+	ids := make([]string, 0, topK)
+	scores := make([]float32, 0, topK)
+	sources := make([]segmentSource, 0, topK)
+	for int64(len(ids)) < topK && h.Len() > 0 {
+		e := heap.Pop(h).(*mergeEntry)
+		pk := e.idStringVal()
+		if _, dup := dedupSet[pk]; !dup {
+			ids = append(ids, pk)
+			scores = append(scores, e.scoreVal())
+			dedupSet[pk] = struct{}{}
+			sources = append(sources, segmentSource{
+				InputIdx:    e.inputIdx,
+				SegOffset:   e.segOffsetVal(),
+				OriginalIdx: e.cursor,
+			})
+		}
+		if e.advance() {
+			heap.Push(h, e)
+		}
+	}
+	return ids, scores, sources
+}
+
+func benchmarkHeapMergeReduce(b *testing.B, stringPK bool) {
+	const (
+		numSegments = 16
+		nq          = 50
+		rowsPerNQ   = 500
+		topK        = 500
+	)
+	pool := memory.NewGoAllocator()
+	dfs := make([]*chain.DataFrame, numSegments)
+	for seg := 0; seg < numSegments; seg++ {
+		scores := make([][]float32, nq)
+		if stringPK {
+			ids := make([][]string, nq)
+			for q := 0; q < nq; q++ {
+				ids[q] = make([]string, rowsPerNQ)
+				scores[q] = make([]float32, rowsPerNQ)
+				for row := 0; row < rowsPerNQ; row++ {
+					ids[q][row] = fmt.Sprintf("%04d-%04d-%04d", q, seg, row)
+					scores[q][row] = float32(rowsPerNQ-row) + float32(numSegments-seg)/float32(numSegments+1)
+				}
+			}
+			dfs[seg] = buildTestDFWithStringPK(pool, ids, scores)
+		} else {
+			ids := make([][]int64, nq)
+			for q := 0; q < nq; q++ {
+				ids[q] = make([]int64, rowsPerNQ)
+				scores[q] = make([]float32, rowsPerNQ)
+				for row := 0; row < rowsPerNQ; row++ {
+					ids[q][row] = int64((q*numSegments+seg)*rowsPerNQ + row)
+					scores[q][row] = float32(rowsPerNQ-row) + float32(numSegments-seg)/float32(numSegments+1)
+				}
+			}
+			dfs[seg] = buildTestDF(pool, ids, scores)
+		}
+	}
+	defer func() {
+		for _, df := range dfs {
+			df.Release()
+		}
+	}()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		result, err := heapMergeReduce(pool, dfs, topK, nil)
+		if err != nil {
+			b.Fatal(err)
+		}
+		result.DF.Release()
+	}
+}
+
+func BenchmarkHeapMergeReduceInt64(b *testing.B) {
+	benchmarkHeapMergeReduce(b, false)
+}
+
+func BenchmarkHeapMergeReduceString(b *testing.B) {
+	benchmarkHeapMergeReduce(b, true)
+}
+
+func benchmarkMergeLoop(b *testing.B, stringPK, legacy bool) {
+	const (
+		numSegments = 16
+		rowsPerNQ   = 500
+		topK        = 500
+	)
+	pool := memory.NewGoAllocator()
+	var intEntries []*mergeEntry
+	var stringEntries []*mergeEntry
+	var arrays []arrow.Array
+	defer func() {
+		for _, arr := range arrays {
+			arr.Release()
+		}
+	}()
+
+	for seg := 0; seg < numSegments; seg++ {
+		scores := make([]float32, rowsPerNQ)
+		offsets := make([]int64, rowsPerNQ)
+		for row := 0; row < rowsPerNQ; row++ {
+			scores[row] = float32(rowsPerNQ-row) + float32(numSegments-seg)/float32(numSegments+1)
+			offsets[row] = int64(row)
+		}
+		scoreBuilder := array.NewFloat32Builder(pool)
+		scoreBuilder.AppendValues(scores, nil)
+		scoreArr := scoreBuilder.NewArray().(*array.Float32)
+		scoreBuilder.Release()
+		offsetBuilder := array.NewInt64Builder(pool)
+		offsetBuilder.AppendValues(offsets, nil)
+		offsetArr := offsetBuilder.NewArray().(*array.Int64)
+		offsetBuilder.Release()
+		arrays = append(arrays, scoreArr, offsetArr)
+
+		if stringPK {
+			ids := make([]string, rowsPerNQ)
+			for row := range ids {
+				ids[row] = fmt.Sprintf("%04d-%04d", seg, row)
+			}
+			idBuilder := array.NewStringBuilder(pool)
+			idBuilder.AppendValues(ids, nil)
+			idArr := idBuilder.NewArray().(*array.String)
+			idBuilder.Release()
+			arrays = append(arrays, idArr)
+			stringEntries = append(stringEntries, &mergeEntry{inputIdx: seg, idString: idArr, scoreArr: scoreArr, segOffsetArr: offsetArr})
+		} else {
+			ids := make([]int64, rowsPerNQ)
+			for row := range ids {
+				ids[row] = int64(seg*rowsPerNQ + row)
+			}
+			idBuilder := array.NewInt64Builder(pool)
+			idBuilder.AppendValues(ids, nil)
+			idArr := idBuilder.NewArray().(*array.Int64)
+			idBuilder.Release()
+			arrays = append(arrays, idArr)
+			intEntries = append(intEntries, &mergeEntry{inputIdx: seg, idInt64: idArr, scoreArr: scoreArr, segOffsetArr: offsetArr})
+		}
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if stringPK {
+			h := make(mergeHeapStringPk, len(stringEntries))
+			for j, entry := range stringEntries {
+				copy := *entry
+				h[j] = &copy
+			}
+			heap.Init(&h)
+			if legacy {
+				legacyMergeStandardStringPk(&h, topK)
+			} else {
+				mergeStandardStringPk(&h, topK)
+			}
+		} else {
+			h := make(mergeHeapInt64Pk, len(intEntries))
+			for j, entry := range intEntries {
+				copy := *entry
+				h[j] = &copy
+			}
+			heap.Init(&h)
+			if legacy {
+				legacyMergeStandardInt64Pk(&h, topK)
+			} else {
+				mergeStandardInt64Pk(&h, topK)
+			}
+		}
+	}
+}
+
+func BenchmarkMergeLoopInt64Legacy(b *testing.B) {
+	benchmarkMergeLoop(b, false, true)
+}
+
+func BenchmarkMergeLoopInt64Optimized(b *testing.B) {
+	benchmarkMergeLoop(b, false, false)
+}
+
+func BenchmarkMergeLoopStringLegacy(b *testing.B) {
+	benchmarkMergeLoop(b, true, true)
+}
+
+func BenchmarkMergeLoopStringOptimized(b *testing.B) {
+	benchmarkMergeLoop(b, true, false)
+}

internal/querynodev2/tasks/merge_entry.go

@@ -107,6 +107,65 @@ func (h *mergeHeapInt64Pk) Pop() interface{} {
 	return item
 }
 
+// advanceRoot consumes the current root and moves that entry to its next row.
+// It deliberately preserves the legacy heap.Pop -> advance -> heap.Push
+// ordering: first remove the current root and repair the remaining heap, then
+// advance and reinsert the entry if it still has rows. This two-phase repair is
+// required because the epsilon-based score comparator is not a strict weak
+// ordering, so advancing the root in place and performing a single sift-down
+// can produce a different result from the legacy merge path.
+func (h *mergeHeapInt64Pk) advanceRoot() {
+	entries := *h
+	entry := entries[0]
+
+	last := len(entries) - 1
+	if last == 0 {
+		entries[0] = nil
+		entries = entries[:0]
+	} else {
+		entries[0] = entries[last]
+		entries[last] = nil
+		entries = entries[:last]
+		siftDownInt64Pk(entries, 0)
+	}
+
+	if entry.advance() {
+		entries = append(entries, entry)
+		siftUpInt64Pk(entries, len(entries)-1)
+	}
+	*h = entries
+}
+
+func siftDownInt64Pk(h mergeHeapInt64Pk, root int) {
+	for {
+		left := root*2 + 1
+		if left >= len(h) {
+			return
+		}
+		best := left
+		right := left + 1
+		if right < len(h) && h[right].greaterInt64Pk(h[left]) {
+			best = right
+		}
+		if !h[best].greaterInt64Pk(h[root]) {
+			return
+		}
+		h[root], h[best] = h[best], h[root]
+		root = best
+	}
+}
+
+func siftUpInt64Pk(h mergeHeapInt64Pk, child int) {
+	for child > 0 {
+		parent := (child - 1) / 2
+		if !h[child].greaterInt64Pk(h[parent]) {
+			return
+		}
+		h[parent], h[child] = h[child], h[parent]
+		child = parent
+	}
+}
+
 // mergeHeapStringPk implements heap.Interface for max-heap with varchar PK.
 type mergeHeapStringPk []*mergeEntry
 
@@ -126,3 +185,56 @@ func (h *mergeHeapStringPk) Pop() interface{} {
 	*h = old[:n-1]
 	return item
 }
+
+// advanceRoot is the varchar counterpart of mergeHeapInt64Pk.advanceRoot.
+func (h *mergeHeapStringPk) advanceRoot() {
+	entries := *h
+	entry := entries[0]
+
+	last := len(entries) - 1
+	if last == 0 {
+		entries[0] = nil
+		entries = entries[:0]
+	} else {
+		entries[0] = entries[last]
+		entries[last] = nil
+		entries = entries[:last]
+		siftDownStringPk(entries, 0)
+	}
+
+	if entry.advance() {
+		entries = append(entries, entry)
+		siftUpStringPk(entries, len(entries)-1)
+	}
+	*h = entries
+}
+
+func siftDownStringPk(h mergeHeapStringPk, root int) {
+	for {
+		left := root*2 + 1
+		if left >= len(h) {
+			return
+		}
+		best := left
+		right := left + 1
+		if right < len(h) && h[right].greaterStringPk(h[left]) {
+			best = right
+		}
+		if !h[best].greaterStringPk(h[root]) {
+			return
+		}
+		h[root], h[best] = h[best], h[root]
+		root = best
+	}
+}
+
+func siftUpStringPk(h mergeHeapStringPk, child int) {
+	for child > 0 {
+		parent := (child - 1) / 2
+		if !h[child].greaterStringPk(h[parent]) {
+			return
+		}
+		h[parent], h[child] = h[child], h[parent]
+		child = parent
+	}
+}

internal/querynodev2/tasks/search_task_go_reduce.go

@@ -336,6 +336,13 @@ func lateMaterializeOutputFields(
 	sources [][]segmentSource,
 	searchResultData *schemapb.SearchResultData,
 ) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+	if !plan.HasTargetEntries() {
+		return nil
+	}
+
 	totalRows := 0
 	for _, chunk := range sources {
 		totalRows += len(chunk)

internal/querynodev2/tasks/search_task_go_reduce_test.go

@@ -484,8 +484,110 @@ func TestLateMaterializeOutputFields_NoOutputFields(t *testing.T) {
 	assert.Empty(t, searchResultData.FieldsData)
 }
 
+func TestLateMaterializeOutputFields_NoOutputFieldsCanceled(t *testing.T) {
+	ts := setupTestSegments(t, 1, 100, setupOpts{NQ: 1, TopK: 5})
+	defer ts.cleanup()
+
+	plan := ts.searchReq.Plan()
+	require.False(t, plan.HasTargetEntries())
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := lateMaterializeOutputFields(
+		ctx,
+		ts.searchResults,
+		plan,
+		nil,
+		&schemapb.SearchResultData{},
+	)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
+func TestLateMaterializeOutputFields_OnlyPrimaryKey(t *testing.T) {
+	const primaryKeyFieldID int64 = 109
+
+	ts := setupTestSegments(t, 1, 100, setupOpts{
+		NQ:             1,
+		TopK:           5,
+		OutputFieldIDs: []int64{primaryKeyFieldID},
+	})
+	defer ts.cleanup()
+
+	reduceResult, segDFs := runGoReducePipeline(t, ts)
+	defer func() {
+		reduceResult.DF.Release()
+		for _, df := range segDFs {
+			df.Release()
+		}
+	}()
+
+	searchResultData, err := marshalReduceResult(reduceResult)
+	require.NoError(t, err)
+
+	plan := ts.searchReq.Plan()
+	require.True(t, plan.HasTargetEntries())
+	require.NoError(t, lateMaterializeOutputFields(
+		context.Background(),
+		ts.searchResults,
+		plan,
+		reduceResult.Sources,
+		searchResultData,
+	))
+
+	require.Len(t, searchResultData.FieldsData, 1)
+	assert.Equal(t, primaryKeyFieldID, searchResultData.FieldsData[0].GetFieldId())
+	pkData := searchResultData.FieldsData[0].GetScalars().GetLongData().GetData()
+	require.NotNil(t, searchResultData.GetIds().GetIntId())
+	assert.Equal(t, searchResultData.GetIds().GetIntId().GetData(), pkData)
+}
+
+func TestLateMaterializeOutputFields_PreservesPrimaryKeyWithOtherFields(t *testing.T) {
+	const (
+		primaryKeyFieldID int64 = 109
+		scalarFieldID     int64 = 103
+	)
+
+	outputFieldIDs := []int64{primaryKeyFieldID, scalarFieldID}
+	ts := setupTestSegments(t, 1, 100, setupOpts{
+		NQ:             1,
+		TopK:           5,
+		OutputFieldIDs: outputFieldIDs,
+	})
+	defer ts.cleanup()
+
+	reduceResult, segDFs := runGoReducePipeline(t, ts)
+	defer func() {
+		reduceResult.DF.Release()
+		for _, df := range segDFs {
+			df.Release()
+		}
+	}()
+
+	searchResultData, err := marshalReduceResult(reduceResult)
+	require.NoError(t, err)
+
+	plan := ts.searchReq.Plan()
+	require.True(t, plan.HasTargetEntries())
+	require.NoError(t, lateMaterializeOutputFields(
+		context.Background(),
+		ts.searchResults,
+		plan,
+		reduceResult.Sources,
+		searchResultData,
+	))
+
+	require.Len(t, searchResultData.FieldsData, 2)
+	assert.Equal(t, primaryKeyFieldID, searchResultData.FieldsData[0].GetFieldId())
+	assert.Equal(t, scalarFieldID, searchResultData.FieldsData[1].GetFieldId())
+
+	pkData := searchResultData.FieldsData[0].GetScalars().GetLongData().GetData()
+	require.NotNil(t, searchResultData.GetIds().GetIntId())
+	assert.Equal(t, searchResultData.GetIds().GetIntId().GetData(), pkData)
+}
+
 func TestLateMaterializeOutputFields_EmptySources(t *testing.T) {
-	outputFieldIDs := []int64{103, 104}
+	outputFieldIDs := []int64{109, 103, 104}
 	ts := setupTestSegments(t, 1, 100, setupOpts{NQ: 1, TopK: 5, OutputFieldIDs: outputFieldIDs})
 	defer ts.cleanup()
 

internal/util/segcore/plan.go

@@ -75,6 +75,10 @@ func (plan *SearchPlan) GetMetricType() string {
 	return metricType
 }
 
+func (plan *SearchPlan) HasTargetEntries() bool {
+	return bool(C.HasTargetEntries(plan.cSearchPlan))
+}
+
 func (plan *SearchPlan) delete() {
 	C.DeleteSearchPlan(plan.cSearchPlan)
 }