snapshot of https://github.com/blevesearch/bleve.git (master, 2026-09-15; blobs > 131072B omitted)

search/collector/bench_test.go

@@ -0,0 +1,50 @@
+//  Copyright (c) 2016 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"context"
+	"math/rand"
+	"strconv"
+	"testing"
+
+	"github.com/blevesearch/bleve/v2/search"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+type createCollector func() search.Collector
+
+func benchHelper(numOfMatches int, cc createCollector, b *testing.B) {
+	matches := make([]*search.DocumentMatch, 0, numOfMatches)
+	for i := 0; i < numOfMatches; i++ {
+		matches = append(matches, &search.DocumentMatch{
+			IndexInternalID: index.IndexInternalID(strconv.Itoa(i)),
+			Score:           rand.Float64(),
+		})
+	}
+
+	b.ResetTimer()
+
+	for run := 0; run < b.N; run++ {
+		searcher := &stubSearcher{
+			matches: matches,
+		}
+		collector := cc()
+		err := collector.Collect(context.Background(), searcher, &stubReader{})
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}

search/collector/eligible.go

@@ -0,0 +1,172 @@
+//  Copyright (c) 2024 Couchbase, Inc.
+//
+// Licensed 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.
+
+//go:build vectors
+// +build vectors
+
+package collector
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/blevesearch/bleve/v2/search"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+type EligibleCollector struct {
+	size             int
+	total            uint64
+	took             time.Duration
+	eligibleSelector index.EligibleDocumentSelector
+}
+
+func NewEligibleCollector(size int) *EligibleCollector {
+	return newEligibleCollector(size)
+}
+
+func newEligibleCollector(size int) *EligibleCollector {
+	// No sort order & skip always 0 since this is only to filter eligible docs.
+	ec := &EligibleCollector{
+		size: size,
+	}
+	return ec
+}
+
+func makeEligibleDocumentMatchHandler(ctx *search.SearchContext, reader index.IndexReader) (search.DocumentMatchHandler, error) {
+	if ec, ok := ctx.Collector.(*EligibleCollector); ok {
+		if vr, ok := reader.(index.VectorIndexReader); ok {
+			// create a new eligible document selector to add eligible document matches
+			ec.eligibleSelector = vr.NewEligibleDocumentSelector()
+			// return a document match handler that adds eligible document matches
+			// to the eligible document selector
+			return func(d *search.DocumentMatch) error {
+				if d == nil {
+					return nil
+				}
+				err := ec.eligibleSelector.AddEligibleDocumentMatch(d.IndexInternalID)
+				if err != nil {
+					return err
+				}
+				// recycle the DocumentMatch
+				ctx.DocumentMatchPool.Put(d)
+				return nil
+			}, nil
+		}
+		return nil, fmt.Errorf("reader is not a VectorIndexReader")
+	}
+
+	return nil, fmt.Errorf("eligiblity collector not available")
+}
+
+func (ec *EligibleCollector) Collect(ctx context.Context, searcher search.Searcher, reader index.IndexReader) error {
+	startTime := time.Now()
+	var err error
+	var next *search.DocumentMatch
+
+	backingSize := ec.size
+	if backingSize > PreAllocSizeSkipCap {
+		backingSize = PreAllocSizeSkipCap + 1
+	}
+	searchContext := &search.SearchContext{
+		DocumentMatchPool: search.NewDocumentMatchPool(backingSize+searcher.DocumentMatchPoolSize(), 0),
+		Collector:         ec,
+		IndexReader:       reader,
+	}
+
+	dmHandler, err := makeEligibleDocumentMatchHandler(searchContext, reader)
+	if err != nil {
+		return err
+	}
+	select {
+	case <-ctx.Done():
+		search.RecordSearchCost(ctx, search.AbortM, 0)
+		return ctx.Err()
+	default:
+		next, err = searcher.Next(searchContext)
+	}
+	for err == nil && next != nil {
+		if ec.total%CheckDoneEvery == 0 {
+			select {
+			case <-ctx.Done():
+				search.RecordSearchCost(ctx, search.AbortM, 0)
+				return ctx.Err()
+			default:
+			}
+		}
+		ec.total++
+
+		err = dmHandler(next)
+		if err != nil {
+			break
+		}
+
+		next, err = searcher.Next(searchContext)
+	}
+	if err != nil {
+		return err
+	}
+
+	// help finalize/flush the results in case
+	// of custom document match handlers.
+	err = dmHandler(nil)
+	if err != nil {
+		return err
+	}
+
+	// compute search duration
+	ec.took = time.Since(startTime)
+
+	return nil
+}
+
+// The eligible collector does not return any document matches and hence
+// this method is a dummy method returning nil, to conform to the
+// search.Collector interface.
+func (ec *EligibleCollector) Results() search.DocumentMatchCollection {
+	return nil
+}
+
+// EligibleSelector returns the eligible document selector, which can be used
+// to retrieve the list of eligible documents from this collector.
+// If the collector has no results, it returns nil.
+func (ec *EligibleCollector) EligibleSelector() index.EligibleDocumentSelector {
+	if ec.total == 0 {
+		return nil
+	}
+	return ec.eligibleSelector
+}
+
+func (ec *EligibleCollector) Total() uint64 {
+	return ec.total
+}
+
+// No concept of scoring in the eligible collector.
+func (ec *EligibleCollector) MaxScore() float64 {
+	return 0
+}
+
+func (ec *EligibleCollector) Took() time.Duration {
+	return ec.took
+}
+
+func (ec *EligibleCollector) SetFacetsBuilder(facetsBuilder *search.FacetsBuilder) {
+	// facet unsupported for pre-filtering in KNN search
+}
+
+func (ec *EligibleCollector) FacetResults() search.FacetResults {
+	// facet unsupported for pre-filtering in KNN search
+	return nil
+}

search/collector/heap.go

@@ -0,0 +1,140 @@
+//  Copyright (c) 2014 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"sort"
+
+	"github.com/blevesearch/bleve/v2/search"
+)
+
+// collectStoreHeap is a min-heap of DocumentMatches where the root (heap[0])
+// is always the *worst* result (lowest score / highest sort key).  Popping the
+// root discards the worst element, so AddNotExceedingSize keeps the best k
+// documents at all times.
+//
+// Implemented as a ternary heap (3 children per node) instead of the standard
+// library's binary heap.  Height is log₃(n) vs log₂(n), so siftDown is
+// shallower and fetches more children per cache line.
+type collectStoreHeap struct {
+	heap    search.DocumentMatchCollection
+	compare collectorCompare
+}
+
+func newStoreHeap(capacity int, compare collectorCompare) *collectStoreHeap {
+	return &collectStoreHeap{
+		heap:    make(search.DocumentMatchCollection, 0, capacity),
+		compare: compare,
+	}
+}
+
+func (c *collectStoreHeap) AddNotExceedingSize(doc *search.DocumentMatch,
+	size int) *search.DocumentMatch {
+	c.add(doc)
+	if len(c.heap) > size {
+		return c.removeLast()
+	}
+	return nil
+}
+
+func (c *collectStoreHeap) add(doc *search.DocumentMatch) {
+	c.heap = append(c.heap, doc)
+	c.siftUp(len(c.heap) - 1)
+}
+
+func (c *collectStoreHeap) removeLast() *search.DocumentMatch {
+	n := len(c.heap)
+	c.heap[0], c.heap[n-1] = c.heap[n-1], c.heap[0]
+	result := c.heap[n-1]
+	c.heap = c.heap[:n-1]
+	if len(c.heap) > 0 {
+		c.siftDown(0)
+	}
+	return result
+}
+
+// siftUp restores the heap invariant after appending at index i.
+// Moves element at i toward the root while it is less than (worse than) its parent.
+func (c *collectStoreHeap) siftUp(i int) {
+	h := c.heap
+	for i > 0 {
+		parent := (i - 1) / 3
+		if c.compare(h[i], h[parent]) > 0 { // h[i] is worse → should be closer to root
+			h[i], h[parent] = h[parent], h[i]
+			i = parent
+		} else {
+			break
+		}
+	}
+}
+
+// siftDown restores the heap invariant after replacing the root.
+// Moves element at i toward the leaves while any child is worse (less) than it.
+func (c *collectStoreHeap) siftDown(i int) {
+	h := c.heap
+	n := len(h)
+	for {
+		first := 3*i + 1
+		if first >= n {
+			break
+		}
+		// Find the worst (least) child among up to three children.
+		worst := first
+		if s := first + 1; s < n && c.compare(h[s], h[worst]) > 0 {
+			worst = s
+		}
+		if t := first + 2; t < n && c.compare(h[t], h[worst]) > 0 {
+			worst = t
+		}
+		if c.compare(h[worst], h[i]) > 0 { // worst child is worse than current → swap
+			h[i], h[worst] = h[worst], h[i]
+			i = worst
+		} else {
+			break
+		}
+	}
+}
+
+func (c *collectStoreHeap) Final(skip int, fixup collectorFixup) (search.DocumentMatchCollection, error) {
+	count := len(c.heap)
+	size := count - skip
+	if size <= 0 {
+		return make(search.DocumentMatchCollection, 0), nil
+	}
+	// Sort in-place so heap[0] = best doc (compare < 0 means "better").
+	// pdqsort (sort.Slice) has much better cache behavior than repeated
+	// removeLast (heapsort) because it accesses the array sequentially,
+	// avoiding the scattered pointer dereferences heapsort pays at each level.
+	sort.Slice(c.heap, func(i, j int) bool {
+		return c.compare(c.heap[i], c.heap[j]) < 0
+	})
+	rv := make(search.DocumentMatchCollection, size)
+	for i := 0; i < size; i++ {
+		doc := c.heap[skip+i]
+		rv[i] = doc
+		if err := fixup(doc); err != nil {
+			return nil, err
+		}
+	}
+	return rv, nil
+}
+
+func (c *collectStoreHeap) Len() int {
+	return len(c.heap)
+}
+
+func (c *collectStoreHeap) Internal() search.DocumentMatchCollection {
+	return c.heap
+}

search/collector/heap_test.go

@@ -0,0 +1,202 @@
+//  Copyright (c) 2026 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"testing"
+
+	"github.com/blevesearch/bleve/v2/search"
+)
+
+// scoreDesc is a collectorCompare for score-descending order.
+// Returns positive when a.Score < b.Score (a is worse — lower score — and
+// should sit closer to the root of the min-heap).
+func scoreDesc(a, b *search.DocumentMatch) int {
+	if a.Score < b.Score {
+		return 1
+	}
+	if a.Score > b.Score {
+		return -1
+	}
+	return 0
+}
+
+func makeScoreDoc(score float64) *search.DocumentMatch {
+	return &search.DocumentMatch{Score: score}
+}
+
+// checkTernaryInvariant verifies the ternary min-heap property: for every node
+// i the parent (at (i-1)/3) must be as bad or worse than the child, i.e.
+// compare(parent, child) >= 0.
+func checkTernaryInvariant(t *testing.T, h *collectStoreHeap) {
+	t.Helper()
+	for i := 1; i < len(h.heap); i++ {
+		parent := (i - 1) / 3
+		if h.compare(h.heap[parent], h.heap[i]) < 0 {
+			t.Errorf("ternary heap invariant violated at index %d (score=%.2f): parent %d (score=%.2f) is better, should be at least as bad",
+				i, h.heap[i].Score, parent, h.heap[parent].Score)
+		}
+	}
+}
+
+// TestTernaryHeapInvariantAfterInserts inserts scores in several orderings and
+// verifies the ternary invariant holds after every insertion.
+func TestTernaryHeapInvariantAfterInserts(t *testing.T) {
+	for _, scores := range [][]float64{
+		{5, 3, 8, 1, 7, 2, 9, 4, 6, 10},  // unsorted
+		{1, 2, 3, 4, 5, 6, 7, 8, 9, 10},  // ascending
+		{10, 9, 8, 7, 6, 5, 4, 3, 2, 1},  // descending
+		{3, 3, 3, 3},                       // all equal
+	} {
+		h := newStoreHeap(len(scores), scoreDesc)
+		for _, s := range scores {
+			h.add(makeScoreDoc(s))
+			checkTernaryInvariant(t, h)
+		}
+		// Root must be the minimum score in the heap.
+		min := scores[0]
+		for _, s := range scores[1:] {
+			if s < min {
+				min = s
+			}
+		}
+		if h.heap[0].Score != min {
+			t.Errorf("root score %.2f, want min %.2f (scores %v)", h.heap[0].Score, min, scores)
+		}
+	}
+}
+
+// TestTernaryHeapRemoveLastAscending verifies that successive removeLast calls
+// return elements in ascending score order (worst first = heap-sort property).
+func TestTernaryHeapRemoveLastAscending(t *testing.T) {
+	scores := []float64{5, 3, 8, 1, 7, 2, 9, 4, 6, 10}
+	h := newStoreHeap(len(scores), scoreDesc)
+	for _, s := range scores {
+		h.add(makeScoreDoc(s))
+	}
+
+	prev := -1.0
+	for len(h.heap) > 0 {
+		doc := h.removeLast()
+		if doc.Score < prev {
+			t.Errorf("removeLast out of order: got %.2f after %.2f", doc.Score, prev)
+		}
+		prev = doc.Score
+		checkTernaryInvariant(t, h)
+	}
+}
+
+// TestTernaryHeapAddNotExceedingSize checks that size enforcement works:
+// adding more than k docs keeps the best k and evicts the rest.
+func TestTernaryHeapAddNotExceedingSize(t *testing.T) {
+	const k = 5
+	h := newStoreHeap(k, scoreDesc)
+	for i := 1; i <= 10; i++ {
+		evicted := h.AddNotExceedingSize(makeScoreDoc(float64(i)), k)
+		if i <= k {
+			if evicted != nil {
+				t.Errorf("insert %d: expected no eviction, got score %.2f", i, evicted.Score)
+			}
+		} else {
+			if evicted == nil {
+				t.Errorf("insert %d: expected eviction", i)
+			}
+		}
+		checkTernaryInvariant(t, h)
+	}
+	if h.Len() != k {
+		t.Errorf("heap size %d, want %d", h.Len(), k)
+	}
+	// The heap should hold the best k scores (6..10); root is their minimum.
+	if h.heap[0].Score != 6.0 {
+		t.Errorf("root score %.2f, want 6.00 (worst of top-%d)", h.heap[0].Score, k)
+	}
+}
+
+// TestTernaryHeapFinalOrder verifies that Final(0, ...) returns results in
+// descending score order (best first).
+func TestTernaryHeapFinalOrder(t *testing.T) {
+	scores := []float64{3, 1, 4, 1, 5, 9, 2, 6, 5, 3}
+	h := newStoreHeap(len(scores), scoreDesc)
+	for _, s := range scores {
+		h.add(makeScoreDoc(s))
+	}
+
+	fixup := func(*search.DocumentMatch) error { return nil }
+	result, err := h.Final(0, fixup)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(result) != len(scores) {
+		t.Fatalf("len(result)=%d, want %d", len(result), len(scores))
+	}
+	for i := 1; i < len(result); i++ {
+		if result[i].Score > result[i-1].Score {
+			t.Errorf("result[%d]=%.2f > result[%d]=%.2f (not sorted descending)",
+				i, result[i].Score, i-1, result[i-1].Score)
+		}
+	}
+}
+
+// TestTernaryHeapFinalWithSkip verifies that Final(skip, ...) skips the skip
+// worst results and returns the rest best-first.
+func TestTernaryHeapFinalWithSkip(t *testing.T) {
+	// Heap holds 1..5; skip=2 should return scores [3, 2, 1] (skipping 4 and 5).
+	scores := []float64{1, 2, 3, 4, 5}
+	h := newStoreHeap(len(scores), scoreDesc)
+	for _, s := range scores {
+		h.add(makeScoreDoc(s))
+	}
+
+	fixup := func(*search.DocumentMatch) error { return nil }
+	result, err := h.Final(2, fixup)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(result) != 3 {
+		t.Fatalf("len(result)=%d, want 3", len(result))
+	}
+	want := []float64{3, 2, 1}
+	for i, w := range want {
+		if result[i].Score != w {
+			t.Errorf("result[%d]=%.2f, want %.2f", i, result[i].Score, w)
+		}
+	}
+}
+
+// TestTernaryHeapLargeN exercises multiple levels of siftDown (n > 3^3 = 27).
+func TestTernaryHeapLargeN(t *testing.T) {
+	const n = 200
+	h := newStoreHeap(n, scoreDesc)
+	for i := n; i >= 1; i-- { // insert descending to stress siftUp
+		h.add(makeScoreDoc(float64(i)))
+		checkTernaryInvariant(t, h)
+	}
+
+	prev := -1.0
+	count := 0
+	for len(h.heap) > 0 {
+		doc := h.removeLast()
+		if doc.Score < prev {
+			t.Errorf("removeLast out of order at position %d: %.2f after %.2f", count, doc.Score, prev)
+		}
+		prev = doc.Score
+		count++
+		checkTernaryInvariant(t, h)
+	}
+	if count != n {
+		t.Errorf("extracted %d elements, want %d", count, n)
+	}
+}

search/collector/knn.go

@@ -0,0 +1,262 @@
+//  Copyright (c) 2023 Couchbase, Inc.
+//
+// Licensed 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.
+
+//go:build vectors
+// +build vectors
+
+package collector
+
+import (
+	"context"
+	"time"
+
+	"github.com/blevesearch/bleve/v2/search"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+type collectStoreKNN struct {
+	internalHeaps []collectorStore
+	kValues       []int64
+	allHits       map[*search.DocumentMatch]struct{}
+	ejectedDocs   map[*search.DocumentMatch]struct{}
+}
+
+func newStoreKNN(internalHeaps []collectorStore, kValues []int64) *collectStoreKNN {
+	return &collectStoreKNN{
+		internalHeaps: internalHeaps,
+		kValues:       kValues,
+		ejectedDocs:   make(map[*search.DocumentMatch]struct{}),
+		allHits:       make(map[*search.DocumentMatch]struct{}),
+	}
+}
+
+// Adds a document to the collector store and returns the documents that were ejected
+// from the store. The documents that were ejected from the store are the ones that
+// were not in the top K documents for any of the heaps.
+// These document are put back into the pool document match pool in the KNN Collector.
+func (c *collectStoreKNN) AddDocument(doc *search.DocumentMatch) []*search.DocumentMatch {
+	for heapIdx := 0; heapIdx < len(c.internalHeaps); heapIdx++ {
+		if _, ok := doc.ScoreBreakdown[heapIdx]; !ok {
+			continue
+		}
+		ejectedDoc := c.internalHeaps[heapIdx].AddNotExceedingSize(doc, int(c.kValues[heapIdx]))
+		if ejectedDoc != nil {
+			delete(ejectedDoc.ScoreBreakdown, heapIdx)
+			c.ejectedDocs[ejectedDoc] = struct{}{}
+		}
+	}
+	var rv []*search.DocumentMatch
+	for doc := range c.ejectedDocs {
+		if len(doc.ScoreBreakdown) == 0 {
+			rv = append(rv, doc)
+		}
+		// clear out the ejectedDocs map to reuse it in the next AddDocument call
+		delete(c.ejectedDocs, doc)
+	}
+	return rv
+}
+
+func (c *collectStoreKNN) Final(fixup collectorFixup) (search.DocumentMatchCollection, error) {
+	for _, heap := range c.internalHeaps {
+		for _, doc := range heap.Internal() {
+			// duplicates may be present across the internal heaps
+			// meaning the same document match may be in the top K
+			// for multiple KNN queries.
+			c.allHits[doc] = struct{}{}
+		}
+	}
+	size := len(c.allHits)
+	if size <= 0 {
+		return make(search.DocumentMatchCollection, 0), nil
+	}
+	rv := make(search.DocumentMatchCollection, size)
+	i := 0
+	for doc := range c.allHits {
+		if fixup != nil {
+			err := fixup(doc)
+			if err != nil {
+				return nil, err
+			}
+		}
+		rv[i] = doc
+		i++
+	}
+	return rv, nil
+}
+
+func MakeKNNDocMatchHandler(ctx *search.SearchContext) (search.DocumentMatchHandler, error) {
+	var hc *KNNCollector
+	var ok bool
+	if hc, ok = ctx.Collector.(*KNNCollector); ok {
+		return func(d *search.DocumentMatch) error {
+			if d == nil {
+				return nil
+			}
+			toRelease := hc.knnStore.AddDocument(d)
+			for _, doc := range toRelease {
+				ctx.DocumentMatchPool.Put(doc)
+			}
+			return nil
+		}, nil
+	}
+	return nil, nil
+}
+
+func GetNewKNNCollectorStore(kArray []int64) *collectStoreKNN {
+	internalHeaps := make([]collectorStore, len(kArray))
+	for knnIdx, k := range kArray {
+		// TODO - Check if the datatype of k can be made into an int instead of int64
+		idx := knnIdx
+		internalHeaps[idx] = getOptimalCollectorStore(int(k), 0, func(i, j *search.DocumentMatch) int {
+			if i.ScoreBreakdown[idx] < j.ScoreBreakdown[idx] {
+				return 1
+			}
+			return -1
+		})
+	}
+	return newStoreKNN(internalHeaps, kArray)
+}
+
+// implements Collector interface
+type KNNCollector struct {
+	knnStore *collectStoreKNN
+	size     int
+	total    uint64
+	took     time.Duration
+	results  search.DocumentMatchCollection
+	maxScore float64
+}
+
+func NewKNNCollector(kArray []int64, size int64) *KNNCollector {
+	return &KNNCollector{
+		knnStore: GetNewKNNCollectorStore(kArray),
+		size:     int(size),
+	}
+}
+
+func (hc *KNNCollector) Collect(ctx context.Context, searcher search.Searcher, reader index.IndexReader) error {
+	startTime := time.Now()
+	var err error
+	var next *search.DocumentMatch
+
+	// pre-allocate enough space in the DocumentMatchPool
+	// unless the sum of K is too large, then cap it
+	// everything should still work, just allocates DocumentMatches on demand
+	backingSize := hc.size
+	if backingSize > PreAllocSizeSkipCap {
+		backingSize = PreAllocSizeSkipCap + 1
+	}
+	searchContext := &search.SearchContext{
+		DocumentMatchPool: search.NewDocumentMatchPool(backingSize+searcher.DocumentMatchPoolSize(), 0),
+		Collector:         hc,
+		IndexReader:       reader,
+	}
+
+	dmHandlerMakerKNN := MakeKNNDocMatchHandler
+	if cv := ctx.Value(search.MakeKNNDocumentMatchHandlerKey); cv != nil {
+		dmHandlerMakerKNN = cv.(search.MakeKNNDocumentMatchHandler)
+	}
+	// use the application given builder for making the custom document match
+	// handler and perform callbacks/invocations on the newly made handler.
+	dmHandler, err := dmHandlerMakerKNN(searchContext)
+	if err != nil {
+		return err
+	}
+	select {
+	case <-ctx.Done():
+		search.RecordSearchCost(ctx, search.AbortM, 0)
+		return ctx.Err()
+	default:
+		next, err = searcher.Next(searchContext)
+	}
+	for err == nil && next != nil {
+		if hc.total%CheckDoneEvery == 0 {
+			select {
+			case <-ctx.Done():
+				search.RecordSearchCost(ctx, search.AbortM, 0)
+				return ctx.Err()
+			default:
+			}
+		}
+		hc.total++
+
+		err = dmHandler(next)
+		if err != nil {
+			break
+		}
+
+		next, err = searcher.Next(searchContext)
+	}
+	if err != nil {
+		return err
+	}
+
+	// help finalize/flush the results in case
+	// of custom document match handlers.
+	err = dmHandler(nil)
+	if err != nil {
+		return err
+	}
+
+	// compute search duration
+	hc.took = time.Since(startTime)
+
+	// finalize actual results
+	err = hc.finalizeResults(reader)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (hc *KNNCollector) finalizeResults(r index.IndexReader) error {
+	var err error
+	hc.results, err = hc.knnStore.Final(func(doc *search.DocumentMatch) error {
+		if doc.ID == "" {
+			// look up the id since we need it for lookup
+			var err error
+			doc.ID, err = r.ExternalID(doc.IndexInternalID)
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	return err
+}
+
+func (hc *KNNCollector) Results() search.DocumentMatchCollection {
+	return hc.results
+}
+
+func (hc *KNNCollector) Total() uint64 {
+	return hc.total
+}
+
+func (hc *KNNCollector) MaxScore() float64 {
+	return hc.maxScore
+}
+
+func (hc *KNNCollector) Took() time.Duration {
+	return hc.took
+}
+
+func (hc *KNNCollector) SetFacetsBuilder(facetsBuilder *search.FacetsBuilder) {
+	// facet unsupported for vector search
+}
+
+func (hc *KNNCollector) FacetResults() search.FacetResults {
+	// facet unsupported for vector search
+	return nil
+}

search/collector/nested.go

@@ -0,0 +1,103 @@
+//  Copyright (c) 2026 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"github.com/blevesearch/bleve/v2/search"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+type collectStoreNested struct {
+	// descAdder is used to customize how descendants are merged into their parent
+	descAdder search.DescendantAdderCallbackFn
+	// nested reader to retrieve ancestor information
+	nr index.NestedReader
+	// the current root document match being built
+	currRoot *search.DocumentMatch
+	// the ancestor ID of the current root document being built
+	currRootAncestorID index.AncestorID
+	// prealloc slice for ancestor IDs
+	ancestors []index.AncestorID
+}
+
+func newStoreNested(nr index.NestedReader, descAdder search.DescendantAdderCallbackFn) *collectStoreNested {
+	rv := &collectStoreNested{
+		descAdder: descAdder,
+		nr:        nr,
+	}
+	return rv
+}
+
+// ProcessNestedDocument adds a document to the nested store, merging it into its root document
+// as needed. If the returned DocumentMatch is nil, the incoming doc has been merged
+// into its parent and should not be processed further. If the returned DocumentMatch
+// is non-nil, it represents a complete root document that should be processed further.
+// NOTE: This implementation assumes that documents are added in increasing order of their internal IDs
+// which is guaranteed by all searchers in bleve.
+func (c *collectStoreNested) ProcessNestedDocument(ctx *search.SearchContext, doc *search.DocumentMatch) (*search.DocumentMatch, error) {
+	// find ancestors for the doc
+	var err error
+	c.ancestors, err = c.nr.Ancestors(doc.IndexInternalID, c.ancestors[:0])
+	if err != nil {
+		return nil, err
+	}
+	if len(c.ancestors) == 0 {
+		// should not happen, every doc should have at least itself as ancestor
+		return nil, nil
+	}
+	// root docID is the last ancestor
+	rootID := c.ancestors[len(c.ancestors)-1]
+	// check if there is an interim root already and if the incoming doc belongs to it
+	if c.currRoot != nil && c.currRootAncestorID.Equals(rootID) {
+		// there is an interim root already, and the incoming doc belongs to it
+		if err := c.descAdder(c.currRoot, doc); err != nil {
+			return nil, err
+		}
+		// recycle the child document now that it's merged into the interim root
+		ctx.DocumentMatchPool.Put(doc)
+		return nil, nil
+	}
+	// completedRoot is the root document match to return, if any
+	var completedRoot *search.DocumentMatch
+	if c.currRoot != nil {
+		// we have an existing interim root, return it for processing
+		completedRoot = c.currRoot
+	}
+	// no interim root for now so either we have a root document incoming
+	// or we have a child doc and need to create an interim root
+	if len(c.ancestors) == 1 {
+		// incoming doc is the root itself
+		c.currRoot = doc
+		c.currRootAncestorID = rootID
+		return completedRoot, nil
+	}
+	// this is a child doc, create interim root
+	newDM := ctx.DocumentMatchPool.Get()
+	newDM.IndexInternalID = rootID.ToIndexInternalID(newDM.IndexInternalID)
+	// merge the incoming doc into the new interim root
+	c.currRoot = newDM
+	c.currRootAncestorID = rootID
+	if err := c.descAdder(c.currRoot, doc); err != nil {
+		return nil, err
+	}
+	// recycle the child document now that it's merged into the interim root
+	ctx.DocumentMatchPool.Put(doc)
+	return completedRoot, nil
+}
+
+// Current returns the current interim root document match being built, if any
+func (c *collectStoreNested) Current() *search.DocumentMatch {
+	return c.currRoot
+}

search/collector/search_test.go

@@ -0,0 +1,187 @@
+//  Copyright (c) 2014 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"context"
+	"reflect"
+
+	"github.com/blevesearch/bleve/v2/search"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+type stubSearcher struct {
+	index   int
+	matches []*search.DocumentMatch
+}
+
+func (ss *stubSearcher) SetBytesRead(val uint64) {
+
+}
+
+func (ss *stubSearcher) BytesRead() uint64 {
+	return 0
+}
+
+func (ss *stubSearcher) Size() int {
+	sizeInBytes := int(reflect.TypeOf(*ss).Size())
+
+	for _, entry := range ss.matches {
+		if entry != nil {
+			sizeInBytes += entry.Size()
+		}
+	}
+
+	return sizeInBytes
+}
+
+func (ss *stubSearcher) Next(ctx *search.SearchContext) (*search.DocumentMatch, error) {
+	if ss.index < len(ss.matches) {
+		rv := ctx.DocumentMatchPool.Get()
+		rv.IndexInternalID = ss.matches[ss.index].IndexInternalID
+		rv.Score = ss.matches[ss.index].Score
+		ss.index++
+		return rv, nil
+	}
+	return nil, nil
+}
+
+func (ss *stubSearcher) Advance(ctx *search.SearchContext, ID index.IndexInternalID) (*search.DocumentMatch, error) {
+
+	for ss.index < len(ss.matches) && ss.matches[ss.index].IndexInternalID.Compare(ID) < 0 {
+		ss.index++
+	}
+	if ss.index < len(ss.matches) {
+		rv := ctx.DocumentMatchPool.Get()
+		rv.IndexInternalID = ss.matches[ss.index].IndexInternalID
+		rv.Score = ss.matches[ss.index].Score
+		ss.index++
+		return rv, nil
+	}
+	return nil, nil
+}
+
+func (ss *stubSearcher) Close() error {
+	return nil
+}
+
+func (ss *stubSearcher) Weight() float64 {
+	return 0.0
+}
+
+func (ss *stubSearcher) SetQueryNorm(float64) {
+}
+
+func (ss *stubSearcher) Count() uint64 {
+	return uint64(len(ss.matches))
+}
+
+func (ss *stubSearcher) Min() int {
+	return 0
+}
+
+func (ss *stubSearcher) DocumentMatchPoolSize() int {
+	return 0
+}
+
+type stubReader struct{}
+
+func (sr *stubReader) Size() int {
+	return 0
+}
+
+func (sr *stubReader) TermFieldReader(ctx context.Context, term []byte, field string, includeFreq, includeNorm, includeTermVectors bool) (index.TermFieldReader, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) DocIDReaderAll() (index.DocIDReader, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) DocIDReaderOnly(ids []string) (index.DocIDReader, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) FieldDict(field string) (index.FieldDict, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) FieldDictRange(field string, startTerm []byte, endTerm []byte) (index.FieldDict, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) FieldDictPrefix(field string, termPrefix []byte) (index.FieldDict, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) Document(id string) (index.Document, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) DocumentVisitFieldTerms(id index.IndexInternalID, fields []string, visitor index.DocValueVisitor) error {
+	return nil
+}
+
+func (sr *stubReader) Fields() ([]string, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) GetInternal(key []byte) ([]byte, error) {
+	return nil, nil
+}
+
+func (sr *stubReader) DocCount() (uint64, error) {
+	return 0, nil
+}
+
+func (sr *stubReader) ExternalID(id index.IndexInternalID) (string, error) {
+	return string(id), nil
+}
+
+func (sr *stubReader) InternalID(id string) (index.IndexInternalID, error) {
+	return []byte(id), nil
+}
+
+func (sr *stubReader) DumpAll() chan interface{} {
+	return nil
+}
+
+func (sr *stubReader) DumpDoc(id string) chan interface{} {
+	return nil
+}
+
+func (sr *stubReader) DumpFields() chan interface{} {
+	return nil
+}
+
+func (sr *stubReader) Close() error {
+	return nil
+}
+
+func (sr *stubReader) DocValueReader(fields []string) (index.DocValueReader, error) {
+	return &DocValueReader{i: sr, fields: fields}, nil
+}
+
+type DocValueReader struct {
+	i      *stubReader
+	fields []string
+}
+
+func (dvr *DocValueReader) VisitDocValues(id index.IndexInternalID, visitor index.DocValueVisitor) error {
+	return dvr.i.DocumentVisitFieldTerms(id, dvr.fields, visitor)
+}
+func (dvr *DocValueReader) BytesRead() uint64 {
+	return 0
+}

search/collector/slice.go

@@ -0,0 +1,83 @@
+//  Copyright (c) 2014 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"github.com/blevesearch/bleve/v2/search"
+)
+
+type collectStoreSlice struct {
+	slice   search.DocumentMatchCollection
+	compare collectorCompare
+}
+
+func newStoreSlice(capacity int, compare collectorCompare) *collectStoreSlice {
+	rv := &collectStoreSlice{
+		slice:   make(search.DocumentMatchCollection, 0, capacity),
+		compare: compare,
+	}
+	return rv
+}
+
+func (c *collectStoreSlice) AddNotExceedingSize(doc *search.DocumentMatch,
+	size int) *search.DocumentMatch {
+	c.add(doc)
+	if c.len() > size {
+		return c.removeLast()
+	}
+	return nil
+}
+
+func (c *collectStoreSlice) add(doc *search.DocumentMatch) {
+	// find where to insert, starting at end (lowest)
+	i := len(c.slice)
+	for ; i > 0; i-- {
+		cmp := c.compare(doc, c.slice[i-1])
+		if cmp >= 0 {
+			break
+		}
+	}
+	// insert at i
+	c.slice = append(c.slice, nil)
+	copy(c.slice[i+1:], c.slice[i:])
+	c.slice[i] = doc
+}
+
+func (c *collectStoreSlice) removeLast() *search.DocumentMatch {
+	var rv *search.DocumentMatch
+	rv, c.slice = c.slice[len(c.slice)-1], c.slice[:len(c.slice)-1]
+	return rv
+}
+
+func (c *collectStoreSlice) Final(skip int, fixup collectorFixup) (search.DocumentMatchCollection, error) {
+	for i := skip; i < len(c.slice); i++ {
+		err := fixup(c.slice[i])
+		if err != nil {
+			return nil, err
+		}
+	}
+	if skip <= len(c.slice) {
+		return c.slice[skip:], nil
+	}
+	return search.DocumentMatchCollection{}, nil
+}
+
+func (c *collectStoreSlice) Internal() search.DocumentMatchCollection {
+	return c.slice
+}
+
+func (c *collectStoreSlice) len() int {
+	return len(c.slice)
+}

search/collector/topn.go

@@ -0,0 +1,739 @@
+//  Copyright (c) 2014 Couchbase, Inc.
+//
+// Licensed 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 collector
+
+import (
+	"context"
+	"reflect"
+	"strconv"
+	"time"
+
+	"github.com/blevesearch/bleve/v2/numeric"
+	"github.com/blevesearch/bleve/v2/search"
+	"github.com/blevesearch/bleve/v2/size"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+var reflectStaticSizeTopNCollector int
+
+func init() {
+	var coll TopNCollector
+	reflectStaticSizeTopNCollector = int(reflect.TypeOf(coll).Size())
+}
+
+type collectorStore interface {
+	// Add the document, and if the new store size exceeds the provided size
+	// the last element is removed and returned.  If the size has not been
+	// exceeded, nil is returned.
+	AddNotExceedingSize(doc *search.DocumentMatch, size int) *search.DocumentMatch
+
+	Final(skip int, fixup collectorFixup) (search.DocumentMatchCollection, error)
+
+	// Provide access the internal heap implementation
+	Internal() search.DocumentMatchCollection
+}
+
+// PreAllocSizeSkipCap will cap preallocation to this amount when
+// size+skip exceeds this value
+var PreAllocSizeSkipCap = 1000
+
+type collectorCompare func(i, j *search.DocumentMatch) int
+
+type collectorFixup func(d *search.DocumentMatch) error
+
+// TopNCollector collects the top N hits, optionally skipping some results
+type TopNCollector struct {
+	size          int
+	skip          int
+	total         uint64
+	bytesRead     uint64
+	maxScore      float64
+	took          time.Duration
+	sort          search.SortOrder
+	results       search.DocumentMatchCollection
+	facetsBuilder *search.FacetsBuilder
+
+	store collectorStore
+	cmp   collectorCompare // specialized or generic; shared by heap + dmHandler
+
+	needDocIds    bool
+	neededFields  []string
+	cachedScoring []bool
+	cachedDesc    []bool
+
+	lowestMatchOutsideResults *search.DocumentMatch
+	updateFieldVisitor        index.DocValueVisitor
+	dvReader                  index.DocValueReader
+	searchAfter               *search.DocumentMatch
+
+	knnHits             map[string]*search.DocumentMatch
+	hybridMergeCallback search.HybridMergeCallbackFn
+
+	nestedStore *collectStoreNested
+
+	fastPrepare bool
+
+	earlyStopN   int
+	earlyStopped bool
+}
+
+// CheckDoneEvery controls how frequently we check the context deadline
+const CheckDoneEvery = uint64(1024)
+
+// NewTopNCollector builds a collector to find the top 'size' hits
+// skipping over the first 'skip' hits
+// ordering hits by the provided sort order
+func NewTopNCollector(size int, skip int, sort search.SortOrder) *TopNCollector {
+	return newTopNCollector(size, skip, sort, nil)
+}
+
+// NewTopNCollectorAfter builds a collector to find the top 'size' hits
+// skipping over the first 'skip' hits
+// ordering hits by the provided sort order
+// starting after the provided 'after' sort values
+func NewTopNCollectorAfter(size int, sort search.SortOrder, after []string) *TopNCollector {
+	rv := newTopNCollector(size, 0, sort, nil)
+	rv.searchAfter = createSearchAfterDocument(sort, after)
+	return rv
+}
+
+// NewNestedTopNCollector builds a collector to find the top 'size' hits
+// skipping over the first 'skip' hits
+// ordering hits by the provided sort order
+// while ensuring the nested documents are handled correctly
+// (i.e. parent document is returned instead of nested document)
+func NewNestedTopNCollector(size int, skip int, sort search.SortOrder, nr index.NestedReader) *TopNCollector {
+	return newTopNCollector(size, skip, sort, nr)
+}
+
+// NewNestedTopNCollectorAfter builds a collector to find the top 'size' hits
+// skipping over the first 'skip' hits
+// ordering hits by the provided sort order
+// starting after the provided 'after' sort values
+// while ensuring the nested documents are handled correctly
+// (i.e. parent document is returned instead of nested document)
+func NewNestedTopNCollectorAfter(size int, sort search.SortOrder, after []string, nr index.NestedReader) *TopNCollector {
+	rv := newTopNCollector(size, 0, sort, nr)
+	rv.searchAfter = createSearchAfterDocument(sort, after)
+	return rv
+}
+
+func newTopNCollector(size int, skip int, sort search.SortOrder, nr index.NestedReader) *TopNCollector {
+	hc := &TopNCollector{
+		size:          size,
+		skip:          skip,
+		sort:          sort,
+		neededFields:  sort.RequiredFields(),
+		cachedScoring: sort.CacheIsScore(),
+		cachedDesc:    sort.CacheDescending(),
+		needDocIds:    sort.RequiresDocID(),
+	}
+
+	hc.cmp = getOptimalCollectorCompare(hc)
+	hc.store = getOptimalCollectorStore(size, skip, hc.cmp)
+
+	if nr != nil {
+		descAdder := func(parent, child *search.DocumentMatch) error {
+			// add descendant score to parent score
+			parent.Score += child.Score
+			// merge explanations
+			parent.Expl = parent.Expl.MergeWith(child.Expl)
+			// merge field term locations
+			parent.FieldTermLocations = search.MergeFieldTermLocationsFromMatch(parent.FieldTermLocations, child)
+			// add child's ID to parent's Descendants
+			// add other as descendant only if it is not the same document
+			if !parent.IndexInternalID.Equals(child.IndexInternalID) {
+				// Add a copy of child.IndexInternalID to descendants, because
+				// child.IndexInternalID will be reset when 'child' is recycled.
+				var descendantID index.IndexInternalID
+				// first check if parent's descendants slice has capacity to reuse
+				if len(parent.Descendants) < cap(parent.Descendants) {
+					// reuse the buffer element at len(parent.Descendants)
+					descendantID = parent.Descendants[:len(parent.Descendants)+1][len(parent.Descendants)]
+				}
+				// copy the contents of id into descendantID, allocating if needed
+				parent.Descendants = append(parent.Descendants, index.NewIndexInternalIDFrom(descendantID, child.IndexInternalID))
+			}
+			return nil
+		}
+		hc.nestedStore = newStoreNested(nr, search.DescendantAdderCallbackFn(descAdder))
+	}
+
+	return hc
+}
+
+func getOptimalCollectorCompare(hc *TopNCollector) collectorCompare {
+	if len(hc.sort) == 1 && hc.cachedScoring[0] && hc.cachedDesc[0] {
+		return search.CompareScoreDescending
+	}
+	return func(i, j *search.DocumentMatch) int {
+		return hc.sort.Compare(hc.cachedScoring, hc.cachedDesc, i, j)
+	}
+}
+
+// Creates a dummy document to compare with for pagination.
+func createSearchAfterDocument(sort search.SortOrder, after []string) *search.DocumentMatch {
+	encodedAfter := make([]string, len(after))
+	for i, ss := range sort {
+		encodedAfter[i] = encodeSearchAfter(ss, after[i])
+	}
+
+	rv := &search.DocumentMatch{
+		Sort: encodedAfter,
+	}
+	for pos, ss := range sort {
+		if ss.RequiresDocID() {
+			rv.ID = after[pos]
+		}
+		if ss.RequiresScoring() {
+			if score, err := strconv.ParseFloat(after[pos], 64); err == nil {
+				rv.Score = score
+			}
+		}
+	}
+	return rv
+}
+
+// encodeSearchAfter applies prefix-coding to SearchAfter
+// if required to enable pagination on numeric, datetime,
+// and geo fields
+func encodeSearchAfter(ss search.SearchSort, after string) string {
+	encodeFloat := func() string {
+		f64, _ := strconv.ParseFloat(after, 64) // error checking in SearchRequest.Validate
+		i64 := numeric.Float64ToInt64(f64)
+		return string(numeric.MustNewPrefixCodedInt64(i64, 0))
+	}
+
+	encodeDate := func() string {
+		t, _ := time.Parse(time.RFC3339Nano, after) // error checking in SearchRequest.Validate
+		i64 := t.UnixNano()
+		return string(numeric.MustNewPrefixCodedInt64(i64, 0))
+	}
+
+	switch ss := ss.(type) {
+	case *search.SortGeoDistance:
+		return encodeFloat()
+	case *search.SortField:
+		switch ss.Type {
+		case search.SortFieldAsNumber:
+			return encodeFloat()
+		case search.SortFieldAsDate:
+			return encodeDate()
+		default:
+			// For SortFieldAsString and SortFieldAuto
+			// NOTE: SortFieldAuto is used if you set Sort with a string
+			// or if the type of the field is not set in the object
+			// in the Sort slice. We cannot perform type inference in
+			// this case, so we return the original string, even if
+			// its actually numeric or date.
+			return after
+		}
+	default:
+		// For SortDocID and SortScore
+		return after
+	}
+}
+
+// Filter document matches based on the SearchAfter field in the SearchRequest.
+func FilterHitsBySearchAfter(hits []*search.DocumentMatch, sort search.SortOrder, after []string) []*search.DocumentMatch {
+	if len(hits) == 0 {
+		return hits
+	}
+	// create a search after document
+	searchAfter := createSearchAfterDocument(sort, after)
+	// filter the hits
+	idx := 0
+	cachedScoring := sort.CacheIsScore()
+	cachedDesc := sort.CacheDescending()
+	for _, hit := range hits {
+		if sort.Compare(cachedScoring, cachedDesc, hit, searchAfter) > 0 {
+			hits[idx] = hit
+			idx++
+		}
+	}
+	return hits[:idx]
+}
+
+func getOptimalCollectorStore(size, skip int, comparator collectorCompare) collectorStore {
+	// pre-allocate space on the store to avoid reslicing
+	// unless the size + skip is too large, then cap it
+	// everything should still work, just reslices as necessary
+	backingSize := size + skip + 1
+	if size+skip > PreAllocSizeSkipCap {
+		backingSize = PreAllocSizeSkipCap + 1
+	}
+
+	if size+skip > 10 {
+		return newStoreHeap(backingSize, comparator)
+	} else {
+		return newStoreSlice(backingSize, comparator)
+	}
+}
+
+func (hc *TopNCollector) Size() int {
+	sizeInBytes := reflectStaticSizeTopNCollector + size.SizeOfPtr
+
+	if hc.facetsBuilder != nil {
+		sizeInBytes += hc.facetsBuilder.Size()
+	}
+
+	for _, entry := range hc.neededFields {
+		sizeInBytes += len(entry) + size.SizeOfString
+	}
+
+	sizeInBytes += len(hc.cachedScoring) + len(hc.cachedDesc)
+
+	return sizeInBytes
+}
+
+// Collect goes to the index to find the matching documents
+func (hc *TopNCollector) Collect(ctx context.Context, searcher search.Searcher, reader index.IndexReader) error {
+	startTime := time.Now()
+	var err error
+	var next *search.DocumentMatch
+
+	// pre-allocate enough space in the DocumentMatchPool
+	// unless the size + skip is too large, then cap it
+	// everything should still work, just allocates DocumentMatches on demand
+	backingSize := hc.size + hc.skip + 1
+	if hc.size+hc.skip > PreAllocSizeSkipCap {
+		backingSize = PreAllocSizeSkipCap + 1
+	}
+	searchContext := &search.SearchContext{
+		DocumentMatchPool: search.NewDocumentMatchPool(backingSize+searcher.DocumentMatchPoolSize(), len(hc.sort)),
+		Collector:         hc,
+		IndexReader:       reader,
+	}
+
+	hc.dvReader, err = reader.DocValueReader(hc.neededFields)
+	if err != nil {
+		return err
+	}
+
+	hc.updateFieldVisitor = func(field string, term []byte) {
+		if hc.facetsBuilder != nil {
+			hc.facetsBuilder.UpdateVisitor(field, term)
+		}
+		hc.sort.UpdateVisitor(field, term)
+	}
+
+	dmHandlerMaker := MakeTopNDocumentMatchHandler
+	if cv := ctx.Value(search.MakeDocumentMatchHandlerKey); cv != nil {
+		dmHandlerMaker = cv.(search.MakeDocumentMatchHandler)
+	}
+	// use the application given builder for making the custom document match
+	// handler and perform callbacks/invocations on the newly made handler.
+	dmHandler, loadID, err := dmHandlerMaker(searchContext)
+	if err != nil {
+		return err
+	}
+
+	hc.needDocIds = hc.needDocIds || loadID
+
+	// fastPrepare is set to true when a hit needs only basicPrepare plus the shared
+	// score sort value (see canFastPrepare). Applies only to score-sorted queries
+	//  with no field-loading needs.
+	hc.fastPrepare = hc.canFastPrepare()
+	select {
+	case <-ctx.Done():
+		search.RecordSearchCost(ctx, search.AbortM, 0)
+		return ctx.Err()
+	default:
+		next, err = searcher.Next(searchContext)
+	}
+	// use a local totalDocs for counting total docs seen
+	// for context deadline checking, as hc.total is only
+	// incremented for actual(root) collected documents, and
+	// we need to check deadline for every document seen (root or nested)
+	var totalDocs uint64
+	for err == nil && next != nil {
+		if totalDocs%CheckDoneEvery == 0 {
+			select {
+			case <-ctx.Done():
+				search.RecordSearchCost(ctx, search.AbortM, 0)
+				return ctx.Err()
+			default:
+			}
+		}
+		totalDocs++
+		if hc.nestedStore != nil {
+			// This may be a nested document — add it to the nested store first.
+			// If the nested store returns nil, the document was merged into its parent
+			// and should not be processed further.
+			// If it returns a non-nil document, it represents a complete root document
+			// and should be processed further.
+			next, err = hc.nestedStore.ProcessNestedDocument(searchContext, next)
+			if err != nil {
+				break
+			}
+		}
+		if next != nil {
+			if hc.knnHits != nil {
+				err = hc.adjustKNNDocumentMatch(searchContext, reader, next)
+				if err != nil {
+					break
+				}
+			}
+			hc.basicPrepare(next)
+			if hc.fastPrepare {
+				next.Sort = sortByScoreOpt
+			} else {
+				err = hc.prepareDocumentMatch(searchContext, reader, next)
+				if err != nil {
+					break
+				}
+			}
+			err = dmHandler(next)
+			if err != nil {
+				break
+			}
+			if hc.earlyStopN > 0 && hc.total >= uint64(hc.earlyStopN) {
+				hc.earlyStopped = true
+				break
+			}
+		}
+		next, err = searcher.Next(searchContext)
+	}
+	if err != nil {
+		return err
+	}
+
+	// if we have a nested store, we may have an interim root
+	// that needs to be returned for processing
+	if hc.nestedStore != nil {
+		currRoot := hc.nestedStore.Current()
+		if currRoot != nil {
+			if hc.knnHits != nil {
+				err = hc.adjustKNNDocumentMatch(searchContext, reader, currRoot)
+				if err != nil {
+					return err
+				}
+			}
+			// no descendants at this point
+			hc.basicPrepare(currRoot)
+			if hc.fastPrepare {
+				currRoot.Sort = sortByScoreOpt
+			} else {
+				err = hc.prepareDocumentMatch(searchContext, reader, currRoot)
+				if err != nil {
+					return err
+				}
+			}
+
+			err = dmHandler(currRoot)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	if hc.knnHits != nil {
+		// we may have some knn hits left that did not match any of the top N tf-idf hits
+		// we need to add them to the collector store to consider them as well.
+		for _, knnDoc := range hc.knnHits {
+			err = hc.prepareKNNDocumentMatch(searchContext, reader, knnDoc)
+			if err != nil {
+				return err
+			}
+			err = dmHandler(knnDoc)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	statsCallbackFn := ctx.Value(search.SearchIOStatsCallbackKey)
+	if statsCallbackFn != nil {
+		// hc.bytesRead corresponds to the
+		// total bytes read as part of docValues being read every hit
+		// which must be accounted by invoking the callback.
+		statsCallbackFn.(search.SearchIOStatsCallbackFunc)(hc.bytesRead)
+
+		search.RecordSearchCost(ctx, search.AddM, hc.bytesRead)
+	}
+
+	// help finalize/flush the results in case
+	// of custom document match handlers.
+	err = dmHandler(nil)
+	if err != nil {
+		return err
+	}
+
+	// compute search duration
+	hc.took = time.Since(startTime)
+
+	// finalize actual results
+	err = hc.finalizeResults(reader)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+var sortByScoreOpt = []string{"_score"}
+
+func (hc *TopNCollector) adjustKNNDocumentMatch(ctx *search.SearchContext,
+	reader index.IndexReader, d *search.DocumentMatch) (err error) {
+	d.ID, err = reader.ExternalID(d.IndexInternalID)
+	if err != nil {
+		return err
+	}
+	if knnHit, ok := hc.knnHits[d.ID]; ok {
+		// merge this document's hit with its knn score, expl etc. and remove it from the map
+		hc.hybridMergeCallback(d, knnHit)
+		delete(hc.knnHits, d.ID)
+	}
+	return nil
+}
+
+func (hc *TopNCollector) basicPrepare(d *search.DocumentMatch) {
+	hc.total++
+	d.HitNumber = hc.total
+	if d.Score > hc.maxScore {
+		hc.maxScore = d.Score
+	}
+}
+
+func (hc *TopNCollector) canFastPrepare() bool {
+	return len(hc.neededFields) == 0 && !hc.needDocIds &&
+		len(hc.sort) == 1 && hc.cachedScoring[0]
+}
+
+// prepareDocumentMatch does the non-fast preparation for a regular (non-KNN)
+// hit: visiting field terms required for sort/facets, loading the docID if
+// needed, and computing the sort value. basicPrepare must be called first.
+func (hc *TopNCollector) prepareDocumentMatch(ctx *search.SearchContext,
+	reader index.IndexReader, d *search.DocumentMatch) (err error) {
+
+	// visit field terms for features that require it (sort, facets)
+	if len(hc.neededFields) > 0 {
+		err = hc.visitFieldTerms(reader, d, hc.updateFieldVisitor)
+		if err != nil {
+			return err
+		}
+	}
+
+	// see if we need to load ID (at this early stage, for example to sort on it)
+	if hc.needDocIds && d.ID == "" {
+		d.ID, err = reader.ExternalID(d.IndexInternalID)
+		if err != nil {
+			return err
+		}
+	}
+
+	// compute this hits sort value
+	if len(hc.sort) == 1 && hc.cachedScoring[0] {
+		d.Sort = sortByScoreOpt
+	} else {
+		hc.sort.Value(d)
+	}
+
+	return nil
+}
+
+// prepareKNNDocumentMatch prepares a KNN hit. Its sort value was already
+// computed by the KNN collector, so only facet field visiting (if any) and the
+// basic bookkeeping are needed.
+func (hc *TopNCollector) prepareKNNDocumentMatch(ctx *search.SearchContext,
+	reader index.IndexReader, d *search.DocumentMatch) (err error) {
+
+	if hc.facetsBuilder != nil {
+		err = hc.visitFieldTerms(reader, d, func(field string, term []byte) {
+			hc.facetsBuilder.UpdateVisitor(field, term)
+		})
+		if err != nil {
+			return err
+		}
+	}
+
+	hc.basicPrepare(d)
+	return nil
+}
+
+func MakeTopNDocumentMatchHandler(
+	ctx *search.SearchContext) (search.DocumentMatchHandler, bool, error) {
+	var hc *TopNCollector
+	var ok bool
+	if hc, ok = ctx.Collector.(*TopNCollector); ok {
+		return func(d *search.DocumentMatch) error {
+			if d == nil {
+				return nil
+			}
+
+			// support search after based pagination,
+			// if this hit is <= the search after sort key
+			// we should skip it
+			if hc.searchAfter != nil {
+				// exact sort order matches use hit number to break tie
+				// but we want to allow for exact match, so we pretend
+				hc.searchAfter.HitNumber = d.HitNumber
+				if hc.cmp(d, hc.searchAfter) <= 0 {
+					ctx.DocumentMatchPool.Put(d)
+					return nil
+				}
+			}
+
+			// optimization, we track lowest sorting hit already removed from heap
+			// with this one comparison, we can avoid all heap operations if
+			// this hit would have been added and then immediately removed
+			if hc.lowestMatchOutsideResults != nil {
+				if hc.cmp(d, hc.lowestMatchOutsideResults) >= 0 {
+					// this hit can't possibly be in the result set, so avoid heap ops
+					ctx.DocumentMatchPool.Put(d)
+					return nil
+				}
+			}
+
+			removed := hc.store.AddNotExceedingSize(d, hc.size+hc.skip)
+			if removed != nil {
+				if hc.lowestMatchOutsideResults == nil {
+					hc.lowestMatchOutsideResults = removed
+				} else {
+					cmp := hc.cmp(removed, hc.lowestMatchOutsideResults)
+					if cmp < 0 {
+						tmp := hc.lowestMatchOutsideResults
+						hc.lowestMatchOutsideResults = removed
+						ctx.DocumentMatchPool.Put(tmp)
+					}
+				}
+			}
+			return nil
+		}, false, nil
+	}
+	return nil, false, nil
+}
+
+// visitFieldTerms is responsible for visiting the field terms of the
+// search hit, and passing visited terms to the sort and facet builder
+func (hc *TopNCollector) visitFieldTerms(reader index.IndexReader, d *search.DocumentMatch, v index.DocValueVisitor) error {
+	if hc.facetsBuilder != nil {
+		hc.facetsBuilder.StartDoc()
+	}
+	if d.ID != "" && d.IndexInternalID == nil {
+		// this document may have been sent over as preSearchData and
+		// we need to look up the internal id to visit the doc values for it
+		var err error
+		d.IndexInternalID, err = reader.InternalID(d.ID)
+		if err != nil {
+			return err
+		}
+	}
+
+	// first visit descendants if any
+	for _, descID := range d.Descendants {
+		err := hc.dvReader.VisitDocValues(descID, v)
+		if err != nil {
+			return err
+		}
+	}
+	// now visit the doc values for this document
+	err := hc.dvReader.VisitDocValues(d.IndexInternalID, v)
+	if hc.facetsBuilder != nil {
+		hc.facetsBuilder.EndDoc()
+	}
+
+	hc.bytesRead += hc.dvReader.BytesRead()
+
+	return err
+}
+
+// SetFacetsBuilder registers a facet builder for this collector
+func (hc *TopNCollector) SetFacetsBuilder(facetsBuilder *search.FacetsBuilder) {
+	hc.facetsBuilder = facetsBuilder
+	fieldsRequiredForFaceting := facetsBuilder.RequiredFields()
+	// for each of these fields, append only if not already there in hc.neededFields.
+	for _, field := range fieldsRequiredForFaceting {
+		found := false
+		for _, neededField := range hc.neededFields {
+			if field == neededField {
+				found = true
+				break
+			}
+		}
+		if !found {
+			hc.neededFields = append(hc.neededFields, field)
+		}
+	}
+}
+
+// finalizeResults starts with the heap containing the final top size+skip
+// it now throws away the results to be skipped
+// and does final doc id lookup (if necessary)
+func (hc *TopNCollector) finalizeResults(r index.IndexReader) error {
+	var err error
+	hc.results, err = hc.store.Final(hc.skip, func(doc *search.DocumentMatch) error {
+		if doc.ID == "" {
+			// look up the id since we need it for lookup
+			var err error
+			doc.ID, err = r.ExternalID(doc.IndexInternalID)
+			if err != nil {
+				return err
+			}
+		}
+		doc.Complete(nil)
+		return nil
+	})
+
+	return err
+}
+
+// Results returns the collected hits
+func (hc *TopNCollector) Results() search.DocumentMatchCollection {
+	return hc.results
+}
+
+// Total returns the total number of hits
+func (hc *TopNCollector) Total() uint64 {
+	return hc.total
+}
+
+// SetEarlyStop makes Collect() stop pulling from the searcher once n hits have
+// been collected; n <= 0 disables. Callers must ensure unseen docs cannot
+// change the result.
+func (hc *TopNCollector) SetEarlyStop(n int) {
+	hc.earlyStopN = n
+}
+
+// EarlyStopped reports whether Collect() stopped early; if true, Total() is a
+// lower bound.
+func (hc *TopNCollector) EarlyStopped() bool {
+	return hc.earlyStopped
+}
+
+// MaxScore returns the maximum score seen across all the hits
+func (hc *TopNCollector) MaxScore() float64 {
+	return hc.maxScore
+}
+
+// Took returns the time spent collecting hits
+func (hc *TopNCollector) Took() time.Duration {
+	return hc.took
+}
+
+// FacetResults returns the computed facets results
+func (hc *TopNCollector) FacetResults() search.FacetResults {
+	if hc.facetsBuilder != nil {
+		return hc.facetsBuilder.Results()
+	}
+	return nil
+}
+
+func (hc *TopNCollector) SetKNNHits(knnHits search.DocumentMatchCollection, hybridMergeCallback search.HybridMergeCallbackFn) {
+	hc.knnHits = make(map[string]*search.DocumentMatch, len(knnHits))
+	for _, hit := range knnHits {
+		hc.knnHits[hit.ID] = hit
+	}
+	hc.hybridMergeCallback = hybridMergeCallback
+}