diff --git a/internal/core/src/segcore/plan_c.cpp b/internal/core/src/segcore/plan_c.cpp index 02b8cd649e..2f05d831bd 100644 --- a/internal/core/src/segcore/plan_c.cpp +++ b/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 +#include #include #include #include #include +#include #include #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(plan); + return !search_plan->target_entries_.empty(); +} + void SetMetricType(CSearchPlan plan, const char* metric_type) { auto search_plan = static_cast(plan); diff --git a/internal/core/src/segcore/plan_c.h b/internal/core/src/segcore/plan_c.h index 3fcea0dac4..c20accf760 100644 --- a/internal/core/src/segcore/plan_c.h +++ b/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); diff --git a/internal/core/src/segcore/search_result_export_c.cpp b/internal/core/src/segcore/search_result_export_c.cpp index 16d049dceb..579592aff5 100644 --- a/internal/core/src/segcore/search_result_export_c.cpp +++ b/internal/core/src/segcore/search_result_export_c.cpp @@ -948,6 +948,16 @@ FillOutputFieldsOrderedImpl(CSearchResult* search_results, milvus::ThreadPoolPriority::MIDDLE); std::vector> 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(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(); } diff --git a/internal/core/src/segcore/search_result_export_c_test.cpp b/internal/core/src/segcore/search_result_export_c_test.cpp index 08d0580669..55d8da751e 100644 --- a/internal/core/src/segcore/search_result_export_c_test.cpp +++ b/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(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(c_proto.proto_blob)); } +TEST(SearchResultExport, HasTargetEntries) { + using namespace milvus; + + auto schema = std::make_shared(); + 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(plan.get()))); + + plan->target_entries_.push_back(pk_fid); + EXPECT_TRUE(HasTargetEntries(reinterpret_cast(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); diff --git a/internal/querynodev2/tasks/heap_merge_reduce.go b/internal/querynodev2/tasks/heap_merge_reduce.go index a387fc46b6..328e8f9c47 100644 --- a/internal/querynodev2/tasks/heap_merge_reduce.go +++ b/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 } diff --git a/internal/querynodev2/tasks/heap_merge_reduce_test.go b/internal/querynodev2/tasks/heap_merge_reduce_test.go index 97ec807242..3f927e97b8 100644 --- a/internal/querynodev2/tasks/heap_merge_reduce_test.go +++ b/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] = © + } + 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] = © + } + 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) +} diff --git a/internal/querynodev2/tasks/merge_entry.go b/internal/querynodev2/tasks/merge_entry.go index 1a96caa63e..be5c94f4e7 100644 --- a/internal/querynodev2/tasks/merge_entry.go +++ b/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 + } +} diff --git a/internal/querynodev2/tasks/search_task_go_reduce.go b/internal/querynodev2/tasks/search_task_go_reduce.go index 6f670b04bf..ea0afc1139 100644 --- a/internal/querynodev2/tasks/search_task_go_reduce.go +++ b/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) diff --git a/internal/querynodev2/tasks/search_task_go_reduce_test.go b/internal/querynodev2/tasks/search_task_go_reduce_test.go index 4cdc1e2b2e..99e91ecffb 100644 --- a/internal/querynodev2/tasks/search_task_go_reduce_test.go +++ b/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() diff --git a/internal/util/segcore/plan.go b/internal/util/segcore/plan.go index bbf2b41fe0..0df390f9e8 100644 --- a/internal/util/segcore/plan.go +++ b/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) }