diff --git a/earlystop_test.go b/earlystop_test.go new file mode 100644 index 000000000..150dba78d --- /dev/null +++ b/earlystop_test.go @@ -0,0 +1,184 @@ +// 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 bleve + +import ( + "fmt" + "os" + "testing" + + "github.com/blevesearch/bleve/v2/mapping" + "github.com/blevesearch/bleve/v2/search/query" +) + +// buildEarlyStopIndex indexes n documents that all match "common", so a bounded +// scan has plenty of matches left undrained. +func buildEarlyStopIndex(t *testing.T, n int) Index { + t.Helper() + + dir, err := os.MkdirTemp("", "earlystop") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.RemoveAll(dir) }) + + im := mapping.NewIndexMapping() + im.DefaultAnalyzer = "standard" + idx, err := New(dir+"/i.bleve", im) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = idx.Close() }) + + batch := idx.NewBatch() + for i := 0; i < n; i++ { + if err := batch.Index(fmt.Sprintf("d%05d", i), map[string]interface{}{ + "body": fmt.Sprintf("common tag%d", i%10), + "num": float64(i), + }); err != nil { + t.Fatal(err) + } + } + if err := idx.Batch(batch); err != nil { + t.Fatal(err) + } + return idx +} + +func earlyStopQuery() query.Query { + q := query.NewTermQuery("common") + q.SetField("body") + return q +} + +// TestEarlyStopBoundedScan is the contract for the bounded scan: with +// Score="none" and a bounded Size, the request means "any Size+From matching +// docs", so collection may stop once that many hits are in hand instead of +// draining every match. +// +// Two things must hold. The caller must still get Size hits — stopping early must +// not lose results it asked for. And Total must be reported as a lower bound +// (TotalRelation "gte"), because the scan genuinely did not count the rest; a +// caller that reads Total as exact would otherwise be silently misled. +func TestEarlyStopBoundedScan(t *testing.T) { + const n = 5000 + idx := buildEarlyStopIndex(t, n) + + for _, size := range []int{1, 10, 100} { + req := NewSearchRequest(earlyStopQuery()) + req.Size = size + req.Score = ScoreNone + res, err := idx.Search(req) + if err != nil { + t.Fatal(err) + } + if len(res.Hits) != size { + t.Errorf("size=%d: got %d hits, want %d — early stop dropped results the "+ + "caller asked for", size, len(res.Hits), size) + } + if res.TotalRelation != TotalRelationGte { + t.Errorf("size=%d: TotalRelation=%q, want %q — Total is a lower bound once "+ + "the scan stops early, and saying otherwise misleads the caller", + size, res.TotalRelation, TotalRelationGte) + } + if res.Total > uint64(n) { + t.Errorf("size=%d: Total=%d exceeds the corpus size %d", size, res.Total, n) + } + } +} + +// TestEarlyStopDoesNotEngageWhenUnsafe pins the preconditions. Each of these +// requests depends on documents the bounded scan would never look at, so it must +// keep draining and report an exact Total. Getting any of these wrong is a silent +// wrong-answer bug, not a slowdown: +// +// facets every match must be counted into the facet buckets +// field sort the top-k by field value can lie anywhere in the match set +// SearchAfter the cursor position depends on the full ordering +// scoring Score != "none" means order depends on scores, not arrival +func TestEarlyStopDoesNotEngageWhenUnsafe(t *testing.T) { + const n = 2000 + idx := buildEarlyStopIndex(t, n) + + cases := []struct { + name string + tweak func(*SearchRequest) + }{ + {"facets", func(r *SearchRequest) { + r.Score = ScoreNone + r.AddFacet("tags", NewFacetRequest("body", 5)) + }}, + {"field-sort", func(r *SearchRequest) { + r.Score = ScoreNone + r.SortBy([]string{"num"}) + }}, + {"search-after", func(r *SearchRequest) { + r.Score = ScoreNone + r.SortBy([]string{"_id"}) + r.SearchAfter = []string{"d00010"} + }}, + {"scoring-enabled", func(r *SearchRequest) { + // Score defaults to full scoring; ordering depends on scores. + }}, + } + + for _, tc := range cases { + req := NewSearchRequest(earlyStopQuery()) + req.Size = 10 + tc.tweak(req) + res, err := idx.Search(req) + if err != nil { + t.Fatalf("%s: %v", tc.name, err) + } + if res.TotalRelation != TotalRelationEq { + t.Errorf("%s: TotalRelation=%q, want %q — the bounded scan engaged on a "+ + "request whose result depends on documents it would not visit", + tc.name, res.TotalRelation, TotalRelationEq) + } + if res.Total != uint64(n) { + t.Errorf("%s: Total=%d, want the exact %d", tc.name, res.Total, n) + } + } +} + +// TestEarlyStopHitsAreRealMatches guards the cheapest way for a bounded scan to be +// wrong: returning documents that do not match. Order is not part of the contract +// under Score="none", but membership is. +func TestEarlyStopHitsAreRealMatches(t *testing.T) { + idx := buildEarlyStopIndex(t, 1000) + + q := query.NewTermQuery("tag3") + q.SetField("body") + req := NewSearchRequest(q) + req.Size = 20 + req.Score = ScoreNone + res, err := idx.Search(req) + if err != nil { + t.Fatal(err) + } + if len(res.Hits) != 20 { + t.Fatalf("got %d hits, want 20", len(res.Hits)) + } + // tag3 was indexed on every doc where i%10 == 3. + for _, h := range res.Hits { + var i int + if _, err := fmt.Sscanf(h.ID, "d%05d", &i); err != nil { + t.Fatalf("unexpected id %q: %v", h.ID, err) + } + if i%10 != 3 { + t.Errorf("id %s does not match tag3 — the bounded scan returned a non-match", h.ID) + } + } +} diff --git a/index_impl.go b/index_impl.go index 1da655512..c781f1ce1 100644 --- a/index_impl.go +++ b/index_impl.go @@ -786,6 +786,19 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr return nil, err } + // score="none" + Size means "return any Size+From matching docs", so the + // collector may stop scanning early — provided nothing below depends on + // unseen matches (facets, KNN, pagination cursor, nested rollup, field sort). + if req.Score == ScoreNone && req.Size > 0 && + len(req.Facets) == 0 && + !requestHasKNN(req) && + req.SearchAfter == nil && !reverseQueryExecution && + len(req.Sort) == 1 && req.Sort[0].RequiresScoring() { + if nestedMode, ok := ctx.Value(search.NestedSearchKey).(bool); !ok || !nestedMode { + coll.SetEarlyStop(req.Size + req.From) + } + } + var knnHits []*search.DocumentMatch var skipKNNCollector bool @@ -1036,16 +1049,21 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr req.SearchAfter = nil } + totalRelation := TotalRelationEq + if coll.EarlyStopped() { + totalRelation = TotalRelationGte + } rv := &SearchResult{ Status: &SearchStatus{ Total: 1, Successful: 1, }, - Hits: hits, - Total: coll.Total(), - MaxScore: coll.MaxScore(), - Took: searchDuration, - Facets: coll.FacetResults(), + Hits: hits, + Total: coll.Total(), + TotalRelation: totalRelation, + MaxScore: coll.MaxScore(), + Took: searchDuration, + Facets: coll.FacetResults(), } // rescore if fusion flag is set diff --git a/search.go b/search.go index 708be0871..a754a5a56 100644 --- a/search.go +++ b/search.go @@ -525,6 +525,15 @@ func (ss *SearchStatus) Merge(other *SearchStatus) { } } +// TotalRelation constants describe the accuracy of SearchResult.Total. +const ( + // TotalRelationEq means Total is an exact count of all matching documents. + TotalRelationEq = "eq" + // TotalRelationGte means Total is a lower bound: collection stopped before + // draining every match, so the true match count is >= Total. + TotalRelationGte = "gte" +) + // A SearchResult describes the results of executing // a SearchRequest. // @@ -540,14 +549,15 @@ func (ss *SearchStatus) Merge(other *SearchStatus) { // Took - The time taken to execute the search. // Facets - The facet results for the search. type SearchResult struct { - Status *SearchStatus `json:"status"` - Request *SearchRequest `json:"request,omitempty"` - Hits search.DocumentMatchCollection `json:"hits"` - Total uint64 `json:"total_hits"` - Cost uint64 `json:"cost"` - MaxScore float64 `json:"max_score"` - Took time.Duration `json:"took"` - Facets search.FacetResults `json:"facets"` + Status *SearchStatus `json:"status"` + Request *SearchRequest `json:"request,omitempty"` + Hits search.DocumentMatchCollection `json:"hits"` + Total uint64 `json:"total_hits"` + TotalRelation string `json:"total_relation"` + Cost uint64 `json:"cost"` + MaxScore float64 `json:"max_score"` + Took time.Duration `json:"took"` + Facets search.FacetResults `json:"facets"` // special fields that are applicable only for search // results that are obtained from a presearch SynonymResult search.FieldTermSynonymMap `json:"synonym_result,omitempty"` @@ -675,6 +685,9 @@ func (sr *SearchResult) Merge(other *SearchResult) { sr.Status.Merge(other.Status) sr.Hits = append(sr.Hits, other.Hits...) sr.Total += other.Total + if other.TotalRelation == TotalRelationGte { + sr.TotalRelation = TotalRelationGte + } sr.Cost += other.Cost if other.MaxScore > sr.MaxScore { sr.MaxScore = other.MaxScore diff --git a/search/collector/heap.go b/search/collector/heap.go index cd662bcf9..6a547bb96 100644 --- a/search/collector/heap.go +++ b/search/collector/heap.go @@ -15,85 +15,126 @@ package collector import ( - "container/heap" + "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 { - rv := &collectStoreHeap{ + return &collectStoreHeap{ heap: make(search.DocumentMatchCollection, 0, capacity), compare: compare, } - heap.Init(rv) - return rv } func (c *collectStoreHeap) AddNotExceedingSize(doc *search.DocumentMatch, size int) *search.DocumentMatch { c.add(doc) - if c.Len() > size { + if len(c.heap) > size { return c.removeLast() } return nil } func (c *collectStoreHeap) add(doc *search.DocumentMatch) { - heap.Push(c, doc) + c.heap = append(c.heap, doc) + c.siftUp(len(c.heap) - 1) } func (c *collectStoreHeap) removeLast() *search.DocumentMatch { - return heap.Pop(c).(*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 := c.Len() + 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 := size - 1; i >= 0; i-- { - doc := heap.Pop(c).(*search.DocumentMatch) + for i := 0; i < size; i++ { + doc := c.heap[skip+i] rv[i] = doc - err := fixup(doc) - if err != nil { + if err := fixup(doc); err != nil { return nil, err } } return rv, nil } -func (c *collectStoreHeap) Internal() search.DocumentMatchCollection { - return c.heap -} - -// heap interface implementation - func (c *collectStoreHeap) Len() int { return len(c.heap) } -func (c *collectStoreHeap) Less(i, j int) bool { - so := c.compare(c.heap[i], c.heap[j]) - return -so < 0 -} - -func (c *collectStoreHeap) Swap(i, j int) { - c.heap[i], c.heap[j] = c.heap[j], c.heap[i] -} - -func (c *collectStoreHeap) Push(x interface{}) { - c.heap = append(c.heap, x.(*search.DocumentMatch)) -} - -func (c *collectStoreHeap) Pop() interface{} { - var rv *search.DocumentMatch - rv, c.heap = c.heap[len(c.heap)-1], c.heap[:len(c.heap)-1] - return rv +func (c *collectStoreHeap) Internal() search.DocumentMatchCollection { + return c.heap } diff --git a/search/collector/heap_test.go b/search/collector/heap_test.go new file mode 100644 index 000000000..31f1c4742 --- /dev/null +++ b/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) + } +} diff --git a/search/collector/list.go b/search/collector/list.go deleted file mode 100644 index f73505e7d..000000000 --- a/search/collector/list.go +++ /dev/null @@ -1,96 +0,0 @@ -// 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 ( - "container/list" - - "github.com/blevesearch/bleve/v2/search" -) - -type collectStoreList struct { - results *list.List - compare collectorCompare -} - -func newStoreList(capacity int, compare collectorCompare) *collectStoreList { - rv := &collectStoreList{ - results: list.New(), - compare: compare, - } - - return rv -} - -func (c *collectStoreList) AddNotExceedingSize(doc *search.DocumentMatch, size int) *search.DocumentMatch { - c.add(doc) - if c.len() > size { - return c.removeLast() - } - return nil -} - -func (c *collectStoreList) add(doc *search.DocumentMatch) { - for e := c.results.Front(); e != nil; e = e.Next() { - curr := e.Value.(*search.DocumentMatch) - if c.compare(doc, curr) >= 0 { - c.results.InsertBefore(doc, e) - return - } - } - // if we got to the end, we still have to add it - c.results.PushBack(doc) -} - -func (c *collectStoreList) removeLast() *search.DocumentMatch { - return c.results.Remove(c.results.Front()).(*search.DocumentMatch) -} - -func (c *collectStoreList) Final(skip int, fixup collectorFixup) (search.DocumentMatchCollection, error) { - if c.results.Len()-skip > 0 { - rv := make(search.DocumentMatchCollection, c.results.Len()-skip) - i := 0 - skipped := 0 - for e := c.results.Back(); e != nil; e = e.Prev() { - if skipped < skip { - skipped++ - continue - } - - rv[i] = e.Value.(*search.DocumentMatch) - err := fixup(rv[i]) - if err != nil { - return nil, err - } - i++ - } - return rv, nil - } - return search.DocumentMatchCollection{}, nil -} - -func (c *collectStoreList) Internal() search.DocumentMatchCollection { - rv := make(search.DocumentMatchCollection, c.results.Len()) - i := 0 - for e := c.results.Front(); e != nil; e = e.Next() { - rv[i] = e.Value.(*search.DocumentMatch) - i++ - } - return rv -} - -func (c *collectStoreList) len() int { - return c.results.Len() -} diff --git a/search/collector/topn.go b/search/collector/topn.go index d7fd27f23..296984fe7 100644 --- a/search/collector/topn.go +++ b/search/collector/topn.go @@ -84,6 +84,9 @@ type TopNCollector struct { nestedStore *collectStoreNested fastPrepare bool + + earlyStopN int + earlyStopped bool } // CheckDoneEvery controls how frequently we check the context deadline @@ -396,6 +399,10 @@ func (hc *TopNCollector) Collect(ctx context.Context, searcher search.Searcher, if err != nil { break } + if hc.earlyStopN > 0 && hc.total >= uint64(hc.earlyStopN) { + hc.earlyStopped = true + break + } } next, err = searcher.Next(searchContext) } @@ -692,6 +699,19 @@ 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 diff --git a/search/search.go b/search/search.go index b89f8acc0..8088c11e3 100644 --- a/search/search.go +++ b/search/search.go @@ -222,18 +222,14 @@ OUTER: // Reset allows an already allocated DocumentMatch to be reused func (dm *DocumentMatch) Reset() *DocumentMatch { - // remember the []byte used for the IndexInternalID + // Save backing arrays for reuse — these are restored below. indexInternalID := dm.IndexInternalID - // remember the []interface{} used for sort - sort := dm.Sort - // remember the []string used for decoded sort + sortBuf := dm.Sort decodedSort := dm.DecodedSort - // remember the FieldTermLocations backing array ftls := dm.FieldTermLocations for i := range ftls { // recycle the ArrayPositions of each location ftls[i].Location.ArrayPositions = ftls[i].Location.ArrayPositions[:0] } - // remember the score breakdown map scoreBreakdown := dm.ScoreBreakdown // clear out the score breakdown map if scoreBreakdown != nil { @@ -244,19 +240,22 @@ func (dm *DocumentMatch) Reset() *DocumentMatch { for i := range descendants { // recycle each IndexInternalID descendants[i] = descendants[i][:0] } - // idiom to copy over from empty DocumentMatch (0 allocations) - *dm = DocumentMatch{} - // reuse the []byte already allocated (and reset len to 0) + // Zero only the fields that are not restored below. + dm.Index = "" + dm.ID = "" + dm.Score = 0 + dm.Expl = nil + dm.Locations = nil + dm.Fragments = nil + dm.Fields = nil + dm.HitNumber = 0 + dm.IndexNames = nil + // Restore reusable allocations. dm.IndexInternalID = indexInternalID[:0] - // reuse the []interface{} already allocated (and reset len to 0) - dm.Sort = sort[:0] - // reuse the []string already allocated (and reset len to 0) + dm.Sort = sortBuf[:0] dm.DecodedSort = decodedSort[:0] - // reuse the FieldTermLocations already allocated (and reset len to 0) dm.FieldTermLocations = ftls[:0] - // reuse the Descendants already allocated (and reset len to 0) dm.Descendants = descendants[:0] - // reuse the score breakdown map already allocated (after clearing it) dm.ScoreBreakdown = scoreBreakdown return dm }