Skip to content

Commit 0d1f64e

Browse files
CascadingRadiumCascadingRadium
committed
revert
1 parent 3f90d48 commit 0d1f64e

File tree

4 files changed

+123
-71
lines changed

4 files changed

+123
-71
lines changed

‎index_impl.go‎

Lines changed: 61 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -640,6 +640,53 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr
640640
}
641641
}
642642

643+
// ------------------------------------------------------------------------------------------
644+
// set up additional contexts for any search operation that will proceed from
645+
// here, such as presearch, knn collector, topn collector etc.
646+
647+
// Scoring model callback to be used to get scoring model
648+
scoringModelCallback := func() string {
649+
if isBM25Enabled(i.m) {
650+
return index.BM25Scoring
651+
}
652+
return index.DefaultScoringModel
653+
}
654+
ctx = context.WithValue(ctx, search.GetScoringModelCallbackKey,
655+
search.GetScoringModelCallbackFn(scoringModelCallback))
656+
657+
// This callback and variable handles the tracking of bytes read
658+
// 1. as part of creation of tfr and its Next() calls which is
659+
// accounted by invoking this callback when the TFR is closed.
660+
// 2. the docvalues portion (accounted in collector) and the retrieval
661+
// of stored fields bytes (by LoadAndHighlightFields)
662+
var totalSearchCost uint64
663+
sendBytesRead := func(bytesRead uint64) {
664+
totalSearchCost += bytesRead
665+
}
666+
667+
ctx = context.WithValue(ctx, search.SearchIOStatsCallbackKey, search.SearchIOStatsCallbackFunc(sendBytesRead))
668+
669+
// Geo buffer pool callback to be used for getting geo buffer pool
670+
var bufPool *s2.GeoBufferPool
671+
getBufferPool := func() *s2.GeoBufferPool {
672+
if bufPool == nil {
673+
bufPool = s2.NewGeoBufferPool(search.MaxGeoBufPoolSize, search.MinGeoBufPoolSize)
674+
}
675+
676+
return bufPool
677+
}
678+
679+
ctx = context.WithValue(ctx, search.GeoBufferPoolCallbackKey, search.GeoBufferPoolCallbackFunc(getBufferPool))
680+
681+
// check if the index mapping has any nested fields, which should force
682+
// all collectors and searchers to be run in nested mode
683+
if nm, ok := i.m.(mapping.NestedMapping); ok {
684+
if nm.CountNested() > 0 {
685+
ctx = context.WithValue(ctx, search.NestedSearchKey, true)
686+
}
687+
}
688+
// ------------------------------------------------------------------------------------------
689+
643690
if _, ok := ctx.Value(search.PreSearchKey).(bool); ok {
644691
preSearchResult, err := i.preSearch(ctx, req, indexReader)
645692
if err != nil {
@@ -666,7 +713,7 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr
666713
req.SearchBefore = nil
667714
}
668715

669-
coll, err := i.buildTopNCollector(req, indexReader)
716+
coll, err := i.buildTopNCollector(ctx, req, indexReader)
670717
if err != nil {
671718
return nil, err
672719
}
@@ -749,44 +796,12 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr
749796
ctx = context.WithValue(ctx, search.FieldTermSynonymMapKey, fts)
750797
}
751798

752-
scoringModelCallback := func() string {
753-
if isBM25Enabled(i.m) {
754-
return index.BM25Scoring
755-
}
756-
return index.DefaultScoringModel
757-
}
758-
ctx = context.WithValue(ctx, search.GetScoringModelCallbackKey,
759-
search.GetScoringModelCallbackFn(scoringModelCallback))
760-
761799
// set the bm25Stats (stats important for consistent scoring) in
762800
// the context object
763801
if bm25Stats != nil {
764802
ctx = context.WithValue(ctx, search.BM25StatsKey, bm25Stats)
765803
}
766804

767-
// This callback and variable handles the tracking of bytes read
768-
// 1. as part of creation of tfr and its Next() calls which is
769-
// accounted by invoking this callback when the TFR is closed.
770-
// 2. the docvalues portion (accounted in collector) and the retrieval
771-
// of stored fields bytes (by LoadAndHighlightFields)
772-
var totalSearchCost uint64
773-
sendBytesRead := func(bytesRead uint64) {
774-
totalSearchCost += bytesRead
775-
}
776-
777-
ctx = context.WithValue(ctx, search.SearchIOStatsCallbackKey, search.SearchIOStatsCallbackFunc(sendBytesRead))
778-
779-
var bufPool *s2.GeoBufferPool
780-
getBufferPool := func() *s2.GeoBufferPool {
781-
if bufPool == nil {
782-
bufPool = s2.NewGeoBufferPool(search.MaxGeoBufPoolSize, search.MinGeoBufPoolSize)
783-
}
784-
785-
return bufPool
786-
}
787-
788-
ctx = context.WithValue(ctx, search.GeoBufferPoolCallbackKey, search.GeoBufferPoolCallbackFunc(getBufferPool))
789-
790805
searcher, err := req.Query.Searcher(ctx, indexReader, i.m, search.SearcherOptions{
791806
Explain: req.Explain,
792807
IncludeTermVectors: req.IncludeLocations || req.Highlight != nil,
@@ -1435,7 +1450,7 @@ func (i *indexImpl) FireIndexEvent() {
14351450
}
14361451
}
14371452

1438-
func (i *indexImpl) buildTopNCollector(req *SearchRequest, reader index.IndexReader) (*collector.TopNCollector, error) {
1453+
func (i *indexImpl) buildTopNCollector(ctx context.Context, req *SearchRequest, reader index.IndexReader) (*collector.TopNCollector, error) {
14391454
newCollector := func() *collector.TopNCollector {
14401455
if req.SearchAfter != nil {
14411456
return collector.NewTopNCollectorAfter(req.Size, req.Sort, req.SearchAfter)
@@ -1450,18 +1465,18 @@ func (i *indexImpl) buildTopNCollector(req *SearchRequest, reader index.IndexRea
14501465
return collector.NewNestedTopNCollector(req.Size, req.From, req.Sort, nr)
14511466
}
14521467

1453-
if nm, ok := i.m.(mapping.NestedMapping); ok {
1454-
if nm.CountNested() > 0 {
1455-
if nr, ok := reader.(index.NestedReader); ok {
1456-
var fs search.FieldSet
1457-
var err error
1458-
fs, err = query.ExtractFields(req.Query, i.m, fs)
1459-
if err != nil {
1460-
return nil, err
1461-
}
1462-
if nm.IntersectsPrefix(fs) {
1463-
return newNestedCollector(nr), nil
1464-
}
1468+
// check if we are in nested mode
1469+
if nestedMode, ok := ctx.Value(search.NestedSearchKey).(bool); ok && nestedMode {
1470+
// get the nested reader from the index reader
1471+
if nr, ok := reader.(index.NestedReader); ok {
1472+
var fs search.FieldSet
1473+
var err error
1474+
fs, err = query.ExtractFields(req.Query, i.m, fs)
1475+
if err != nil {
1476+
return nil, err
1477+
}
1478+
if nm.IntersectsPrefix(fs) {
1479+
return newNestedCollector(nr), nil
14651480
}
14661481
}
14671482
}

‎search/searcher/search_match_all.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ type MatchAllSearcher struct {
3636
reader index.DocIDReader
3737
scorer *scorer.ConstantScorer
3838
count uint64
39+
nested bool
3940
}
4041

4142
func NewMatchAllSearcher(ctx context.Context, indexReader index.IndexReader, boost float64, options search.SearcherOptions) (*MatchAllSearcher, error) {
@@ -50,11 +51,18 @@ func NewMatchAllSearcher(ctx context.Context, indexReader index.IndexReader, boo
5051
}
5152
scorer := scorer.NewConstantScorer(1.0, boost, options)
5253

54+
// check if we are in nested mode
55+
nested := false
56+
if nestedMode, ok := ctx.Value(search.NestedSearchKey).(bool); ok && nestedMode {
57+
nested = true
58+
}
59+
5360
return &MatchAllSearcher{
5461
indexReader: indexReader,
5562
reader: reader,
5663
scorer: scorer,
5764
count: count,
65+
nested: nested,
5866
}, nil
5967
}
6068

@@ -76,6 +84,22 @@ func (s *MatchAllSearcher) SetQueryNorm(qnorm float64) {
7684
s.scorer.SetQueryNorm(qnorm)
7785
}
7886

87+
func (s *MatchAllSearcher) isNested(id index.IndexInternalID) bool {
88+
// if not running in nested mode, always return false
89+
if !s.nested {
90+
return false
91+
}
92+
// check if this doc has ancestors, if so it is nested
93+
if nr, ok := s.reader.(index.NestedReader); ok {
94+
anc, err := nr.Ancestors(id)
95+
if err != nil {
96+
return false
97+
}
98+
return len(anc) > 1
99+
}
100+
return false
101+
}
102+
79103
func (s *MatchAllSearcher) Next(ctx *search.SearchContext) (*search.DocumentMatch, error) {
80104
id, err := s.reader.Next()
81105
if err != nil {
@@ -86,6 +110,11 @@ func (s *MatchAllSearcher) Next(ctx *search.SearchContext) (*search.DocumentMatc
86110
return nil, nil
87111
}
88112

113+
if s.isNested(id) {
114+
// if nested then skip and get next
115+
return s.Next(ctx)
116+
}
117+
89118
// score match
90119
docMatch := s.scorer.Score(ctx, id)
91120
// return doc match
@@ -103,6 +132,11 @@ func (s *MatchAllSearcher) Advance(ctx *search.SearchContext, ID index.IndexInte
103132
return nil, nil
104133
}
105134

135+
if s.isNested(id) {
136+
// if nested then return next
137+
return s.Next(ctx)
138+
}
139+
106140
// score match
107141
docMatch := s.scorer.Score(ctx, id)
108142

‎search/util.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,10 @@ const (
156156
// ScoreFusionKey is used to communicate whether KNN hits need to be preserved for
157157
// hybrid search algorithms (like RRF)
158158
ScoreFusionKey ContextKey = "_fusion_rescoring_key"
159+
160+
// NestedSearchKey is used to communicate whether the search is performed
161+
// in an index with nested documents
162+
NestedSearchKey ContextKey = "_nested_search_key"
159163
)
160164

161165
func RecordSearchCost(ctx context.Context,

‎search_knn.go‎

Lines changed: 24 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import (
2424
"sort"
2525

2626
"github.com/blevesearch/bleve/v2/document"
27-
"github.com/blevesearch/bleve/v2/mapping"
2827
"github.com/blevesearch/bleve/v2/search"
2928
"github.com/blevesearch/bleve/v2/search/collector"
3029
"github.com/blevesearch/bleve/v2/search/query"
@@ -669,39 +668,39 @@ func (r *rescorer) restoreKnnRequest() {
669668
}
670669
}
671670

672-
func (i *indexImpl) buildKNNCollector(KNNQuery query.Query, reader index.IndexReader, kArray []int64, somOfK int64) (*collector.KNNCollector, error) {
673-
if nm, ok := i.m.(mapping.NestedMapping); ok {
671+
func (i *indexImpl) buildKNNCollector(ctx context.Context, KNNQuery query.Query, reader index.IndexReader, kArray []int64, somOfK int64) (*collector.KNNCollector, error) {
672+
// check if we are in nested mode
673+
if nestedMode, ok := ctx.Value(search.NestedSearchKey).(bool); ok && nestedMode {
674+
// get the nested reader from the index reader
674675
if nr, ok := reader.(index.NestedReader); ok {
675-
if nm.CountNested() > 0 {
676-
var fs search.FieldSet
677-
var err error
678-
fs, err = query.ExtractFields(KNNQuery, i.m, fs)
679-
if err != nil {
680-
return nil, err
681-
}
682-
if nm.IntersectsPrefix(fs) {
683-
return collector.NewNestedKNNCollector(nr, kArray, somOfK), nil
684-
}
676+
var fs search.FieldSet
677+
var err error
678+
fs, err = query.ExtractFields(KNNQuery, i.m, fs)
679+
if err != nil {
680+
return nil, err
681+
}
682+
if nm.IntersectsPrefix(fs) {
683+
return collector.NewNestedKNNCollector(nr, kArray, somOfK), nil
685684
}
686685
}
687686
}
688687

689688
return collector.NewKNNCollector(kArray, somOfK), nil
690689
}
691690

692-
func (i *indexImpl) buildEligibleCollector(filterQuery query.Query, reader index.IndexReader, size int) (*collector.EligibleCollector, error) {
693-
if nm, ok := i.m.(mapping.NestedMapping); ok {
691+
func (i *indexImpl) buildEligibleCollector(ctx context.Context, filterQuery query.Query, reader index.IndexReader, size int) (*collector.EligibleCollector, error) {
692+
// check if we are in nested mode
693+
if nestedMode, ok := ctx.Value(search.NestedSearchKey).(bool); ok && nestedMode {
694+
// get the nested reader from the index reader
694695
if nr, ok := reader.(index.NestedReader); ok {
695-
if nm.CountNested() > 0 {
696-
var fs search.FieldSet
697-
var err error
698-
fs, err = query.ExtractFields(filterQuery, i.m, fs)
699-
if err != nil {
700-
return nil, err
701-
}
702-
if nm.IntersectsPrefix(fs) {
703-
return collector.NewNestedEligibleCollector(nr, size), nil
704-
}
696+
var fs search.FieldSet
697+
var err error
698+
fs, err = query.ExtractFields(filterQuery, i.m, fs)
699+
if err != nil {
700+
return nil, err
701+
}
702+
if nm.IntersectsPrefix(fs) {
703+
return collector.NewNestedEligibleCollector(nr, size), nil
705704
}
706705
}
707706
}

0 commit comments

Comments
 (0)