Skip to main content

hermes_core/query/
boolean.rs

1//! Boolean query with MUST, SHOULD, and MUST_NOT clauses
2
3use std::sync::Arc;
4
5use crate::segment::SegmentReader;
6use crate::structures::TERMINATED;
7use crate::{DocId, Score};
8
9use super::planner::{
10    build_combined_bitset, build_sparse_bmp_results, build_sparse_bmp_results_filtered,
11    build_sparse_maxscore_executor, chain_predicates, combine_sparse_results, compute_idf,
12    extract_all_sparse_infos, finish_text_maxscore, prepare_per_field_grouping,
13    prepare_text_maxscore,
14};
15use super::{CountFuture, EmptyScorer, GlobalStats, Query, Scorer, ScorerFuture};
16
17/// Boolean query with MUST, SHOULD, and MUST_NOT clauses
18///
19/// When all clauses are SHOULD term queries on the same field, automatically
20/// uses MaxScore optimization for efficient top-k retrieval.
21#[derive(Default, Clone)]
22pub struct BooleanQuery {
23    pub must: Vec<Arc<dyn Query>>,
24    pub should: Vec<Arc<dyn Query>>,
25    pub must_not: Vec<Arc<dyn Query>>,
26    /// Optional global statistics for cross-segment IDF
27    global_stats: Option<Arc<GlobalStats>>,
28}
29
30impl std::fmt::Debug for BooleanQuery {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        f.debug_struct("BooleanQuery")
33            .field("must_count", &self.must.len())
34            .field("should_count", &self.should.len())
35            .field("must_not_count", &self.must_not.len())
36            .field("has_global_stats", &self.global_stats.is_some())
37            .finish()
38    }
39}
40
41impl std::fmt::Display for BooleanQuery {
42    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
43        write!(f, "Boolean(")?;
44        let mut first = true;
45        for q in &self.must {
46            if !first {
47                write!(f, " ")?;
48            }
49            write!(f, "+{}", q)?;
50            first = false;
51        }
52        for q in &self.should {
53            if !first {
54                write!(f, " ")?;
55            }
56            write!(f, "{}", q)?;
57            first = false;
58        }
59        for q in &self.must_not {
60            if !first {
61                write!(f, " ")?;
62            }
63            write!(f, "-{}", q)?;
64            first = false;
65        }
66        write!(f, ")")
67    }
68}
69
70impl BooleanQuery {
71    pub fn new() -> Self {
72        Self::default()
73    }
74
75    pub fn must(mut self, query: impl Query + 'static) -> Self {
76        self.must.push(Arc::new(query));
77        self
78    }
79
80    pub fn should(mut self, query: impl Query + 'static) -> Self {
81        self.should.push(Arc::new(query));
82        self
83    }
84
85    pub fn must_not(mut self, query: impl Query + 'static) -> Self {
86        self.must_not.push(Arc::new(query));
87        self
88    }
89
90    /// Set global statistics for cross-segment IDF
91    pub fn with_global_stats(mut self, stats: Arc<GlobalStats>) -> Self {
92        self.global_stats = Some(stats);
93        self
94    }
95}
96
97/// Build a SHOULD-only scorer from a vec of optimized scorers.
98fn build_should_scorer<'a>(scorers: Vec<Box<dyn Scorer + 'a>>) -> Box<dyn Scorer + 'a> {
99    if scorers.is_empty() {
100        return Box::new(EmptyScorer);
101    }
102    if scorers.len() == 1 {
103        return scorers.into_iter().next().unwrap();
104    }
105    let mut scorer = BooleanScorer {
106        must: vec![],
107        should: scorers,
108        must_not: vec![],
109        current_doc: 0,
110    };
111    scorer.current_doc = scorer.find_next_match();
112    Box::new(scorer)
113}
114
115// ── Planner macro ────────────────────────────────────────────────────────
116//
117// Unified planner for both async and sync paths.  Parameterised on:
118//   $scorer_fn      – scorer_with_options | scorer_sync_with_options
119//   $get_postings_fn – get_postings | get_postings_sync
120//   $execute_fn     – execute | execute_sync
121//   $($aw)*         – .await  (present for async, absent for sync)
122//
123// Decision order:
124//   1. Single-clause unwrap
125//   2. Pure OR → text MaxScore | sparse MaxScore | per-field MaxScore
126//   3. Filter push-down → predicate-aware sparse MaxScore | PredicatedScorer
127//   4. Standard BooleanScorer fallback
128macro_rules! boolean_plan {
129    ($must:expr, $should:expr, $must_not:expr, $global_stats:expr,
130     $reader:expr, $limit:expr, $scorer_options:expr,
131     $scorer_fn:ident, $get_postings_fn:ident, $execute_fn:ident
132     $(, $aw:tt)*) => {{
133        let must: &[Arc<dyn Query>] = &$must;
134        let should_all: &[Arc<dyn Query>] = &$should;
135        let must_not: &[Arc<dyn Query>] = &$must_not;
136        let global_stats: Option<&Arc<GlobalStats>> = $global_stats;
137        let reader: &SegmentReader = $reader;
138        let limit: usize = $limit;
139        let scorer_options: super::ScorerOptions = $scorer_options;
140
141        // Cap SHOULD clauses to MAX_QUERY_TERMS, but only count queries that need
142        // posting-list cursors. Fast-field predicates (O(1) per doc) are exempt.
143        let should_capped: Vec<Arc<dyn Query>>;
144        let should: &[Arc<dyn Query>] = if should_all.len() > super::MAX_QUERY_TERMS {
145            let is_predicate: Vec<bool> = should_all
146                .iter()
147                .map(|q| q.is_filter() || q.as_doc_predicate(reader).is_some())
148                .collect();
149            let cursor_count = is_predicate.iter().filter(|&&p| !p).count();
150
151            if cursor_count > super::MAX_QUERY_TERMS {
152                let mut kept = Vec::with_capacity(should_all.len());
153                let mut cursor_kept = 0usize;
154                for (q, &is_pred) in should_all.iter().zip(is_predicate.iter()) {
155                    if is_pred {
156                        kept.push(q.clone());
157                    } else if cursor_kept < super::MAX_QUERY_TERMS {
158                        kept.push(q.clone());
159                        cursor_kept += 1;
160                    }
161                }
162                log::debug!(
163                    "BooleanQuery: capping cursor SHOULD from {} to {} ({} fast-field predicates exempt)",
164                    cursor_count,
165                    super::MAX_QUERY_TERMS,
166                    kept.len() - cursor_kept,
167                );
168                should_capped = kept;
169                &should_capped
170            } else {
171                log::debug!(
172                    "BooleanQuery: {} SHOULD clauses OK ({} need cursors, {} fast-field predicates)",
173                    should_all.len(),
174                    cursor_count,
175                    should_all.len() - cursor_count,
176                );
177                should_all
178            }
179        } else {
180            should_all
181        };
182
183        // ── 1. Single-clause optimisation ────────────────────────────────
184        if must_not.is_empty() {
185            if must.len() == 1 && should.is_empty() {
186                return must[0].$scorer_fn(reader, limit, scorer_options) $(.  $aw)* ;
187            }
188            if should.len() == 1 && must.is_empty() {
189                return should[0].$scorer_fn(reader, limit, scorer_options) $(. $aw)* ;
190            }
191        }
192
193        // ── 2. Pure OR → MaxScore optimisations ──────────────────────────
194        if must.is_empty() && must_not.is_empty() && should.len() >= 2 {
195            // 2a. Text MaxScore (single-field, all term queries)
196            if !scorer_options.collect_positions
197                && let Some((mut infos, text_field, avg_field_len, num_docs)) =
198                prepare_text_maxscore(should, reader, global_stats)
199            {
200                let mut posting_lists = Vec::with_capacity(infos.len());
201                for info in infos.drain(..) {
202                    if let Some(pl) = reader.$get_postings_fn(info.field, &info.term)
203                        $(. $aw)* ?
204                    {
205                        let idf = compute_idf(&pl, info.field, &info.term, num_docs, global_stats);
206                        posting_lists.push((pl, idf));
207                    }
208                }
209                // Seed from the cross-segment floor: this path scores final
210                // per-doc BM25 into a top-`limit` heap, so a floor carried from
211                // an already-searched segment prunes exactly (see
212                // SharedThreshold). The per-field path below stays at 0.0 —
213                // its per-field partial scores are not the final doc score.
214                let shared_threshold = std::cell::Cell::new(scorer_options.initial_threshold);
215                return finish_text_maxscore(
216                    posting_lists,
217                    avg_field_len,
218                    limit,
219                    &shared_threshold,
220                    reader.schema().index_label(),
221                    reader.schema().get_field_name(text_field).unwrap_or("?"),
222                );
223            }
224
225            // 2b. Sparse (single-field, all sparse term queries)
226            // Auto-detect: BMP executor if field has BMP index, else MaxScore
227            if let Some(infos) = extract_all_sparse_infos(should) {
228                if let Some((raw, info)) =
229                    build_sparse_bmp_results(&infos, reader, limit)
230                {
231                    return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
232                }
233                if let Some((executor, info)) =
234                    build_sparse_maxscore_executor(&infos, reader, limit, None)
235                {
236                    let raw = executor.$execute_fn() $(. $aw)* ?;
237                    return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
238                }
239            }
240
241            // 2c. Per-field text MaxScore (multi-field term grouping)
242            if !scorer_options.collect_positions
243                && let Some(grouping) =
244                    prepare_per_field_grouping(should, reader, limit, global_stats)
245            {
246                let mut scorers: Vec<Box<dyn Scorer + '_>> = Vec::new();
247                // Query-local cross-group threshold seeding (see finish_text_maxscore)
248                let shared_threshold = std::cell::Cell::new(0.0f32);
249                for (field, avg_field_len, infos) in &grouping.multi_term_groups {
250                    let mut posting_lists = Vec::with_capacity(infos.len());
251                    for info in infos {
252                        if let Some(pl) = reader.$get_postings_fn(info.field, &info.term)
253                            $(. $aw)* ?
254                        {
255                            let idf = compute_idf(
256                                &pl, *field, &info.term, grouping.num_docs, global_stats,
257                            );
258                            posting_lists.push((pl, idf));
259                        }
260                    }
261                    if !posting_lists.is_empty() {
262                        scorers.push(finish_text_maxscore(
263                            posting_lists,
264                            *avg_field_len,
265                            grouping.per_field_limit,
266                            &shared_threshold,
267                            reader.schema().index_label(),
268                            reader.schema().get_field_name(*field).unwrap_or("?"),
269                        )?);
270                    }
271                }
272                for &idx in &grouping.fallback_indices {
273                    scorers.push(should[idx].$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
274                }
275                return Ok(build_should_scorer(scorers));
276            }
277        }
278
279        // ── 3. Filter push-down (MUST + SHOULD) ─────────────────────────
280        if !scorer_options.collect_positions
281            && !should.is_empty()
282            && !must.is_empty()
283            && limit < usize::MAX / 4
284        {
285            // Pre-check: is SHOULD all-sparse? This determines whether we can
286            // use bitset fallback for MUST clauses that lack fast-field predicates.
287            // For sparse SHOULD, the predicate is pushed into BMP/MaxScore traversal
288            // so all qualifying docs are found. For text SHOULD, we must NOT convert
289            // MUST to a predicate (PredicatedScorer would drop MUST-only docs that
290            // don't match SHOULD), so those go to verifier → BooleanScorer.
291            let should_is_sparse = extract_all_sparse_infos(should).is_some();
292
293            // 3a. Compile MUST → predicates (O(1)) vs verifier scorers (seek)
294            //
295            // Priority: as_doc_predicate (fast-field O(1)) > as_doc_bitset
296            // (posting-list materialization, O(1) lookup, sparse-SHOULD only)
297            // > verifier scorer (seek).
298            let mut predicates: Vec<super::DocPredicate<'_>> = Vec::new();
299            let mut must_verifiers: Vec<Box<dyn super::Scorer + '_>> = Vec::new();
300            for q in must {
301                if let Some(pred) = q.as_doc_predicate(reader) {
302                    log::debug!("BooleanQuery planner 3a: MUST clause → predicate ({})", q);
303                    predicates.push(pred);
304                } else if should_is_sparse {
305                    if let Some(bitset) = q.as_doc_bitset(reader) {
306                        log::debug!("BooleanQuery planner 3a: MUST clause → bitset predicate ({})", q);
307                        predicates.push(Box::new(move |doc_id| bitset.contains(doc_id)));
308                    } else {
309                        log::debug!("BooleanQuery planner 3a: MUST clause → verifier scorer ({})", q);
310                        must_verifiers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
311                    }
312                } else {
313                    log::debug!("BooleanQuery planner 3a: MUST clause → verifier scorer ({})", q);
314                    must_verifiers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
315                }
316            }
317            // Compile MUST_NOT → negated predicates vs verifier scorers
318            let mut must_not_verifiers: Vec<Box<dyn super::Scorer + '_>> = Vec::new();
319            for q in must_not {
320                if let Some(pred) = q.as_doc_predicate(reader) {
321                    let negated: super::DocPredicate<'_> =
322                        Box::new(move |doc_id| !pred(doc_id));
323                    predicates.push(negated);
324                } else if should_is_sparse {
325                    if let Some(bitset) = q.as_doc_bitset(reader) {
326                        log::debug!("BooleanQuery planner 3a: MUST_NOT clause → bitset predicate ({})", q);
327                        predicates.push(Box::new(move |doc_id| !bitset.contains(doc_id)));
328                    } else {
329                        must_not_verifiers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
330                    }
331                } else {
332                    must_not_verifiers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
333                }
334            }
335
336            // 3b. Fast path: pure predicates + sparse SHOULD → BMP or MaxScore w/ predicate
337            if must_verifiers.is_empty()
338                && must_not_verifiers.is_empty()
339                && !predicates.is_empty()
340            {
341                if let Some(infos) = extract_all_sparse_infos(should) {
342                    // Try BMP with bitset first: build compact bitset from MUST/MUST_NOT
343                    // posting lists (O(M) for term queries) for fast per-slot lookup.
344                    let bitset_result = build_combined_bitset(must, must_not, reader);
345                    if let Some(ref bitset) = bitset_result {
346                        let bitset_pred = |doc_id: crate::DocId| bitset.contains(doc_id);
347                        if let Some((raw, info)) =
348                            build_sparse_bmp_results_filtered(&infos, reader, limit, &bitset_pred)
349                        {
350                            log::debug!(
351                                "BooleanQuery planner: bitset-aware sparse BMP, {} dims, {} matching docs",
352                                infos.len(),
353                                bitset.count()
354                            );
355                            return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
356                        }
357                    }
358
359                    // Fallback: closure predicate (for queries that don't support bitsets)
360                    let combined = chain_predicates(predicates);
361                    if let Some((raw, info)) =
362                        build_sparse_bmp_results_filtered(&infos, reader, limit, &*combined)
363                    {
364                        log::debug!(
365                            "BooleanQuery planner: predicate-aware sparse BMP, {} dims",
366                            infos.len()
367                        );
368                        return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
369                    }
370                    // Try MaxScore with predicate
371                    if let Some((executor, info)) =
372                        build_sparse_maxscore_executor(&infos, reader, limit, Some(combined))
373                    {
374                        log::debug!(
375                            "BooleanQuery planner: predicate-aware sparse MaxScore, {} dims",
376                            infos.len()
377                        );
378                        let raw = executor.$execute_fn() $(. $aw)* ?;
379                        return Ok(combine_sparse_results(raw, info.combiner, info.field, limit));
380                    }
381                    // predicates consumed — cannot fall through; rebuild them
382                    // (this path only triggers if neither sparse index exists)
383                    // should_is_sparse is true here (we're inside extract_all_sparse_infos)
384                    predicates = Vec::new();
385                    for q in must {
386                        if let Some(pred) = q.as_doc_predicate(reader) {
387                            predicates.push(pred);
388                        } else if let Some(bitset) = q.as_doc_bitset(reader) {
389                            predicates.push(Box::new(move |doc_id| bitset.contains(doc_id)));
390                        }
391                    }
392                    for q in must_not {
393                        if let Some(pred) = q.as_doc_predicate(reader) {
394                            let negated: super::DocPredicate<'_> =
395                                Box::new(move |doc_id| !pred(doc_id));
396                            predicates.push(negated);
397                        } else if let Some(bitset) = q.as_doc_bitset(reader) {
398                            predicates.push(Box::new(move |doc_id| !bitset.contains(doc_id)));
399                        }
400                    }
401                }
402            }
403
404            // 3c. PredicatedScorer fallback (over-fetch 4x when any filter is present)
405            let has_filters = !predicates.is_empty()
406                || !must_verifiers.is_empty()
407                || !must_not_verifiers.is_empty();
408            let should_limit = if has_filters { limit * 4 } else { limit };
409            let should_scorer = if should.len() == 1 {
410                should[0].$scorer_fn(reader, should_limit, scorer_options) $(. $aw)* ?
411            } else {
412                let sub = BooleanQuery {
413                    must: Vec::new(),
414                    should: should.to_vec(),
415                    must_not: Vec::new(),
416                    global_stats: global_stats.cloned(),
417                };
418                sub.$scorer_fn(reader, should_limit, scorer_options) $(. $aw)* ?
419            };
420
421            let use_predicated =
422                must_verifiers.is_empty() || should_scorer.size_hint() >= limit as u32;
423
424            if use_predicated {
425                log::debug!(
426                    "BooleanQuery planner: PredicatedScorer {} preds + {} must_v + {} must_not_v, \
427                     SHOULD size_hint={}, over_fetch={}",
428                    predicates.len(), must_verifiers.len(), must_not_verifiers.len(),
429                    should_scorer.size_hint(), should_limit
430                );
431                return Ok(Box::new(super::PredicatedScorer::new(
432                    should_scorer, predicates, must_verifiers, must_not_verifiers,
433                )));
434            }
435
436            // size_hint < limit with verifiers → BooleanScorer
437            log::debug!(
438                "BooleanQuery planner: BooleanScorer fallback, size_hint={} < limit={}, \
439                 {} must_v + {} must_not_v",
440                should_scorer.size_hint(), limit,
441                must_verifiers.len(), must_not_verifiers.len()
442            );
443            let mut scorer = BooleanScorer {
444                must: must_verifiers,
445                should: vec![should_scorer],
446                must_not: must_not_verifiers,
447                current_doc: 0,
448            };
449            scorer.current_doc = scorer.find_next_match();
450            return Ok(Box::new(scorer));
451        }
452
453        // ── 4. Standard BooleanScorer fallback ───────────────────────────
454        let mut must_scorers = Vec::with_capacity(must.len());
455        for q in must {
456            must_scorers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
457        }
458        let mut should_scorers = Vec::with_capacity(should.len());
459        for q in should {
460            should_scorers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
461        }
462        let mut must_not_scorers = Vec::with_capacity(must_not.len());
463        for q in must_not {
464            must_not_scorers.push(q.$scorer_fn(reader, limit, scorer_options) $(. $aw)* ?);
465        }
466        let mut scorer = BooleanScorer {
467            must: must_scorers,
468            should: should_scorers,
469            must_not: must_not_scorers,
470            current_doc: 0,
471        };
472        scorer.current_doc = scorer.find_next_match();
473        Ok(Box::new(scorer) as Box<dyn Scorer + '_>)
474    }};
475}
476
477impl Query for BooleanQuery {
478    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
479        self.scorer_with_options(reader, limit, super::ScorerOptions::with_positions())
480    }
481
482    fn scorer_with_options<'a>(
483        &self,
484        reader: &'a SegmentReader,
485        limit: usize,
486        options: super::ScorerOptions,
487    ) -> ScorerFuture<'a> {
488        let must = self.must.clone();
489        let should = self.should.clone();
490        let must_not = self.must_not.clone();
491        let global_stats = self.global_stats.clone();
492        Box::pin(async move {
493            boolean_plan!(
494                must,
495                should,
496                must_not,
497                global_stats.as_ref(),
498                reader,
499                limit,
500                options,
501                scorer_with_options,
502                get_postings,
503                execute,
504                await
505            )
506        })
507    }
508
509    #[cfg(feature = "sync")]
510    fn scorer_sync<'a>(
511        &self,
512        reader: &'a SegmentReader,
513        limit: usize,
514    ) -> crate::Result<Box<dyn Scorer + 'a>> {
515        self.scorer_sync_with_options(reader, limit, super::ScorerOptions::with_positions())
516    }
517
518    #[cfg(feature = "sync")]
519    fn scorer_sync_with_options<'a>(
520        &self,
521        reader: &'a SegmentReader,
522        limit: usize,
523        options: super::ScorerOptions,
524    ) -> crate::Result<Box<dyn Scorer + 'a>> {
525        boolean_plan!(
526            self.must,
527            self.should,
528            self.must_not,
529            self.global_stats.as_ref(),
530            reader,
531            limit,
532            options,
533            scorer_sync_with_options,
534            get_postings_sync,
535            execute_sync
536        )
537    }
538
539    fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
540        if self.must.is_empty() && self.should.is_empty() {
541            return None;
542        }
543
544        let num_docs = reader.num_docs();
545
546        // MUST clauses: intersect bitsets (AND)
547        let mut result: Option<super::DocBitset> = None;
548        for q in &self.must {
549            let bs = q.as_doc_bitset(reader)?;
550            match result {
551                None => result = Some(bs),
552                Some(ref mut acc) => acc.intersect_with(&bs),
553            }
554        }
555
556        // SHOULD clauses: union bitsets (OR), then intersect with MUST result
557        if !self.should.is_empty() {
558            let mut should_union = super::DocBitset::new(num_docs);
559            for q in &self.should {
560                let bs = q.as_doc_bitset(reader)?;
561                should_union.union_with(&bs);
562            }
563            match result {
564                None => result = Some(should_union),
565                Some(ref mut acc) => {
566                    // When MUST clauses exist, SHOULD is optional (doesn't filter).
567                    // When no MUST clauses, at least one SHOULD must match.
568                    if self.must.is_empty() {
569                        *acc = should_union;
570                    }
571                }
572            }
573        }
574
575        // MUST_NOT clauses: subtract bitsets (ANDNOT)
576        if let Some(ref mut acc) = result {
577            for q in &self.must_not {
578                {
579                    let bs = q.as_doc_bitset(reader)?;
580                    acc.subtract(&bs);
581                }
582            }
583        }
584
585        result
586    }
587
588    fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
589        // Need at least some clauses
590        if self.must.is_empty() && self.should.is_empty() {
591            return None;
592        }
593
594        // Try converting all clauses to predicates; bail if any child can't
595        let must_preds: Vec<_> = self
596            .must
597            .iter()
598            .map(|q| q.as_doc_predicate(reader))
599            .collect::<Option<Vec<_>>>()?;
600        let should_preds: Vec<_> = self
601            .should
602            .iter()
603            .map(|q| q.as_doc_predicate(reader))
604            .collect::<Option<Vec<_>>>()?;
605        let must_not_preds: Vec<_> = self
606            .must_not
607            .iter()
608            .map(|q| q.as_doc_predicate(reader))
609            .collect::<Option<Vec<_>>>()?;
610
611        let has_must = !must_preds.is_empty();
612
613        Some(Box::new(move |doc_id| {
614            // All MUST predicates must pass
615            if !must_preds.iter().all(|p| p(doc_id)) {
616                return false;
617            }
618            // When there are no MUST clauses, at least one SHOULD must pass
619            if !has_must && !should_preds.is_empty() && !should_preds.iter().any(|p| p(doc_id)) {
620                return false;
621            }
622            // No MUST_NOT predicate should pass
623            must_not_preds.iter().all(|p| !p(doc_id))
624        }))
625    }
626
627    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
628        let must = self.must.clone();
629        let should = self.should.clone();
630
631        Box::pin(async move {
632            if !must.is_empty() {
633                let mut estimates = Vec::with_capacity(must.len());
634                for q in &must {
635                    estimates.push(q.count_estimate(reader).await?);
636                }
637                estimates
638                    .into_iter()
639                    .min()
640                    .ok_or_else(|| crate::Error::Corruption("Empty must clause".to_string()))
641            } else if !should.is_empty() {
642                let mut sum = 0u32;
643                for q in &should {
644                    sum = sum.saturating_add(q.count_estimate(reader).await?);
645                }
646                Ok(sum)
647            } else {
648                Ok(0)
649            }
650        })
651    }
652}
653
654struct BooleanScorer<'a> {
655    must: Vec<Box<dyn Scorer + 'a>>,
656    should: Vec<Box<dyn Scorer + 'a>>,
657    must_not: Vec<Box<dyn Scorer + 'a>>,
658    current_doc: DocId,
659}
660
661impl BooleanScorer<'_> {
662    fn find_next_match(&mut self) -> DocId {
663        if self.must.is_empty() && self.should.is_empty() {
664            return TERMINATED;
665        }
666
667        loop {
668            let candidate = if !self.must.is_empty() {
669                let mut max_doc = self
670                    .must
671                    .iter()
672                    .map(|s| s.doc())
673                    .max()
674                    .unwrap_or(TERMINATED);
675
676                if max_doc == TERMINATED {
677                    return TERMINATED;
678                }
679
680                loop {
681                    let mut all_match = true;
682                    for scorer in &mut self.must {
683                        let doc = scorer.seek(max_doc);
684                        if doc == TERMINATED {
685                            return TERMINATED;
686                        }
687                        if doc > max_doc {
688                            max_doc = doc;
689                            all_match = false;
690                            break;
691                        }
692                    }
693                    if all_match {
694                        break;
695                    }
696                }
697                max_doc
698            } else {
699                self.should
700                    .iter()
701                    .map(|s| s.doc())
702                    .filter(|&d| d != TERMINATED)
703                    .min()
704                    .unwrap_or(TERMINATED)
705            };
706
707            if candidate == TERMINATED {
708                return TERMINATED;
709            }
710
711            let excluded = self.must_not.iter_mut().any(|scorer| {
712                let doc = scorer.seek(candidate);
713                doc == candidate
714            });
715
716            if !excluded {
717                // Seek SHOULD scorers to candidate so score() can see their contributions
718                for scorer in &mut self.should {
719                    scorer.seek(candidate);
720                }
721                self.current_doc = candidate;
722                return candidate;
723            }
724
725            // Advance past excluded candidate
726            if !self.must.is_empty() {
727                for scorer in &mut self.must {
728                    scorer.advance();
729                }
730            } else {
731                // For SHOULD-only: seek all scorers past the excluded candidate
732                for scorer in &mut self.should {
733                    if scorer.doc() <= candidate && scorer.doc() != TERMINATED {
734                        scorer.seek(candidate + 1);
735                    }
736                }
737            }
738        }
739    }
740}
741
742impl super::docset::DocSet for BooleanScorer<'_> {
743    fn doc(&self) -> DocId {
744        self.current_doc
745    }
746
747    fn advance(&mut self) -> DocId {
748        if !self.must.is_empty() {
749            for scorer in &mut self.must {
750                scorer.advance();
751            }
752        } else {
753            for scorer in &mut self.should {
754                if scorer.doc() == self.current_doc {
755                    scorer.advance();
756                }
757            }
758        }
759
760        self.current_doc = self.find_next_match();
761        self.current_doc
762    }
763
764    fn seek(&mut self, target: DocId) -> DocId {
765        for scorer in &mut self.must {
766            scorer.seek(target);
767        }
768
769        for scorer in &mut self.should {
770            scorer.seek(target);
771        }
772
773        self.current_doc = self.find_next_match();
774        self.current_doc
775    }
776
777    fn size_hint(&self) -> u32 {
778        if !self.must.is_empty() {
779            self.must.iter().map(|s| s.size_hint()).min().unwrap_or(0)
780        } else {
781            self.should.iter().map(|s| s.size_hint()).sum()
782        }
783    }
784}
785
786impl Scorer for BooleanScorer<'_> {
787    fn score(&self) -> Score {
788        let mut total = 0.0;
789
790        for scorer in &self.must {
791            if scorer.doc() == self.current_doc {
792                total += scorer.score();
793            }
794        }
795
796        for scorer in &self.should {
797            if scorer.doc() == self.current_doc {
798                total += scorer.score();
799            }
800        }
801
802        total
803    }
804
805    fn matched_positions(&self) -> Option<super::MatchedPositions> {
806        let mut all_positions: super::MatchedPositions = Vec::new();
807
808        for scorer in &self.must {
809            if scorer.doc() == self.current_doc
810                && let Some(positions) = scorer.matched_positions()
811            {
812                all_positions.extend(positions);
813            }
814        }
815
816        for scorer in &self.should {
817            if scorer.doc() == self.current_doc
818                && let Some(positions) = scorer.matched_positions()
819            {
820                all_positions.extend(positions);
821            }
822        }
823
824        if all_positions.is_empty() {
825            None
826        } else {
827            Some(all_positions)
828        }
829    }
830}
831
832#[cfg(test)]
833mod tests {
834    use super::*;
835    use crate::dsl::Field;
836    use crate::query::{QueryDecomposition, TermQuery};
837
838    #[test]
839    fn test_maxscore_eligible_pure_or_same_field() {
840        // Pure OR query with multiple terms in same field should be MaxScore-eligible
841        let query = BooleanQuery::new()
842            .should(TermQuery::text(Field(0), "hello"))
843            .should(TermQuery::text(Field(0), "world"))
844            .should(TermQuery::text(Field(0), "foo"));
845
846        // All clauses should return term info
847        assert!(
848            query
849                .should
850                .iter()
851                .all(|q| matches!(q.decompose(), QueryDecomposition::TextTerm(_)))
852        );
853
854        // All should be same field
855        let infos: Vec<_> = query
856            .should
857            .iter()
858            .filter_map(|q| match q.decompose() {
859                QueryDecomposition::TextTerm(info) => Some(info),
860                _ => None,
861            })
862            .collect();
863        assert_eq!(infos.len(), 3);
864        assert!(infos.iter().all(|i| i.field == Field(0)));
865    }
866
867    #[test]
868    fn test_maxscore_not_eligible_different_fields() {
869        // OR query with terms in different fields should NOT use MaxScore
870        let query = BooleanQuery::new()
871            .should(TermQuery::text(Field(0), "hello"))
872            .should(TermQuery::text(Field(1), "world")); // Different field!
873
874        let infos: Vec<_> = query
875            .should
876            .iter()
877            .filter_map(|q| match q.decompose() {
878                QueryDecomposition::TextTerm(info) => Some(info),
879                _ => None,
880            })
881            .collect();
882        assert_eq!(infos.len(), 2);
883        // Fields are different, MaxScore should not be used
884        assert!(infos[0].field != infos[1].field);
885    }
886
887    #[test]
888    fn test_maxscore_not_eligible_with_must() {
889        // Query with MUST clause should NOT use MaxScore optimization
890        let query = BooleanQuery::new()
891            .must(TermQuery::text(Field(0), "required"))
892            .should(TermQuery::text(Field(0), "hello"))
893            .should(TermQuery::text(Field(0), "world"));
894
895        // Has MUST clause, so MaxScore optimization should not kick in
896        assert!(!query.must.is_empty());
897    }
898
899    #[test]
900    fn test_maxscore_not_eligible_with_must_not() {
901        // Query with MUST_NOT clause should NOT use MaxScore optimization
902        let query = BooleanQuery::new()
903            .should(TermQuery::text(Field(0), "hello"))
904            .should(TermQuery::text(Field(0), "world"))
905            .must_not(TermQuery::text(Field(0), "excluded"));
906
907        // Has MUST_NOT clause, so MaxScore optimization should not kick in
908        assert!(!query.must_not.is_empty());
909    }
910
911    #[test]
912    fn test_maxscore_not_eligible_single_term() {
913        // Single SHOULD clause should NOT use MaxScore (no benefit)
914        let query = BooleanQuery::new().should(TermQuery::text(Field(0), "hello"));
915
916        // Only one term, MaxScore not beneficial
917        assert_eq!(query.should.len(), 1);
918    }
919
920    #[test]
921    fn test_term_query_info_extraction() {
922        let term_query = TermQuery::text(Field(42), "test");
923        match term_query.decompose() {
924            QueryDecomposition::TextTerm(info) => {
925                assert_eq!(info.field, Field(42));
926                assert_eq!(info.term, b"test");
927            }
928            _ => panic!("Expected TextTerm decomposition"),
929        }
930    }
931
932    #[test]
933    fn test_boolean_query_no_term_info() {
934        // BooleanQuery itself should not return term info
935        let query = BooleanQuery::new().should(TermQuery::text(Field(0), "hello"));
936
937        assert!(matches!(query.decompose(), QueryDecomposition::Opaque));
938    }
939}