/*
 * Licensed to the LF AI & Data foundation under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License 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.
 */

package tasks

import (
	"container/heap"
	"fmt"
	"testing"

	"github.com/apache/arrow/go/v17/arrow"
	"github.com/apache/arrow/go/v17/arrow/array"
	"github.com/apache/arrow/go/v17/arrow/memory"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/milvus-io/milvus/internal/util/function/chain"
	"github.com/milvus-io/milvus/internal/util/reduce"
)

func appendDefaultSegOffsetChunk(pool memory.Allocator, collector *chain.ChunkCollector, chunkIdx, rows int) {
	offsets := make([]int64, rows)
	for i := range offsets {
		offsets[i] = int64(i)
	}
	offsetBuilder := array.NewInt64Builder(pool)
	offsetBuilder.AppendValues(offsets, nil)
	collector.Set(segOffsetCol, chunkIdx, offsetBuilder.NewArray())
	offsetBuilder.Release()
}

// buildTestDF creates a simple test DataFrame with int64 PKs.
// ids and scores are per-chunk (NQ) arrays.
func buildTestDF(pool memory.Allocator, idsPerChunk [][]int64, scoresPerChunk [][]float32) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	builder := chain.NewDataFrameBuilder()
	colNames := []string{idFieldName, scoreFieldName, segOffsetCol}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	collector.Release()

	return builder.Build()
}

func buildTestDFWithoutSegOffset(pool memory.Allocator, idsPerChunk [][]int64, scoresPerChunk [][]float32) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	builder := chain.NewDataFrameBuilder()
	colNames := []string{idFieldName, scoreFieldName}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()
	}

	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	collector.Release()

	return builder.Build()
}

// buildTestDFWithGroupBy creates a test DataFrame with int64 PKs and int64 group-by values.
func buildTestDFWithGroupBy(pool memory.Allocator, idsPerChunk [][]int64, scoresPerChunk [][]float32, groupByPerChunk [][]int64, groupByFieldName string) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, groupByFieldName}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()

		gbBuilder := array.NewInt64Builder(pool)
		gbBuilder.AppendValues(groupByPerChunk[i], nil)
		collector.Set(groupByFieldName, i, gbBuilder.NewArray())
		gbBuilder.Release()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(groupByFieldName, collector.Consume(groupByFieldName))
	collector.Release()

	return builder.Build()
}

// buildTestDFWithStringPKAndGroupBy creates a test DataFrame with varchar PKs
// and int64 group-by values.
func buildTestDFWithStringPKAndGroupBy(pool memory.Allocator, idsPerChunk [][]string, scoresPerChunk [][]float32, groupByPerChunk [][]int64, groupByFieldName string) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, 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()

		gbBuilder := array.NewInt64Builder(pool)
		gbBuilder.AppendValues(groupByPerChunk[i], nil)
		collector.Set(groupByFieldName, i, gbBuilder.NewArray())
		gbBuilder.Release()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(groupByFieldName, collector.Consume(groupByFieldName))
	collector.Release()

	return builder.Build()
}

func buildTestDFWithStringPK(pool memory.Allocator, idsPerChunk [][]string, scoresPerChunk [][]float32) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol}
	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()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	collector.Release()

	return builder.Build()
}

func buildTestDFWithStringPKAndElementIndices(pool memory.Allocator, idsPerChunk [][]string, scoresPerChunk [][]float32, elemIdxPerChunk [][]int32) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, elementIndicesCol}
	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()

		eiBuilder := array.NewInt32Builder(pool)
		eiBuilder.AppendValues(elemIdxPerChunk[i], nil)
		collector.Set(elementIndicesCol, i, eiBuilder.NewArray())
		eiBuilder.Release()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(elementIndicesCol, collector.Consume(elementIndicesCol))
	collector.Release()

	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)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()

		offsetBuilder := array.NewInt64Builder(pool)
		offsetBuilder.AppendValues(offsetsPerChunk[i], nil)
		collector.Set(segOffsetCol, i, offsetBuilder.NewArray())
		offsetBuilder.Release()
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	collector.Release()

	return builder.Build()
}

// buildTestDFWithElementIndices creates a test DataFrame with int64 PKs and
// int32 element indices (simulating element-level search results).
func buildTestDFWithElementIndices(pool memory.Allocator, idsPerChunk [][]int64, scoresPerChunk [][]float32, elemIdxPerChunk [][]int32) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, elementIndicesCol}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()

		eiBuilder := array.NewInt32Builder(pool)
		eiBuilder.AppendValues(elemIdxPerChunk[i], nil)
		collector.Set(elementIndicesCol, i, eiBuilder.NewArray())
		eiBuilder.Release()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(elementIndicesCol, collector.Consume(elementIndicesCol))
	collector.Release()

	return builder.Build()
}

func buildTestDFWithElementIndicesAndGroupBy(
	pool memory.Allocator,
	idsPerChunk [][]int64,
	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.NewInt64Builder(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()

		eiBuilder := array.NewInt32Builder(pool)
		eiBuilder.AppendValues(elemIdxPerChunk[i], nil)
		collector.Set(elementIndicesCol, i, eiBuilder.NewArray())
		eiBuilder.Release()

		gbBuilder := array.NewInt64Builder(pool)
		gbBuilder.AppendValues(groupByPerChunk[i], nil)
		collector.Set(groupByFieldName, i, gbBuilder.NewArray())
		gbBuilder.Release()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(elementIndicesCol, collector.Consume(elementIndicesCol))
	_ = builder.AddColumnFromChunks(groupByFieldName, collector.Consume(groupByFieldName))
	collector.Release()

	return builder.Build()
}

func TestHeapMergeReduce_BasicMerge(t *testing.T) {
	pool := memory.NewGoAllocator()

	// Segment 0: NQ=1, results [id=1,score=0.9], [id=2,score=0.7], [id=3,score=0.5]
	df0 := buildTestDF(pool, [][]int64{{1, 2, 3}}, [][]float32{{0.9, 0.7, 0.5}})
	defer df0.Release()

	// Segment 1: NQ=1, results [id=4,score=0.8], [id=5,score=0.6], [id=6,score=0.4]
	df1 := buildTestDF(pool, [][]int64{{4, 5, 6}}, [][]float32{{0.8, 0.6, 0.4}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 4, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Expected merged order: id=1(0.9), id=4(0.8), id=2(0.7), id=5(0.6)
	assert.Equal(t, 1, result.DF.NumChunks())
	assert.Equal(t, int64(4), result.DF.NumRows())

	idCol := result.DF.Column(idFieldName)
	scoreCol := result.DF.Column(scoreFieldName)

	ids := idCol.Chunk(0).(*array.Int64)
	scores := scoreCol.Chunk(0).(*array.Float32)

	assert.Equal(t, int64(1), ids.Value(0))
	assert.Equal(t, int64(4), ids.Value(1))
	assert.Equal(t, int64(2), ids.Value(2))
	assert.Equal(t, int64(5), ids.Value(3))

	assert.InDelta(t, float32(0.9), scores.Value(0), 0.001)
	assert.InDelta(t, float32(0.8), scores.Value(1), 0.001)
	assert.InDelta(t, float32(0.7), scores.Value(2), 0.001)
	assert.InDelta(t, float32(0.6), scores.Value(3), 0.001)
}

func TestHeapMergeReduce_PKDedup(t *testing.T) {
	pool := memory.NewGoAllocator()

	// Same PK=1 in both segments, different scores
	df0 := buildTestDF(pool, [][]int64{{1, 2}}, [][]float32{{0.9, 0.7}})
	defer df0.Release()

	df1 := buildTestDF(pool, [][]int64{{1, 3}}, [][]float32{{0.8, 0.6}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 3, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// id=1 from df0 (score=0.9) wins, id=1 from df1 (score=0.8) is deduped
	// Expected: id=1(0.9), id=2(0.7), id=3(0.6)
	assert.Equal(t, int64(3), result.DF.NumRows())

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0))
	assert.Equal(t, int64(2), ids.Value(1))
	assert.Equal(t, int64(3), ids.Value(2))
}

func TestHeapMergeReduce_EqualScoreDeterminism(t *testing.T) {
	pool := memory.NewGoAllocator()

	// heapMergeReduce expects inputs already normalized (score DESC, ties by
	// PK ASC) — the C++ exporter does this. Tests must mirror that contract.
	df0 := buildTestDF(pool, [][]int64{{3, 5}}, [][]float32{{0.9, 0.9}})
	defer df0.Release()

	df1 := buildTestDF(pool, [][]int64{{1, 7}}, [][]float32{{0.9, 0.9}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 4, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Equal score → PK ASC: id=1, id=3, id=5, id=7
	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0))
	assert.Equal(t, int64(3), ids.Value(1))
	assert.Equal(t, int64(5), ids.Value(2))
	assert.Equal(t, int64(7), ids.Value(3))
}

// TestHeapMergeReduce_SamePKEqualScore verifies that when the same PK appears
// in multiple segments with equal scores, dedup deterministically keeps one
// entry. With PK ASC tie-breaking, the "first seen" pop picks the sole row,
// and the duplicate from the later segment is dropped.
func TestHeapMergeReduce_SamePKEqualScore(t *testing.T) {
	pool := memory.NewGoAllocator()

	// Both segments have pk=1 at the same score (0.9). Dedup must drop one.
	df0 := buildTestDF(pool, [][]int64{{1, 2}}, [][]float32{{0.9, 0.7}})
	defer df0.Release()

	df1 := buildTestDF(pool, [][]int64{{1, 3}}, [][]float32{{0.9, 0.6}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 5, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Expect 3 distinct rows: pk=1 (once), pk=2, pk=3. Not 4.
	assert.Equal(t, int64(3), result.DF.NumRows())
	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0)) // pk=1 at 0.9 (from whichever segment)
	assert.Equal(t, int64(2), ids.Value(1)) // pk=2 at 0.7
	assert.Equal(t, int64(3), ids.Value(2)) // pk=3 at 0.6
}

func TestHeapMergeReduce_MultiNQ(t *testing.T) {
	pool := memory.NewGoAllocator()

	// NQ=2: chunk0=[id=1,0.9 | id=2,0.7], chunk1=[id=10,0.8 | id=11,0.6]
	df0 := buildTestDF(pool,
		[][]int64{{1, 2}, {10, 11}},
		[][]float32{{0.9, 0.7}, {0.8, 0.6}})
	defer df0.Release()

	df1 := buildTestDF(pool,
		[][]int64{{3, 4}, {12, 13}},
		[][]float32{{0.85, 0.65}, {0.75, 0.55}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 3, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	assert.Equal(t, 2, result.DF.NumChunks())

	// Chunk 0: id=1(0.9), id=3(0.85), id=2(0.7)
	ids0 := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids0.Value(0))
	assert.Equal(t, int64(3), ids0.Value(1))
	assert.Equal(t, int64(2), ids0.Value(2))

	// Chunk 1: id=10(0.8), id=12(0.75), id=11(0.6)
	ids1 := result.DF.Column(idFieldName).Chunk(1).(*array.Int64)
	assert.Equal(t, int64(10), ids1.Value(0))
	assert.Equal(t, int64(12), ids1.Value(1))
	assert.Equal(t, int64(11), ids1.Value(2))
}

// TestHeapMergeReduce_RaggedPerNQ verifies that heapMergeReduce handles
// inputs where per-NQ chunk sizes differ both across NQs within a segment
// AND across segments for the same NQ. This mirrors the post-Truncate
// (Global Refine) layout where survivors per (segment, NQ) are uneven, and
// the GroupBy path where SearchGroupByOperator emits variable-size chunks.
// The invariant heapMergeReduce relies on is only "same number of chunks",
// not "same chunk size".
func TestHeapMergeReduce_RaggedPerNQ(t *testing.T) {
	pool := memory.NewGoAllocator()

	// 3 NQs across 2 segments. Per-NQ row counts (shown as seg0 / seg1):
	//   NQ0: 3 / 1
	//   NQ1: 0 / 4   (seg0 empty for this NQ)
	//   NQ2: 2 / 2
	df0 := buildTestDF(pool,
		[][]int64{{1, 2, 3}, {}, {7, 8}},
		[][]float32{{0.9, 0.8, 0.5}, {}, {0.6, 0.4}})
	defer df0.Release()

	df1 := buildTestDF(pool,
		[][]int64{{10}, {20, 21, 22, 23}, {30, 31}},
		[][]float32{{0.85}, {0.95, 0.75, 0.55, 0.35}, {0.65, 0.45}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 5, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	assert.Equal(t, 3, result.DF.NumChunks(), "must preserve NQ count")

	// NQ0: seg0 contributes 3, seg1 contributes 1, topK=5 → 4 rows.
	// Global score-sorted order: 0.9, 0.85, 0.8, 0.5
	ids0 := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, 4, ids0.Len())
	assert.Equal(t, int64(1), ids0.Value(0))
	assert.Equal(t, int64(10), ids0.Value(1))
	assert.Equal(t, int64(2), ids0.Value(2))
	assert.Equal(t, int64(3), ids0.Value(3))

	// NQ1: seg0 empty, seg1 contributes 4 → 4 rows (topK=5 isn't hit).
	ids1 := result.DF.Column(idFieldName).Chunk(1).(*array.Int64)
	assert.Equal(t, 4, ids1.Len())
	assert.Equal(t, int64(20), ids1.Value(0))
	assert.Equal(t, int64(21), ids1.Value(1))
	assert.Equal(t, int64(22), ids1.Value(2))
	assert.Equal(t, int64(23), ids1.Value(3))

	// NQ2: 2 + 2 = 4 rows. Order: 0.65, 0.6, 0.45, 0.4
	ids2 := result.DF.Column(idFieldName).Chunk(2).(*array.Int64)
	assert.Equal(t, 4, ids2.Len())
	assert.Equal(t, int64(30), ids2.Value(0))
	assert.Equal(t, int64(7), ids2.Value(1))
	assert.Equal(t, int64(31), ids2.Value(2))
	assert.Equal(t, int64(8), ids2.Value(3))
}

func TestHeapMergeReduce_EmptyInput(t *testing.T) {
	pool := memory.NewGoAllocator()

	// One segment has results, one is empty
	df0 := buildTestDF(pool, [][]int64{{1, 2}}, [][]float32{{0.9, 0.7}})
	defer df0.Release()

	df1 := buildTestDF(pool, [][]int64{{}}, [][]float32{{}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 3, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Only results from df0
	assert.Equal(t, int64(2), result.DF.NumRows())
	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0))
	assert.Equal(t, int64(2), ids.Value(1))
}

func TestHeapMergeReduce_SingleSegment(t *testing.T) {
	pool := memory.NewGoAllocator()

	// Inputs must be pre-sorted (score DESC, ties by PK ASC).
	df0 := buildTestDF(pool, [][]int64{{1, 2, 3}}, [][]float32{{0.9, 0.7, 0.5}})
	defer df0.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0}, 2, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Should sort by score DESC and take top 2
	assert.Equal(t, int64(2), result.DF.NumRows())
	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0)) // score=0.9
	assert.Equal(t, int64(2), ids.Value(1)) // score=0.7
}

func TestHeapMergeReduce_StringPKStandard(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithStringPK(pool,
		[][]string{{"b", "d"}},
		[][]float32{{0.9, 0.7}})
	defer df0.Release()

	df1 := buildTestDFWithStringPK(pool,
		[][]string{{"a", "b", "c"}},
		[][]float32{{0.9, 0.85, 0.6}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 4, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.String)
	require.Equal(t, 4, ids.Len())
	assert.Equal(t, []string{"a", "b", "d", "c"}, []string{
		ids.Value(0), ids.Value(1), ids.Value(2), ids.Value(3),
	})
}

func TestHeapMergeReduce_StringPKElementIndicesDedup(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithStringPKAndElementIndices(pool,
		[][]string{{"doc", "doc", "z"}},
		[][]float32{{0.99, 0.97, 0.95}},
		[][]int32{{0, 1, 0}})
	defer df0.Release()

	df1 := buildTestDFWithStringPKAndElementIndices(pool,
		[][]string{{"doc", "a"}},
		[][]float32{{0.98, 0.96}},
		[][]int32{{1, 0}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 4, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.String)
	ei := result.DF.Column(elementIndicesCol).Chunk(0).(*array.Int32)
	require.Equal(t, 4, ids.Len())
	assert.Equal(t, []string{"doc", "doc", "a", "z"}, []string{
		ids.Value(0), ids.Value(1), ids.Value(2), ids.Value(3),
	})
	assert.Equal(t, []int32{0, 1, 0, 0}, []int32{
		ei.Value(0), ei.Value(1), ei.Value(2), ei.Value(3),
	})
}

func TestHeapMergeReduce_GroupByWithoutColumnsFallsBackToStandard(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDF(pool,
		[][]int64{{1, 2, 3}},
		[][]float32{{0.9, 0.8, 0.7}})
	defer df0.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0}, 2, &groupByOptions{
		GroupSize: 2,
	})
	require.NoError(t, err)
	defer result.DF.Release()

	assert.False(t, result.DF.HasColumn(groupByCol))
	assert.Equal(t, []int64{2}, result.DF.ChunkSizes())
	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, []int64{1, 2}, []int64{ids.Value(0), ids.Value(1)})
}

func TestHeapMergeReduce_SegOffsetSourceTracking(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithSegOffset(pool,
		[][]int64{{1, 3}},
		[][]float32{{0.9, 0.7}},
		[][]int64{{101, 103}})
	defer df0.Release()

	df1 := buildTestDFWithSegOffset(pool,
		[][]int64{{2}},
		[][]float32{{0.8}},
		[][]int64{{202}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 3, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	require.Len(t, result.Sources, 1)
	require.Len(t, result.Sources[0], 3)
	assert.Equal(t, []int64{101, 202, 103}, []int64{
		result.Sources[0][0].SegOffset,
		result.Sources[0][1].SegOffset,
		result.Sources[0][2].SegOffset,
	})
	assert.Equal(t, []int{0, 0, 1}, []int{
		result.Sources[0][0].OriginalIdx,
		result.Sources[0][1].OriginalIdx,
		result.Sources[0][2].OriginalIdx,
	})
}

func TestHeapMergeReduce_GroupByBasic(t *testing.T) {
	pool := memory.NewGoAllocator()

	// Segment 0: group A has 3 items, group B has 1
	df0 := buildTestDFWithGroupBy(pool,
		[][]int64{{1, 2, 3, 4}},
		[][]float32{{0.9, 0.8, 0.7, 0.6}},
		[][]int64{{100, 100, 100, 200}}, // group 100 and 200
		groupByCol)
	defer df0.Release()

	// Segment 1: group B has 2 items
	df1 := buildTestDFWithGroupBy(pool,
		[][]int64{{5, 6}},
		[][]float32{{0.85, 0.75}},
		[][]int64{{200, 200}},
		groupByCol)
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 2, &groupByOptions{
		GroupSize: 2,
	})
	require.NoError(t, err)
	defer result.DF.Release()

	// topK=2, groupSize=2 → max 4 results
	// group 100: id=1(0.9), id=2(0.8) → 2 items (full)
	// group 200: id=5(0.85), id=4(0.6) or id=6(0.75) → 2 items (full)
	// Expected order by score: 1(0.9), 5(0.85), 2(0.8), 6(0.75)
	assert.True(t, result.DF.NumRows() <= 4)

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0))
	assert.Equal(t, int64(5), ids.Value(1))
	assert.Equal(t, int64(2), ids.Value(2))
	assert.Equal(t, int64(6), ids.Value(3))

	// Verify group_by column exists
	assert.True(t, result.DF.HasColumn(groupByCol))
}

func TestHeapMergeReduce_GroupByMaxGroups(t *testing.T) {
	pool := memory.NewGoAllocator()

	// 4 groups, but topK=2 → only 2 groups allowed
	df0 := buildTestDFWithGroupBy(pool,
		[][]int64{{1, 2, 3, 4}},
		[][]float32{{0.9, 0.8, 0.7, 0.6}},
		[][]int64{{10, 20, 30, 40}},
		groupByCol)
	defer df0.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0}, 2, &groupByOptions{
		GroupSize: 1,
	})
	require.NoError(t, err)
	defer result.DF.Release()

	// topK=2, groupSize=1 → max 2 results, 2 groups
	assert.Equal(t, int64(2), result.DF.NumRows())
	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	assert.Equal(t, int64(1), ids.Value(0)) // group 10
	assert.Equal(t, int64(2), ids.Value(1)) // group 20
}

func TestHeapMergeReduce_GroupByCompositeKey(t *testing.T) {
	pool := memory.NewGoAllocator()
	groupByCol0 := groupByColumnName(105)
	groupByCol1 := groupByColumnName(106)

	df0 := buildTestDFWithTwoGroupByColumns(pool,
		[][]int64{{1, 2, 3, 4}},
		[][]float32{{0.95, 0.9, 0.85, 0.8}},
		[][]int64{{10, 10, 10, 20}},
		[][]string{{"a", "b", "a", "a"}},
		groupByCol0,
		groupByCol1)
	defer df0.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0}, 3, &groupByOptions{
		GroupSize: 1,
		Columns:   []string{groupByCol0, groupByCol1},
	})
	require.NoError(t, err)
	defer result.DF.Release()

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	require.Equal(t, 3, ids.Len())
	assert.Equal(t, int64(1), ids.Value(0)) // (10, a)
	assert.Equal(t, int64(2), ids.Value(1)) // (10, b), distinct from (10, a)
	assert.Equal(t, int64(4), ids.Value(2)) // (20, a); id=3 is duplicate composite group

	require.True(t, result.DF.HasColumn(groupByCol0))
	require.True(t, result.DF.HasColumn(groupByCol1))
	gb0 := result.DF.Column(groupByCol0).Chunk(0).(*array.Int64)
	gb1 := result.DF.Column(groupByCol1).Chunk(0).(*array.String)
	assert.Equal(t, []int64{10, 10, 20}, []int64{gb0.Value(0), gb0.Value(1), gb0.Value(2)})
	assert.Equal(t, []string{"a", "b", "a"}, []string{gb1.Value(0), gb1.Value(1), gb1.Value(2)})
}

func TestHeapMergeReduce_GroupByNullKey(t *testing.T) {
	pool := memory.NewGoAllocator()

	gb0Builder := array.NewInt64Builder(pool)
	gb0Builder.AppendValues([]int64{0, 0, 10, 10}, []bool{false, false, true, true})
	gb0 := gb0Builder.NewArray()
	gb0Builder.Release()
	defer gb0.Release()

	gb1Builder := array.NewInt64Builder(pool)
	gb1Builder.AppendValues([]int64{0, 10}, []bool{false, true})
	gb1 := gb1Builder.NewArray()
	gb1Builder.Release()
	defer gb1.Release()

	df0 := buildTestDFWithGroupByArr(pool,
		[][]int64{{1, 2, 3, 4}},
		[][]float32{{0.9, 0.8, 0.7, 0.6}},
		[]arrow.Array{gb0})
	defer df0.Release()

	df1 := buildTestDFWithGroupByArr(pool,
		[][]int64{{5, 6}},
		[][]float32{{0.85, 0.75}},
		[]arrow.Array{gb1})
	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.Int64)
	require.Equal(t, 4, ids.Len())
	assert.Equal(t, []int64{1, 5, 6, 3}, []int64{
		ids.Value(0), ids.Value(1), ids.Value(2), ids.Value(3),
	})

	gb := result.DF.Column(groupByCol).Chunk(0).(*array.Int64)
	require.Equal(t, 4, gb.Len())
	assert.True(t, gb.IsNull(0), "first null group row should be preserved")
	assert.True(t, gb.IsNull(1), "second null group row should be preserved")
	assert.False(t, gb.IsNull(2))
	assert.False(t, gb.IsNull(3))
	assert.Equal(t, int64(10), gb.Value(2))
	assert.Equal(t, int64(10), gb.Value(3))
}

func TestHeapMergeReduce_GroupByStringPK(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithStringPKAndGroupBy(pool,
		[][]string{{"a", "b", "c", "d"}},
		[][]float32{{0.9, 0.8, 0.7, 0.6}},
		[][]int64{{1, 1, 2, 3}},
		groupByCol)
	defer df0.Release()

	df1 := buildTestDFWithStringPKAndGroupBy(pool,
		[][]string{{"e", "f", "b"}},
		[][]float32{{0.85, 0.75, 0.65}},
		[][]int64{{2, 2, 1}},
		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)
	require.Equal(t, 4, ids.Len())
	assert.Equal(t, []string{"a", "e", "b", "f"}, []string{
		ids.Value(0), ids.Value(1), ids.Value(2), ids.Value(3),
	})

	gb := result.DF.Column(groupByCol).Chunk(0).(*array.Int64)
	assert.Equal(t, []int64{1, 2, 1, 2}, []int64{
		gb.Value(0), gb.Value(1), gb.Value(2), gb.Value(3),
	})
}

func TestHeapMergeReduce_GroupByStringPKSkipsDuplicateAndFullGroup(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithStringPKAndGroupBy(pool,
		[][]string{{"a", "b", "c"}},
		[][]float32{{0.99, 0.97, 0.95}},
		[][]int64{{1, 1, 2}},
		groupByCol)
	defer df0.Release()

	df1 := buildTestDFWithStringPKAndGroupBy(pool,
		[][]string{{"a", "d"}},
		[][]float32{{0.98, 0.96}},
		[][]int64{{1, 2}},
		groupByCol)
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 2, &groupByOptions{
		GroupSize: 1,
	})
	require.NoError(t, err)
	defer result.DF.Release()

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.String)
	require.Equal(t, 2, ids.Len())
	assert.Equal(t, []string{"a", "d"}, []string{ids.Value(0), ids.Value(1)})

	gb := result.DF.Column(groupByCol).Chunk(0).(*array.Int64)
	assert.Equal(t, []int64{1, 2}, []int64{gb.Value(0), gb.Value(1)})
}

func TestMergeGroupByNoGroupFieldsFallsBackToStandard(t *testing.T) {
	pool := memory.NewGoAllocator()

	intIDBuilder := array.NewInt64Builder(pool)
	intIDBuilder.AppendValues([]int64{1, 2}, nil)
	intIDs := intIDBuilder.NewArray().(*array.Int64)
	intIDBuilder.Release()
	defer intIDs.Release()

	intScoreBuilder := array.NewFloat32Builder(pool)
	intScoreBuilder.AppendValues([]float32{0.9, 0.8}, nil)
	intScores := intScoreBuilder.NewArray().(*array.Float32)
	intScoreBuilder.Release()
	defer intScores.Release()

	intOffsetBuilder := array.NewInt64Builder(pool)
	intOffsetBuilder.AppendValues([]int64{10, 11}, nil)
	intOffsets := intOffsetBuilder.NewArray().(*array.Int64)
	intOffsetBuilder.Release()
	defer intOffsets.Release()

	intHeap := &mergeHeapInt64Pk{&mergeEntry{idInt64: intIDs, scoreArr: intScores, segOffsetArr: intOffsets}}
	heap.Init(intHeap)
	intIDsOut, _, intSources, err := mergeGroupByInt64Pk(intHeap, 2, 2, 0)
	require.NoError(t, err)
	assert.Equal(t, []int64{1, 2}, intIDsOut)
	assert.Equal(t, []segmentSource{{SegOffset: 10, OriginalIdx: 0}, {SegOffset: 11, OriginalIdx: 1}}, intSources)

	stringIDBuilder := array.NewStringBuilder(pool)
	stringIDBuilder.AppendValues([]string{"a", "b"}, nil)
	stringIDs := stringIDBuilder.NewArray().(*array.String)
	stringIDBuilder.Release()
	defer stringIDs.Release()

	stringScoreBuilder := array.NewFloat32Builder(pool)
	stringScoreBuilder.AppendValues([]float32{0.7, 0.6}, nil)
	stringScores := stringScoreBuilder.NewArray().(*array.Float32)
	stringScoreBuilder.Release()
	defer stringScores.Release()

	stringOffsetBuilder := array.NewInt64Builder(pool)
	stringOffsetBuilder.AppendValues([]int64{20, 21}, nil)
	stringOffsets := stringOffsetBuilder.NewArray().(*array.Int64)
	stringOffsetBuilder.Release()
	defer stringOffsets.Release()

	stringHeap := &mergeHeapStringPk{&mergeEntry{idString: stringIDs, scoreArr: stringScores, segOffsetArr: stringOffsets}}
	heap.Init(stringHeap)
	stringIDsOut, _, stringSources, err := mergeGroupByStringPk(stringHeap, 2, 2, 0)
	require.NoError(t, err)
	assert.Equal(t, []string{"a", "b"}, stringIDsOut)
	assert.Equal(t, []segmentSource{{SegOffset: 20, OriginalIdx: 0}, {SegOffset: 21, OriginalIdx: 1}}, stringSources)
}

func TestExtractCompositeGroupValuesMissingArray(t *testing.T) {
	pool := memory.NewGoAllocator()

	floatBuilder := array.NewFloat32Builder(pool)
	floatBuilder.AppendValues([]float32{1.25}, nil)
	floatArr := floatBuilder.NewArray()
	floatBuilder.Release()
	defer floatArr.Release()

	values, err := extractCompositeGroupValues(&mergeEntry{
		cursor:      0,
		groupByArrs: []arrow.Array{floatArr},
	}, 2)

	require.Error(t, err)
	assert.Nil(t, values)
	assert.Contains(t, err.Error(), "unsupported group-by arrow type float32")
}

func TestExtractCompositeGroupValuesMissingArrayErrors(t *testing.T) {
	pool := memory.NewGoAllocator()

	intBuilder := array.NewInt64Builder(pool)
	intBuilder.AppendValues([]int64{1}, nil)
	intArr := intBuilder.NewArray()
	intBuilder.Release()
	defer intArr.Release()

	values, err := extractCompositeGroupValues(&mergeEntry{
		cursor:      0,
		groupByArrs: []arrow.Array{intArr},
	}, 2)

	require.Error(t, err)
	assert.Nil(t, values)
	assert.Contains(t, err.Error(), "missing group-by column")
}

func buildTestDFWithTwoGroupByColumns(
	pool memory.Allocator,
	idsPerChunk [][]int64,
	scoresPerChunk [][]float32,
	groupBy0PerChunk [][]int64,
	groupBy1PerChunk [][]string,
	groupByCol0 string,
	groupByCol1 string,
) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, groupByCol0, groupByCol1}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()

		gb0Builder := array.NewInt64Builder(pool)
		gb0Builder.AppendValues(groupBy0PerChunk[i], nil)
		collector.Set(groupByCol0, i, gb0Builder.NewArray())
		gb0Builder.Release()

		gb1Builder := array.NewStringBuilder(pool)
		gb1Builder.AppendValues(groupBy1PerChunk[i], nil)
		collector.Set(groupByCol1, i, gb1Builder.NewArray())
		gb1Builder.Release()

		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(groupByCol0, collector.Consume(groupByCol0))
	_ = builder.AddColumnFromChunks(groupByCol1, collector.Consume(groupByCol1))
	collector.Release()

	return builder.Build()
}

// buildTestDFWithGroupByArr accepts a pre-built arrow.Array per chunk so tests
// can exercise group-by value types beyond the int64 covered by buildTestDFWithGroupBy.
func buildTestDFWithGroupByArr(pool memory.Allocator, idsPerChunk [][]int64, scoresPerChunk [][]float32, groupByArrPerChunk []arrow.Array) *chain.DataFrame {
	numChunks := len(idsPerChunk)
	chunkSizes := make([]int64, numChunks)

	colNames := []string{idFieldName, scoreFieldName, segOffsetCol, groupByCol}
	collector := chain.NewChunkCollector(colNames, numChunks)

	for i := 0; i < numChunks; i++ {
		chunkSizes[i] = int64(len(idsPerChunk[i]))

		idBuilder := array.NewInt64Builder(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()

		collector.Set(groupByCol, i, groupByArrPerChunk[i])
		appendDefaultSegOffsetChunk(pool, collector, i, len(idsPerChunk[i]))
	}

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes(chunkSizes)
	_ = builder.AddColumnFromChunks(idFieldName, collector.Consume(idFieldName))
	_ = builder.AddColumnFromChunks(scoreFieldName, collector.Consume(scoreFieldName))
	_ = builder.AddColumnFromChunks(segOffsetCol, collector.Consume(segOffsetCol))
	_ = builder.AddColumnFromChunks(groupByCol, collector.Consume(groupByCol))
	collector.Release()

	return builder.Build()
}

// TestHeapMergeReduce_GroupByValueTypes covers the dispatcher cases (Int8/16/32,
// String) not exercised by the int64/bool tests above.
func TestHeapMergeReduce_GroupByValueTypes(t *testing.T) {
	idsSeg0 := [][]int64{{1, 2, 3, 4}}
	scoresSeg0 := [][]float32{{0.9, 0.8, 0.7, 0.6}}
	idsSeg1 := [][]int64{{5, 6}}
	scoresSeg1 := [][]float32{{0.85, 0.75}}

	expectedIDs := []int64{1, 5, 2, 6}

	tests := []struct {
		name      string
		buildSeg0 func(pool memory.Allocator) arrow.Array
		buildSeg1 func(pool memory.Allocator) arrow.Array
		assertCol func(t *testing.T, arr arrow.Array)
	}{
		{
			name: "Int8",
			buildSeg0: func(pool memory.Allocator) arrow.Array {
				b := array.NewInt8Builder(pool)
				defer b.Release()
				b.AppendValues([]int8{10, 10, 10, 20}, nil)
				return b.NewArray()
			},
			buildSeg1: func(pool memory.Allocator) arrow.Array {
				b := array.NewInt8Builder(pool)
				defer b.Release()
				b.AppendValues([]int8{20, 20}, nil)
				return b.NewArray()
			},
			assertCol: func(t *testing.T, arr arrow.Array) {
				a, ok := arr.(*array.Int8)
				require.True(t, ok, "group-by column should be *array.Int8")
				// Group values follow same order as expectedIDs: A, B, A, B.
				assert.Equal(t, int8(10), a.Value(0))
				assert.Equal(t, int8(20), a.Value(1))
				assert.Equal(t, int8(10), a.Value(2))
				assert.Equal(t, int8(20), a.Value(3))
			},
		},
		{
			name: "Int16",
			buildSeg0: func(pool memory.Allocator) arrow.Array {
				b := array.NewInt16Builder(pool)
				defer b.Release()
				b.AppendValues([]int16{100, 100, 100, 200}, nil)
				return b.NewArray()
			},
			buildSeg1: func(pool memory.Allocator) arrow.Array {
				b := array.NewInt16Builder(pool)
				defer b.Release()
				b.AppendValues([]int16{200, 200}, nil)
				return b.NewArray()
			},
			assertCol: func(t *testing.T, arr arrow.Array) {
				a, ok := arr.(*array.Int16)
				require.True(t, ok, "group-by column should be *array.Int16")
				assert.Equal(t, int16(100), a.Value(0))
				assert.Equal(t, int16(200), a.Value(1))
				assert.Equal(t, int16(100), a.Value(2))
				assert.Equal(t, int16(200), a.Value(3))
			},
		},
		{
			name: "Int32",
			buildSeg0: func(pool memory.Allocator) arrow.Array {
				b := array.NewInt32Builder(pool)
				defer b.Release()
				b.AppendValues([]int32{1000, 1000, 1000, 2000}, nil)
				return b.NewArray()
			},
			buildSeg1: func(pool memory.Allocator) arrow.Array {
				b := array.NewInt32Builder(pool)
				defer b.Release()
				b.AppendValues([]int32{2000, 2000}, nil)
				return b.NewArray()
			},
			assertCol: func(t *testing.T, arr arrow.Array) {
				a, ok := arr.(*array.Int32)
				require.True(t, ok, "group-by column should be *array.Int32")
				assert.Equal(t, int32(1000), a.Value(0))
				assert.Equal(t, int32(2000), a.Value(1))
				assert.Equal(t, int32(1000), a.Value(2))
				assert.Equal(t, int32(2000), a.Value(3))
			},
		},
		{
			name: "String",
			buildSeg0: func(pool memory.Allocator) arrow.Array {
				b := array.NewStringBuilder(pool)
				defer b.Release()
				b.AppendValues([]string{"groupA", "groupA", "groupA", "groupB"}, nil)
				return b.NewArray()
			},
			buildSeg1: func(pool memory.Allocator) arrow.Array {
				b := array.NewStringBuilder(pool)
				defer b.Release()
				b.AppendValues([]string{"groupB", "groupB"}, nil)
				return b.NewArray()
			},
			assertCol: func(t *testing.T, arr arrow.Array) {
				a, ok := arr.(*array.String)
				require.True(t, ok, "group-by column should be *array.String")
				assert.Equal(t, "groupA", a.Value(0))
				assert.Equal(t, "groupB", a.Value(1))
				assert.Equal(t, "groupA", a.Value(2))
				assert.Equal(t, "groupB", a.Value(3))
			},
		},
		{
			name: "Bool",
			buildSeg0: func(pool memory.Allocator) arrow.Array {
				b := array.NewBooleanBuilder(pool)
				defer b.Release()
				b.AppendValues([]bool{true, true, true, false}, nil)
				return b.NewArray()
			},
			buildSeg1: func(pool memory.Allocator) arrow.Array {
				b := array.NewBooleanBuilder(pool)
				defer b.Release()
				b.AppendValues([]bool{false, false}, nil)
				return b.NewArray()
			},
			assertCol: func(t *testing.T, arr arrow.Array) {
				a, ok := arr.(*array.Boolean)
				require.True(t, ok, "group-by column should be *array.Boolean")
				assert.Equal(t, true, a.Value(0))
				assert.Equal(t, false, a.Value(1))
				assert.Equal(t, true, a.Value(2))
				assert.Equal(t, false, a.Value(3))
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			pool := memory.NewGoAllocator()

			df0 := buildTestDFWithGroupByArr(pool,
				idsSeg0, scoresSeg0,
				[]arrow.Array{tc.buildSeg0(pool)},
			)
			defer df0.Release()

			df1 := buildTestDFWithGroupByArr(pool,
				idsSeg1, scoresSeg1,
				[]arrow.Array{tc.buildSeg1(pool)},
			)
			defer df1.Release()

			result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 2, &groupByOptions{
				GroupSize: 2,
			})
			require.NoError(t, err)
			defer result.DF.Release()

			require.Equal(t, int64(4), result.DF.NumRows(),
				"topK=2 * groupSize=2 = 4 rows expected")

			ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
			for i, want := range expectedIDs {
				assert.Equal(t, want, ids.Value(i), "id mismatch at row %d", i)
			}

			require.True(t, result.DF.HasColumn(groupByCol),
				"group-by column should be propagated to output")
			tc.assertCol(t, result.DF.Column(groupByCol).Chunk(0))
		})
	}
}

func TestHeapMergeReduce_SourceTracking(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDF(pool, [][]int64{{1, 2}}, [][]float32{{0.9, 0.7}})
	defer df0.Release()

	df1 := buildTestDF(pool, [][]int64{{3}}, [][]float32{{0.8}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 3, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Verify source tracking
	require.Len(t, result.Sources, 1) // 1 chunk
	sources := result.Sources[0]
	require.Len(t, sources, 3)

	// id=1(0.9) from input 0
	assert.Equal(t, 0, sources[0].InputIdx)
	// id=3(0.8) from input 1
	assert.Equal(t, 1, sources[1].InputIdx)
	// id=2(0.7) from input 0
	assert.Equal(t, 0, sources[2].InputIdx)
}

func TestHeapMergeReduce_LargeScale(t *testing.T) {
	pool := memory.NewGoAllocator()

	// 4 segments, each with 100 results
	numSegments := 4
	numResults := 100
	dfs := make([]*chain.DataFrame, numSegments)
	for seg := 0; seg < numSegments; seg++ {
		ids := make([]int64, numResults)
		scores := make([]float32, numResults)
		for i := 0; i < numResults; i++ {
			ids[i] = int64(seg*1000 + i)
			scores[i] = float32(numResults-i) / float32(numResults) // descending scores
		}
		dfs[seg] = buildTestDF(pool, [][]int64{ids}, [][]float32{scores})
		defer dfs[seg].Release()
	}

	result, err := heapMergeReduce(pool, dfs, 50, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	assert.Equal(t, int64(50), result.DF.NumRows())

	// Verify scores are in descending order
	scoreCol := result.DF.Column(scoreFieldName).Chunk(0).(*array.Float32)
	for i := 1; i < 50; i++ {
		assert.GreaterOrEqual(t, scoreCol.Value(i-1), scoreCol.Value(i),
			"scores should be descending at index %d", i)
	}
}

func TestHeapMergeReduce_NoInputs(t *testing.T) {
	pool := memory.NewGoAllocator()
	_, err := heapMergeReduce(pool, nil, 10, nil)
	assert.Error(t, err)
}

func TestHeapMergeReduce_MissingSegOffsetErrors(t *testing.T) {
	pool := memory.NewGoAllocator()

	df := buildTestDFWithoutSegOffset(pool, [][]int64{{1}}, [][]float32{{0.9}})
	defer df.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df}, 1, nil)
	require.Error(t, err)
	assert.Nil(t, result)
	assert.Contains(t, err.Error(), "missing required $seg_offset column")
}

func TestHeapMergeReduce_InvalidGroupSizeErrors(t *testing.T) {
	pool := memory.NewGoAllocator()

	df := buildTestDFWithGroupBy(pool,
		[][]int64{{1}},
		[][]float32{{0.9}},
		[][]int64{{10}},
		groupByCol)
	defer df.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df}, 1, &groupByOptions{
		GroupSize: 0,
		Columns:   []string{groupByCol},
	})
	require.Error(t, err)
	assert.Nil(t, result)
	assert.Contains(t, err.Error(), "group size must be positive")
}

func TestHeapMergeReduce_MismatchedChunks(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDF(pool, [][]int64{{1}, {2}}, [][]float32{{0.9}, {0.8}})
	defer df0.Release()

	df1 := buildTestDF(pool, [][]int64{{3}}, [][]float32{{0.7}})
	defer df1.Release()

	_, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 10, nil)
	assert.Error(t, err)
}

func TestHeapMergeReduce_MissingIDColumn(t *testing.T) {
	pool := memory.NewGoAllocator()

	scoreBuilder := array.NewFloat32Builder(pool)
	scoreBuilder.AppendValues([]float32{0.9, 0.8}, nil)
	scoreArr := scoreBuilder.NewArray()
	scoreBuilder.Release()

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes([]int64{2})
	require.NoError(t, builder.AddColumnFromChunks(scoreFieldName, []arrow.Array{scoreArr}))
	df := builder.Build()
	defer df.Release()

	_, err := heapMergeReduce(pool, []*chain.DataFrame{df}, 2, nil)
	require.Error(t, err)
	assert.Contains(t, err.Error(), "$id column not found")
}

func TestHeapMergeReduce_MissingScoreColumnProducesEmptyResult(t *testing.T) {
	pool := memory.NewGoAllocator()

	idBuilder := array.NewInt64Builder(pool)
	idBuilder.AppendValues([]int64{1, 2}, nil)
	idArr := idBuilder.NewArray()
	idBuilder.Release()

	builder := chain.NewDataFrameBuilder()
	builder.SetChunkSizes([]int64{2})
	require.NoError(t, builder.AddColumnFromChunks(idFieldName, []arrow.Array{idArr}))
	df := builder.Build()
	defer df.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df}, 2, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	assert.Equal(t, []int64{0}, result.DF.ChunkSizes())
	assert.Zero(t, result.DF.NumRows())
	require.Len(t, result.Sources, 1)
	assert.Empty(t, result.Sources[0])
}

func TestHeapMergeReduceImplDuplicateOutputColumnError(t *testing.T) {
	pool := memory.NewGoAllocator()

	df := buildTestDF(pool, [][]int64{{1}}, [][]float32{{0.9}})
	defer df.Release()

	result, err := heapMergeReduceImpl(pool, []*chain.DataFrame{df}, 1, &groupByOptions{
		GroupSize: 1,
		Columns:   []string{idFieldName},
	}, 0, 1, false)
	require.Error(t, err)
	assert.Nil(t, result)
	assert.Contains(t, err.Error(), "already exists")
}

func TestPickGroupByValuesEmptyAndUnsupportedArrays(t *testing.T) {
	pool := memory.NewGoAllocator()

	intBuilder := array.NewInt32Builder(pool)
	intBuilder.AppendValues([]int32{10, 20}, nil)
	intArr := intBuilder.NewArray()
	intBuilder.Release()
	defer intArr.Release()

	intDF := buildTestDFWithGroupByArr(pool,
		[][]int64{{1, 2}},
		[][]float32{{0.9, 0.8}},
		[]arrow.Array{intArr})
	defer intDF.Release()
	intCols := resolveInputCols([]*chain.DataFrame{intDF}, &groupByOptions{Columns: []string{groupByCol}}, false)

	empty, err := pickGroupByValues(pool, intCols, 0, nil, 0)
	require.NoError(t, err)
	defer empty.Release()
	assert.Equal(t, arrow.INT32, empty.DataType().ID())
	assert.Zero(t, empty.Len())

	floatBuilder := array.NewFloat32Builder(pool)
	floatBuilder.AppendValues([]float32{1.5}, nil)
	floatArr := floatBuilder.NewArray()
	floatBuilder.Release()
	defer floatArr.Release()

	floatDF := buildTestDFWithGroupByArr(pool,
		[][]int64{{1}},
		[][]float32{{0.9}},
		[]arrow.Array{floatArr})
	defer floatDF.Release()
	floatCols := resolveInputCols([]*chain.DataFrame{floatDF}, &groupByOptions{Columns: []string{groupByCol}}, false)

	unsupported, err := pickGroupByValues(pool, floatCols, 0, []segmentSource{{InputIdx: 0, OriginalIdx: 0}}, 0)
	require.Error(t, err)
	assert.Nil(t, unsupported)
	assert.Contains(t, err.Error(), "unsupported group-by arrow type float32")

	fallback := buildEmptyGroupByArray(pool, []inputCols{{}}, 0)
	defer fallback.Release()
	assert.Equal(t, arrow.INT64, fallback.DataType().ID())
	assert.Zero(t, fallback.Len())

	noArrays, err := pickGroupByValues(pool, []inputCols{{}}, 0, []segmentSource{{InputIdx: 0}}, 0)
	require.NoError(t, err)
	defer noArrays.Release()
	assert.Equal(t, arrow.INT64, noArrays.DataType().ID())
	assert.Zero(t, noArrays.Len())
}

func TestHeapMergeReduceUnsupportedGroupByTypeReleasesMemory(t *testing.T) {
	pool := memory.NewCheckedAllocator(memory.NewGoAllocator())
	defer pool.AssertSize(t, 0)

	floatBuilder := array.NewFloat32Builder(pool)
	floatBuilder.AppendValues([]float32{1.5}, nil)
	floatArr := floatBuilder.NewArray()
	floatBuilder.Release()

	df := buildTestDFWithGroupByArr(pool,
		[][]int64{{1}},
		[][]float32{{0.9}},
		[]arrow.Array{floatArr})
	defer df.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df}, 1, &groupByOptions{
		GroupSize: 1,
		Columns:   []string{groupByCol},
	})
	require.Error(t, err)
	assert.Nil(t, result)
	assert.Contains(t, err.Error(), "unsupported group-by arrow type float32")
}

func TestHeapMergeReduceMissingGroupByColumnErrors(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithGroupBy(pool,
		[][]int64{{1}},
		[][]float32{{0.9}},
		[][]int64{{10}},
		groupByCol)
	defer df0.Release()

	df1 := buildTestDF(pool,
		[][]int64{{2}},
		[][]float32{{0.8}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 2, &groupByOptions{
		GroupSize: 1,
		Columns:   []string{groupByCol},
	})
	require.Error(t, err)
	assert.Nil(t, result)
	assert.Contains(t, err.Error(), "missing group-by column")
}

func TestHeapMergeReduceElementIndicesMustBePresentOnAllInputs(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithElementIndices(pool,
		[][]int64{{1}},
		[][]float32{{0.9}},
		[][]int32{{3}})
	defer df0.Release()

	df1 := buildTestDF(pool,
		[][]int64{{2}},
		[][]float32{{0.8}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 2, nil)
	require.Error(t, err)
	assert.Nil(t, result)
	assert.Contains(t, err.Error(), "missing $element_indices column")
}

func TestCompositeGroupCounterHashBucketCollision(t *testing.T) {
	values := []any{int64(2)}
	hash := reduce.HashGroupValues(values)
	counter := &compositeGroupCounter{
		groups: map[uint64][]*compositeGroup{
			hash: {
				{values: []any{int64(1)}, count: 1},
			},
		},
		distinct:       1,
		saturatedCount: 1,
		topK:           2,
		groupSize:      1,
	}

	require.True(t, counter.shouldAccept(values))
	require.Len(t, counter.groups[hash], 2)
	assert.Equal(t, values, counter.groups[hash][1].values)
	assert.Equal(t, int64(2), counter.distinct)
}

func TestCompositeGroupCounterAllSaturated(t *testing.T) {
	counter := newCompositeGroupCounter(2, 2)
	assert.False(t, counter.allSaturated())
	assert.Equal(t, int64(0), counter.saturatedCount)

	require.True(t, counter.shouldAccept([]any{int64(10)}))
	require.True(t, counter.shouldAccept([]any{int64(10)}))
	assert.False(t, counter.allSaturated())
	assert.Equal(t, int64(1), counter.saturatedCount)

	require.True(t, counter.shouldAccept([]any{int64(20)}))
	assert.False(t, counter.allSaturated())
	assert.Equal(t, int64(1), counter.saturatedCount)

	require.True(t, counter.shouldAccept([]any{int64(20)}))
	assert.True(t, counter.allSaturated())
	assert.Equal(t, int64(2), counter.saturatedCount)
}

// TestHeapMergeReduce_ElementIndicesBasic verifies that the $element_indices
// column is propagated through the heap merge. Each result row should carry
// the element index from its original segment.
func TestHeapMergeReduce_ElementIndicesBasic(t *testing.T) {
	pool := memory.NewGoAllocator()

	// Segment 0: doc 1 matched element 2, doc 2 matched element 0
	df0 := buildTestDFWithElementIndices(pool,
		[][]int64{{1, 2}},
		[][]float32{{0.9, 0.7}},
		[][]int32{{2, 0}})
	defer df0.Release()

	// Segment 1: doc 3 matched element 5, doc 4 matched element 1
	df1 := buildTestDFWithElementIndices(pool,
		[][]int64{{3, 4}},
		[][]float32{{0.8, 0.6}},
		[][]int32{{5, 1}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 4, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	// Merged order by score: doc1(0.9,ei=2), doc3(0.8,ei=5), doc2(0.7,ei=0), doc4(0.6,ei=1)
	require.True(t, result.DF.HasColumn(elementIndicesCol),
		"merged DataFrame should have $element_indices column")

	eiCol := result.DF.Column(elementIndicesCol).Chunk(0).(*array.Int32)
	require.Equal(t, 4, eiCol.Len())

	assert.Equal(t, int32(2), eiCol.Value(0)) // doc1 from seg0
	assert.Equal(t, int32(5), eiCol.Value(1)) // doc3 from seg1
	assert.Equal(t, int32(0), eiCol.Value(2)) // doc2 from seg0
	assert.Equal(t, int32(1), eiCol.Value(3)) // doc4 from seg1
}

// TestHeapMergeReduce_ElementIndicesDedup verifies that element-level search
// deduplicates by (PK, element_index): same PK with different elements is kept,
// while the same element hit is dropped.
func TestHeapMergeReduce_ElementIndicesDedup(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithElementIndices(pool,
		[][]int64{{5, 5, 6}},
		[][]float32{{0.99, 0.97, 0.96}},
		[][]int32{{0, 1, 0}})
	defer df0.Release()

	df1 := buildTestDFWithElementIndices(pool,
		[][]int64{{5, 7}},
		[][]float32{{0.98, 0.95}},
		[][]int32{{1, 0}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 4, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	ids := result.DF.Column(idFieldName).Chunk(0).(*array.Int64)
	eiCol := result.DF.Column(elementIndicesCol).Chunk(0).(*array.Int32)
	require.Equal(t, 4, ids.Len())

	assert.Equal(t, int64(5), ids.Value(0))
	assert.Equal(t, int32(0), eiCol.Value(0))

	assert.Equal(t, int64(5), ids.Value(1))
	assert.Equal(t, int32(1), eiCol.Value(1))

	assert.Equal(t, int64(6), ids.Value(2))
	assert.Equal(t, int32(0), eiCol.Value(2))

	assert.Equal(t, int64(7), ids.Value(3))
	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()

	df0 := buildTestDFWithElementIndicesAndGroupBy(pool,
		[][]int64{{5, 5, 6}},
		[][]float32{{0.99, 0.97, 0.96}},
		[][]int32{{0, 1, 0}},
		[][]int64{{5, 5, 6}},
		groupByCol)
	defer df0.Release()

	df1 := buildTestDFWithElementIndicesAndGroupBy(pool,
		[][]int64{{5, 7}},
		[][]float32{{0.98, 0.95}},
		[][]int32{{1, 0}},
		[][]int64{{5, 7}},
		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.Int64)
	eiCol := result.DF.Column(elementIndicesCol).Chunk(0).(*array.Int32)
	require.Equal(t, 3, ids.Len())

	assert.Equal(t, []int64{5, 5, 6}, []int64{ids.Value(0), ids.Value(1), ids.Value(2)})
	assert.Equal(t, []int32{0, 1, 0}, []int32{eiCol.Value(0), eiCol.Value(1), eiCol.Value(2)})
}

// TestHeapMergeReduce_ElementIndicesMultiNQ verifies element indices propagation
// across multiple NQ chunks.
func TestHeapMergeReduce_ElementIndicesMultiNQ(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDFWithElementIndices(pool,
		[][]int64{{1, 2}, {10, 11}},
		[][]float32{{0.9, 0.7}, {0.8, 0.6}},
		[][]int32{{0, 1}, {2, 3}})
	defer df0.Release()

	df1 := buildTestDFWithElementIndices(pool,
		[][]int64{{3}, {12}},
		[][]float32{{0.85}, {0.75}},
		[][]int32{{4}, {5}})
	defer df1.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0, df1}, 3, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	assert.Equal(t, 2, result.DF.NumChunks())

	// Chunk 0: doc1(0.9,ei=0), doc3(0.85,ei=4), doc2(0.7,ei=1)
	ei0 := result.DF.Column(elementIndicesCol).Chunk(0).(*array.Int32)
	assert.Equal(t, int32(0), ei0.Value(0))
	assert.Equal(t, int32(4), ei0.Value(1))
	assert.Equal(t, int32(1), ei0.Value(2))

	// Chunk 1: doc10(0.8,ei=2), doc12(0.75,ei=5), doc11(0.6,ei=3)
	ei1 := result.DF.Column(elementIndicesCol).Chunk(1).(*array.Int32)
	assert.Equal(t, int32(2), ei1.Value(0))
	assert.Equal(t, int32(5), ei1.Value(1))
	assert.Equal(t, int32(3), ei1.Value(2))
}

// TestHeapMergeReduce_NoElementIndices verifies that when inputs don't have
// $element_indices, the output also doesn't — no spurious column is created.
func TestHeapMergeReduce_NoElementIndices(t *testing.T) {
	pool := memory.NewGoAllocator()

	df0 := buildTestDF(pool, [][]int64{{1, 2}}, [][]float32{{0.9, 0.7}})
	defer df0.Release()

	result, err := heapMergeReduce(pool, []*chain.DataFrame{df0}, 2, nil)
	require.NoError(t, err)
	defer result.DF.Release()

	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)
}
