Skip to main content

hermes_core/query/
phrase.rs

1//! Phrase query - matches documents containing terms in consecutive positions
2
3use std::sync::Arc;
4
5use crate::dsl::Field;
6use crate::segment::SegmentReader;
7use crate::structures::{BlockPostingIterator, BlockPostingList, TERMINATED, TermPositions};
8use crate::{DocId, Score};
9
10use super::docset::DocSet;
11use super::{CountFuture, EmptyScorer, GlobalStats, Query, Scorer, ScorerFuture};
12
13/// Phrase query - matches documents containing terms in consecutive positions
14///
15/// Example: "quick brown fox" matches only if all three terms appear
16/// consecutively in the document.
17#[derive(Clone)]
18pub struct PhraseQuery {
19    pub field: Field,
20    /// Terms in the phrase, in order
21    pub terms: Vec<Vec<u8>>,
22    /// Token offset of each term inside the phrase, ascending, one per term.
23    /// `offsets[i + 1] - offsets[i]` is the required distance between two
24    /// consecutive terms: 1 for adjacent words, more when index-time stop
25    /// words were dropped between them (`quantum@0 art@3`). [`PhraseQuery::new`]
26    /// makes every term adjacent.
27    pub offsets: Vec<u32>,
28    /// Optional slop (max distance between terms, 0 = exact phrase)
29    pub slop: u32,
30    /// Optional global statistics for cross-segment IDF
31    global_stats: Option<Arc<GlobalStats>>,
32}
33
34impl std::fmt::Display for PhraseQuery {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        let terms: Vec<String> = self
37            .terms
38            .iter()
39            .zip(&self.offsets)
40            .map(|(term, offset)| {
41                if self.is_adjacent() {
42                    String::from_utf8_lossy(term).into_owned()
43                } else {
44                    format!("{}@{offset}", String::from_utf8_lossy(term))
45                }
46            })
47            .collect();
48        write!(f, "Phrase({}:\"{}\"", self.field.0, terms.join(" "))?;
49        if self.slop > 0 {
50            write!(f, "~{}", self.slop)?;
51        }
52        write!(f, ")")
53    }
54}
55
56impl std::fmt::Debug for PhraseQuery {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        let terms: Vec<_> = self
59            .terms
60            .iter()
61            .map(|t| String::from_utf8_lossy(t).to_string())
62            .collect();
63        f.debug_struct("PhraseQuery")
64            .field("field", &self.field)
65            .field("terms", &terms)
66            .field("offsets", &self.offsets)
67            .field("slop", &self.slop)
68            .finish()
69    }
70}
71
72impl PhraseQuery {
73    /// Create a new exact phrase query of adjacent terms.
74    pub fn new(field: Field, terms: Vec<Vec<u8>>) -> Self {
75        let offsets = (0..terms.len() as u32).collect();
76        Self {
77            field,
78            terms,
79            offsets,
80            slop: 0,
81            global_stats: None,
82        }
83    }
84
85    /// Create a phrase whose terms carry their token offsets, as produced by
86    /// a tokenizer that drops stop words without renumbering (`(0, quantum)`,
87    /// `(3, art)` for "quantum of the art"). Offsets must be ascending.
88    pub fn with_offsets(field: Field, terms: Vec<(u32, Vec<u8>)>) -> Self {
89        debug_assert!(
90            terms.windows(2).all(|pair| pair[0].0 < pair[1].0),
91            "phrase offsets must be strictly ascending"
92        );
93        let (offsets, terms): (Vec<u32>, Vec<Vec<u8>>) = terms.into_iter().unzip();
94        Self {
95            field,
96            terms,
97            offsets,
98            slop: 0,
99            global_stats: None,
100        }
101    }
102
103    /// Create from text using the simple tokenizer (whitespace split,
104    /// punctuation stripped, lowercased). Fields with a stemming tokenizer
105    /// should tokenize the phrase themselves and call
106    /// [`PhraseQuery::with_offsets`] so the query terms match the indexed
107    /// stems and keep the gaps of dropped stop words.
108    pub fn text(field: Field, phrase: &str) -> Self {
109        use crate::tokenizer::Tokenizer;
110        let terms: Vec<(u32, Vec<u8>)> = crate::tokenizer::SimpleTokenizer
111            .tokenize(phrase)
112            .into_iter()
113            .map(|token| (token.position, token.text.into_bytes()))
114            .collect();
115        Self::with_offsets(field, terms)
116    }
117
118    /// Whether every term must directly follow the previous one.
119    fn is_adjacent(&self) -> bool {
120        self.offsets.windows(2).all(|pair| pair[1] == pair[0] + 1)
121    }
122
123    /// Set slop (max distance between terms)
124    pub fn with_slop(mut self, slop: u32) -> Self {
125        self.slop = slop;
126        self
127    }
128
129    /// Set global statistics for cross-segment IDF
130    pub fn with_global_stats(mut self, stats: Arc<GlobalStats>) -> Self {
131        self.global_stats = Some(stats);
132        self
133    }
134}
135
136/// Phrase over a chunked field: match and score every chunk (posting ids are
137/// virtual chunk ids, positions restart per chunk so a phrase never spans two
138/// chunks), then fold the chunk hits into documents with per-ordinal scores.
139///
140/// The scorer intersects the term postings and verifies positions in each
141/// conjunction candidate; only actual phrase hits enter the document fold.
142fn build_chunked_phrase_scorer<'a>(
143    term_data: Vec<(BlockPostingList, TermPositions)>,
144    offsets: &[u32],
145    slop: u32,
146    reader: &SegmentReader,
147    field: Field,
148    budget: Option<super::SharedThreshold>,
149) -> crate::Result<Box<dyn Scorer + 'a>> {
150    let Some(chunk_map) = reader.chunk_map(field) else {
151        return Err(crate::Error::Corruption(format!(
152            "chunked text field '{}' has postings but segment {:016x} carries no chunk map",
153            reader.schema().get_field_name(field).unwrap_or("?"),
154            reader.meta().id,
155        )));
156    };
157    let num_chunks = chunk_map.num_chunks() as f32;
158    let idf: f32 = term_data
159        .iter()
160        .map(|(p, _)| super::bm25_idf(p.doc_count() as f32, num_chunks))
161        .sum();
162    let (postings, positions): (Vec<_>, Vec<_>) = term_data.into_iter().unzip();
163    let scorer = PhraseScorer::new(
164        postings,
165        positions,
166        offsets,
167        slop,
168        idf,
169        chunk_map.avg_len(),
170        budget.clone(),
171    )
172    .with_lengths(Lengths::Chunks(chunk_map.clone()))
173    .with_params(super::Bm25Params::for_field(reader.schema(), field));
174
175    Ok(fold_chunked_phrase_scorer(
176        scorer,
177        chunk_map.clone(),
178        field.0,
179        budget,
180    ))
181}
182
183/// Ordered maps need only one document's ordinals and one matching-chunk
184/// lookahead. Reordered maps retain the stable, all-document aggregation.
185fn fold_chunked_phrase_scorer<'a, S: Scorer + 'a>(
186    mut scorer: S,
187    chunk_map: crate::segment::chunk_map::ChunkMap,
188    field_id: u32,
189    budget: Option<super::SharedThreshold>,
190) -> Box<dyn Scorer + 'a> {
191    if chunk_map.is_doc_ordered() {
192        let mut folded = ChunkedPhraseScorer {
193            inner: scorer,
194            chunk_map,
195            field_id,
196            budget,
197            current_doc: TERMINATED,
198            score: 0.0,
199            ordinals: crate::segment::VectorOrdinals::new(),
200        };
201        folded.fold_next_document();
202        return Box::new(folded);
203    }
204    let mut raw: Vec<(u32, u16, f32)> = Vec::new();
205    while scorer.doc() != TERMINATED {
206        if budget
207            .as_ref()
208            .is_some_and(super::SharedThreshold::stop_if_expired)
209        {
210            return Box::new(EmptyScorer);
211        }
212        let (doc_id, ordinal) = chunk_map.resolve(scorer.doc());
213        raw.push((doc_id, ordinal, scorer.score()));
214        scorer.advance();
215    }
216    if budget
217        .as_ref()
218        .is_some_and(super::SharedThreshold::stop_if_expired)
219    {
220        // Do not start an all-hit sort/fold after cancellation.
221        return Box::new(EmptyScorer);
222    }
223    // Every matching document is kept: a phrase is also used as a MUST
224    // constraint (verifier or bitset), where truncating to `limit` would
225    // silently reject documents that do contain the phrase.
226    let combined =
227        crate::segment::combine_ordinal_results(raw, super::MultiValueCombiner::Max, usize::MAX);
228    Box::new(super::vector::VectorResultScorer::new(combined, field_id))
229}
230
231struct ChunkedPhraseScorer<S> {
232    inner: S,
233    chunk_map: crate::segment::chunk_map::ChunkMap,
234    field_id: u32,
235    budget: Option<super::SharedThreshold>,
236    current_doc: DocId,
237    score: Score,
238    ordinals: crate::segment::VectorOrdinals,
239}
240
241impl<S: Scorer> ChunkedPhraseScorer<S> {
242    fn finish(&mut self) -> DocId {
243        self.current_doc = TERMINATED;
244        self.score = 0.0;
245        self.ordinals.clear();
246        TERMINATED
247    }
248
249    fn expired(&self) -> bool {
250        self.budget
251            .as_ref()
252            .is_some_and(super::SharedThreshold::stop_if_expired)
253    }
254
255    fn fold_next_document(&mut self) -> DocId {
256        if self.expired() || self.inner.doc() == TERMINATED {
257            return self.finish();
258        }
259        let doc = self.chunk_map.doc_id(self.inner.doc());
260        self.ordinals.clear();
261        loop {
262            self.ordinals.push((
263                u32::from(self.chunk_map.ordinal(self.inner.doc())),
264                self.inner.score(),
265            ));
266            self.inner.advance();
267            // Never expose a partial document's max score or ordinal list.
268            if self.expired() {
269                return self.finish();
270            }
271            if self.inner.doc() == TERMINATED || self.chunk_map.doc_id(self.inner.doc()) != doc {
272                break;
273            }
274        }
275        self.current_doc = doc;
276        self.score = super::MultiValueCombiner::Max.combine(&self.ordinals);
277        doc
278    }
279}
280
281impl<S: Scorer> DocSet for ChunkedPhraseScorer<S> {
282    fn doc(&self) -> DocId {
283        self.current_doc
284    }
285
286    fn advance(&mut self) -> DocId {
287        if self.current_doc == TERMINATED {
288            return TERMINATED;
289        }
290        self.fold_next_document()
291    }
292
293    fn seek(&mut self, target: DocId) -> DocId {
294        if self.current_doc >= target {
295            return self.current_doc;
296        }
297        if target == TERMINATED || self.expired() {
298            return self.finish();
299        }
300        let vid = self.chunk_map.lower_bound_doc(target);
301        if vid == self.chunk_map.num_chunks() {
302            return self.finish();
303        }
304        self.inner.seek(vid);
305        self.fold_next_document()
306    }
307
308    fn size_hint(&self) -> u32 {
309        if self.current_doc == TERMINATED {
310            0
311        } else {
312            self.inner.size_hint().saturating_add(1)
313        }
314    }
315}
316
317impl<S: Scorer> Scorer for ChunkedPhraseScorer<S> {
318    fn score(&self) -> Score {
319        self.score
320    }
321
322    fn matched_positions(&self) -> Option<super::MatchedPositions> {
323        (self.current_doc != TERMINATED).then(|| {
324            vec![(
325                self.field_id,
326                self.ordinals
327                    .iter()
328                    .map(|&(ordinal, score)| super::ScoredPosition::new(ordinal, score))
329                    .collect(),
330            )]
331        })
332    }
333}
334
335/// Build a PhraseScorer from already-fetched term data.
336fn build_phrase_scorer<'a>(
337    term_data: Vec<(BlockPostingList, TermPositions)>,
338    offsets: &[u32],
339    slop: u32,
340    reader: &SegmentReader,
341    field: Field,
342    budget: Option<super::SharedThreshold>,
343) -> Box<dyn Scorer + 'a> {
344    let idf: f32 = term_data
345        .iter()
346        .map(|(p, _)| {
347            let num_docs = reader.num_docs() as f32;
348            let doc_freq = p.doc_count() as f32;
349            super::bm25_idf(doc_freq, num_docs)
350        })
351        .sum();
352    let avg_field_len = reader.avg_field_len(field);
353    let (postings, positions): (Vec<_>, Vec<_>) = term_data.into_iter().unzip();
354    let mut scorer = PhraseScorer::new(
355        postings,
356        positions,
357        offsets,
358        slop,
359        idf,
360        avg_field_len,
361        budget,
362    )
363    .with_params(super::Bm25Params::for_field(reader.schema(), field));
364    if let Some(lengths) = reader.doc_lengths(field) {
365        scorer = scorer.with_lengths(Lengths::Docs(lengths.clone()));
366    }
367    Box::new(scorer)
368}
369
370// ── Shared early-return checks for phrase scorer ─────────────────────────
371//
372// Handles: empty terms, single-term delegation, no-positions fallback.
373// Parameterised on the option-aware scorer function plus async/sync awaiting.
374macro_rules! phrase_early_returns {
375    ($field:expr, $terms:expr, $reader:expr, $limit:expr,
376     $scorer_fn:ident, $options:expr $(, $aw:tt)*) => {
377        if $options.stop_if_expired() || $terms.is_empty() {
378            return Ok(Box::new(EmptyScorer) as Box<dyn Scorer + '_>);
379        }
380        if $terms.len() == 1 {
381            let tq = super::TermQuery::new($field, $terms[0].clone());
382            return tq.$scorer_fn($reader, $limit, $options) $(. $aw)* ;
383        }
384        if !$reader.has_positions($field) {
385            let mut bq = super::BooleanQuery::new();
386            for t in $terms.iter() {
387                bq = bq.must(super::TermQuery::new($field, t.clone()));
388            }
389            return bq.$scorer_fn($reader, $limit, $options) $(. $aw)* ;
390        }
391    };
392}
393
394impl Query for PhraseQuery {
395    fn text_terms(&self, out: &mut Vec<(Field, Vec<u8>)>) {
396        for term in &self.terms {
397            out.push((self.field, term.clone()));
398        }
399    }
400
401    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
402        self.scorer_with_options(reader, limit, super::ScorerOptions::with_positions())
403    }
404
405    fn scorer_with_options<'a>(
406        &self,
407        reader: &'a SegmentReader,
408        limit: usize,
409        options: super::ScorerOptions,
410    ) -> ScorerFuture<'a> {
411        let field = self.field;
412        let terms = self.terms.clone();
413        let offsets = self.offsets.clone();
414        let slop = self.slop;
415
416        Box::pin(async move {
417            phrase_early_returns!(
418                field,
419                terms,
420                reader,
421                limit,
422                scorer_with_options,
423                options,
424                await
425            );
426
427            // Fetch postings + positions in parallel per term via futures::join!
428            let mut term_data = Vec::with_capacity(terms.len());
429            for term in &terms {
430                let (postings, positions) = futures::join!(
431                    reader.get_postings(field, term),
432                    reader.get_positions(field, term)
433                );
434                match (postings?, positions?) {
435                    (Some(p), Some(pos)) => term_data.push((p, pos)),
436                    _ => return Ok(Box::new(EmptyScorer) as Box<dyn Scorer + 'a>),
437                }
438            }
439
440            if reader.is_chunked_field(field) {
441                return build_chunked_phrase_scorer(
442                    term_data,
443                    &offsets,
444                    slop,
445                    reader,
446                    field,
447                    options.shared_threshold,
448                );
449            }
450            Ok(build_phrase_scorer(
451                term_data,
452                &offsets,
453                slop,
454                reader,
455                field,
456                options.shared_threshold,
457            ))
458        })
459    }
460
461    #[cfg(feature = "sync")]
462    fn scorer_sync<'a>(
463        &self,
464        reader: &'a SegmentReader,
465        limit: usize,
466    ) -> crate::Result<Box<dyn Scorer + 'a>> {
467        self.scorer_sync_with_options(reader, limit, super::ScorerOptions::with_positions())
468    }
469
470    #[cfg(feature = "sync")]
471    fn scorer_sync_with_options<'a>(
472        &self,
473        reader: &'a SegmentReader,
474        limit: usize,
475        options: super::ScorerOptions,
476    ) -> crate::Result<Box<dyn Scorer + 'a>> {
477        phrase_early_returns!(
478            self.field,
479            self.terms,
480            reader,
481            limit,
482            scorer_sync_with_options,
483            options
484        );
485
486        // Parallel fetch across all terms via rayon
487        use rayon::prelude::*;
488        let pairs: crate::Result<Vec<Option<(BlockPostingList, TermPositions)>>> = self
489            .terms
490            .par_iter()
491            .map(|term| {
492                let postings = reader.get_postings_sync(self.field, term)?;
493                let positions = reader.get_positions_sync(self.field, term)?;
494                Ok(match (postings, positions) {
495                    (Some(p), Some(pos)) => Some((p, pos)),
496                    _ => None,
497                })
498            })
499            .collect();
500        let mut term_data = Vec::with_capacity(self.terms.len());
501        for entry in pairs? {
502            match entry {
503                Some(pair) => term_data.push(pair),
504                None => return Ok(Box::new(EmptyScorer) as Box<dyn Scorer + 'a>),
505            }
506        }
507
508        if reader.is_chunked_field(self.field) {
509            return build_chunked_phrase_scorer(
510                term_data,
511                &self.offsets,
512                self.slop,
513                reader,
514                self.field,
515                options.shared_threshold,
516            );
517        }
518        Ok(build_phrase_scorer(
519            term_data,
520            &self.offsets,
521            self.slop,
522            reader,
523            self.field,
524            options.shared_threshold,
525        ))
526    }
527
528    /// Every document containing the phrase, as a bitset (documents, also
529    /// for chunked fields). Lets the planner push a quoted span into the
530    /// MaxScore executors as an O(1) predicate instead of a verifier.
531    #[cfg(feature = "sync")]
532    fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
533        self.as_doc_bitset_with_options(reader, &super::ScorerOptions::default())
534    }
535
536    #[cfg(feature = "sync")]
537    fn as_doc_bitset_with_options(
538        &self,
539        reader: &SegmentReader,
540        options: &super::ScorerOptions,
541    ) -> Option<super::DocBitset> {
542        if options.stop_if_expired() || self.terms.is_empty() {
543            return None;
544        }
545        let mut bitset = super::DocBitset::new(reader.num_docs());
546        if self.terms.len() == 1 {
547            // A one-term phrase is the term itself; walk its postings and
548            // resolve chunk ids to documents where needed.
549            let list = reader
550                .get_postings_sync(self.field, &self.terms[0])
551                .ok()??;
552            let chunk_map = reader.chunk_map(self.field);
553            let mut it = list.iterator();
554            while it.doc() != TERMINATED {
555                if options.stop_if_expired() {
556                    return None;
557                }
558                let doc = chunk_map.map_or(it.doc(), |map| map.doc_id(it.doc()));
559                bitset.set(doc);
560                it.advance();
561            }
562            return Some(bitset);
563        }
564        let mut scorer = self
565            .scorer_sync_with_options(reader, usize::MAX, options.without_threshold())
566            .ok()?;
567        while scorer.doc() != TERMINATED {
568            if options.stop_if_expired() {
569                return None;
570            }
571            bitset.set(scorer.doc());
572            scorer.advance();
573        }
574        if options.stop_if_expired() {
575            None
576        } else {
577            Some(bitset)
578        }
579    }
580
581    /// Matches are at most the rarest term's postings; the planner only
582    /// needs the order of magnitude to pick which clause to materialize.
583    #[cfg(feature = "sync")]
584    fn bitset_cardinality_estimate(&self, reader: &SegmentReader) -> Option<u64> {
585        let mut min = u64::MAX;
586        for term in &self.terms {
587            let list = reader.get_postings_sync(self.field, term).ok()??;
588            min = min.min(u64::from(list.doc_count()));
589        }
590        Some((min / 10).max(1))
591    }
592
593    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
594        let field = self.field;
595        let terms = self.terms.clone();
596
597        Box::pin(async move {
598            if terms.is_empty() {
599                return Ok(0);
600            }
601
602            // Estimate based on minimum posting list size
603            let mut min_count = u32::MAX;
604            for term in &terms {
605                match reader.get_postings(field, term).await? {
606                    Some(list) => min_count = min_count.min(list.doc_count()),
607                    None => return Ok(0),
608                }
609            }
610
611            // Phrase matching will typically match fewer docs than the minimum
612            // Estimate ~10% of the smallest posting list
613            Ok((min_count / 10).max(1))
614        })
615    }
616}
617
618/// Real lengths of the scoring units of a phrase: chunk lengths of a chunked
619/// field or persisted document lengths of a plain field.
620enum Lengths {
621    Chunks(crate::segment::chunk_map::ChunkMap),
622    Docs(crate::segment::chunk_map::DocLengths),
623}
624
625impl Lengths {
626    fn length(&self, id: u32) -> u32 {
627        match self {
628            Lengths::Chunks(map) => map.bm25_length(id),
629            Lengths::Docs(lengths) => lengths.length(id),
630        }
631    }
632}
633
634/// Scorer that checks phrase positions
635struct PhraseScorer {
636    budget: Option<super::SharedThreshold>,
637    /// Posting iterators for each term
638    posting_iters: Vec<BlockPostingIterator<'static>>,
639    /// Positions of each term (legacy list or cursor-addressed stream)
640    position_lists: Vec<crate::structures::postings::TermPositionCursor>,
641    /// Position-list cursors advance monotonically within a matching unit.
642    position_indices: Vec<usize>,
643    /// Required distance of each term from the first one (`offsets[i] -
644    /// offsets[0]`); `deltas[0]` is 0.
645    deltas: Vec<u32>,
646    /// Max slop between terms
647    slop: u32,
648    /// Current matching document
649    current_doc: DocId,
650    /// Number of phrase occurrences in the current document (phrase
651    /// frequency), the `tf` of the phrase for BM25.
652    current_matches: u32,
653    /// Combined IDF
654    idf: f32,
655    /// Per-field k1/b.
656    params: super::Bm25Params,
657    /// Average field length
658    avg_field_len: f32,
659    /// Real lengths of the scoring units. `None` keeps the historic
660    /// `tf`-as-length approximation.
661    lengths: Option<Lengths>,
662    /// Reusable position buffers (one per term, avoids per-document allocation)
663    position_bufs: Vec<Vec<u32>>,
664}
665
666impl PhraseScorer {
667    #[allow(clippy::too_many_arguments)]
668    fn new(
669        posting_lists: Vec<BlockPostingList>,
670        position_lists: Vec<TermPositions>,
671        offsets: &[u32],
672        slop: u32,
673        idf: f32,
674        avg_field_len: f32,
675        budget: Option<super::SharedThreshold>,
676    ) -> Self {
677        let posting_iters: Vec<_> = posting_lists
678            .into_iter()
679            .map(|p| p.into_iterator())
680            .collect();
681
682        let num_terms = position_lists.len();
683        // Offsets are optional for callers that built the query term by
684        // term; missing entries mean adjacency.
685        let first = offsets.first().copied().unwrap_or(0);
686        let deltas: Vec<u32> = (0..num_terms)
687            .map(|i| offsets.get(i).map_or(i as u32, |o| o - first))
688            .collect();
689        let mut scorer = Self {
690            budget: budget.filter(|b| b.deadline().is_some()),
691            posting_iters,
692            position_lists: position_lists
693                .into_iter()
694                .map(TermPositions::into_cursor)
695                .collect(),
696            position_indices: vec![0; num_terms],
697            deltas,
698            slop,
699            current_doc: 0,
700            current_matches: 0,
701            params: super::Bm25Params::default(),
702            idf,
703            avg_field_len,
704            lengths: None,
705            position_bufs: (0..num_terms).map(|_| Vec::new()).collect(),
706        };
707
708        scorer.find_next_phrase_match();
709        scorer
710    }
711
712    /// Score with the real length of each scoring unit.
713    fn with_lengths(mut self, lengths: Lengths) -> Self {
714        self.lengths = Some(lengths);
715        self
716    }
717
718    /// Score with the field's BM25 parameters.
719    fn with_params(mut self, params: super::Bm25Params) -> Self {
720        self.params = params;
721        self
722    }
723
724    /// Find next document where all terms appear as a phrase
725    fn find_next_phrase_match(&mut self) {
726        loop {
727            // First, find a document where all terms appear (AND semantics)
728            let doc = self.find_next_and_match();
729            if doc == TERMINATED {
730                self.current_doc = TERMINATED;
731                return;
732            }
733
734            // Check if positions form a valid phrase
735            if self.check_phrase_positions(doc) {
736                self.current_doc = doc;
737                return;
738            }
739
740            // Advance and try again
741            self.posting_iters[0].advance();
742        }
743    }
744
745    /// Find next document where all terms appear
746    fn find_next_and_match(&mut self) -> DocId {
747        if self.posting_iters.is_empty() {
748            return TERMINATED;
749        }
750
751        loop {
752            if self
753                .budget
754                .as_ref()
755                .is_some_and(super::SharedThreshold::stop_if_expired)
756            {
757                return TERMINATED;
758            }
759            let max_doc = self.posting_iters.iter().map(|it| it.doc()).max().unwrap();
760
761            if max_doc == TERMINATED {
762                return TERMINATED;
763            }
764
765            let mut all_match = true;
766            for it in &mut self.posting_iters {
767                let doc = it.seek(max_doc);
768                if doc != max_doc {
769                    all_match = false;
770                    if doc == TERMINATED {
771                        return TERMINATED;
772                    }
773                }
774            }
775
776            if all_match {
777                return max_doc;
778            }
779        }
780    }
781
782    /// Check if positions form a valid phrase for the given document
783    fn check_phrase_positions(&mut self, doc_id: DocId) -> bool {
784        // Get positions for each term into reusable buffers (zero allocation).
785        // The doc-posting iterator of every term is parked on `doc_id`, so
786        // its cursor and term frequency address the term's position stream.
787        for i in 0..self.position_lists.len() {
788            let cursor = self.posting_iters[i].position_cursor();
789            let tf = self.posting_iters[i].term_freq();
790            if !self.position_lists[i].read_into(doc_id, cursor, tf, &mut self.position_bufs[i]) {
791                return false;
792            }
793        }
794
795        // Count the occurrences: every position of the first term that
796        // starts a full match. The count is the phrase frequency BM25 scores.
797        self.current_matches = self.count_phrase_matches_in_bufs();
798        self.current_matches > 0
799    }
800
801    /// Number of phrase occurrences in the internal reusable buffers.
802    fn count_phrase_matches_in_bufs(&mut self) -> u32 {
803        count_phrase_matches(
804            &self.position_bufs,
805            &self.deltas,
806            self.slop,
807            &mut self.position_indices,
808        )
809    }
810}
811
812/// Count starts with a match in every term's independent slop interval.
813/// Ascending starts make each interval monotone: O(terms * starts + positions),
814/// preserving repeated terms and the existing (not edit-distance) slop rule.
815fn count_phrase_matches(
816    bufs: &[Vec<u32>],
817    deltas: &[u32],
818    slop: u32,
819    indices: &mut [usize],
820) -> u32 {
821    let Some(first) = bufs.first() else {
822        return 0;
823    };
824    indices.fill(0);
825    let mut matches = 0;
826    'starts: for &start in first {
827        for i in 1..bufs.len() {
828            let expected = u64::from(start) + u64::from(deltas[i]);
829            let low = expected.saturating_sub(u64::from(slop));
830            let high = expected + u64::from(slop);
831            while indices[i] < bufs[i].len() && u64::from(bufs[i][indices[i]]) < low {
832                indices[i] += 1;
833            }
834            let Some(&position) = bufs[i].get(indices[i]) else {
835                return matches;
836            };
837            if u64::from(position) > high {
838                continue 'starts;
839            }
840        }
841        matches += 1;
842    }
843    matches
844}
845
846impl super::docset::DocSet for PhraseScorer {
847    fn doc(&self) -> DocId {
848        self.current_doc
849    }
850
851    fn advance(&mut self) -> DocId {
852        if self.current_doc == TERMINATED {
853            return TERMINATED;
854        }
855
856        self.posting_iters[0].advance();
857        self.find_next_phrase_match();
858        self.current_doc
859    }
860
861    fn seek(&mut self, target: DocId) -> DocId {
862        if target == TERMINATED {
863            self.current_doc = TERMINATED;
864            return TERMINATED;
865        }
866
867        self.posting_iters[0].seek(target);
868        self.find_next_phrase_match();
869        self.current_doc
870    }
871
872    fn size_hint(&self) -> u32 {
873        0
874    }
875}
876
877impl Scorer for PhraseScorer {
878    fn score(&self) -> Score {
879        if self.current_doc == TERMINATED {
880            return 0.0;
881        }
882
883        // BM25 over the phrase frequency with the summed idf of the terms
884        // (Lucene semantics): a document with two occurrences of the phrase
885        // outranks one with a single occurrence at equal length.
886        let tf = self.current_matches.max(1) as f32;
887
888        // Real unit length when the segment has it; otherwise the summed
889        // term frequency stands in for the length (legacy segments).
890        let doc_len = match &self.lengths {
891            Some(lengths) => (lengths.length(self.current_doc) as f32).max(1.0),
892            None => self
893                .posting_iters
894                .iter()
895                .map(|it| it.term_freq() as f32)
896                .sum::<f32>()
897                .max(tf),
898        };
899
900        self.params.score(tf, self.idf, doc_len, self.avg_field_len)
901    }
902}
903
904#[cfg(test)]
905mod tests {
906    use super::*;
907
908    struct ChunkHits {
909        hits: Vec<(u32, f32)>,
910        at: usize,
911        advances: Arc<std::sync::atomic::AtomicUsize>,
912    }
913
914    impl DocSet for ChunkHits {
915        fn doc(&self) -> DocId {
916            self.hits.get(self.at).map_or(TERMINATED, |h| h.0)
917        }
918        fn advance(&mut self) -> DocId {
919            self.advances
920                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
921            self.at = (self.at + 1).min(self.hits.len());
922            self.doc()
923        }
924        fn seek(&mut self, target: DocId) -> DocId {
925            self.at += self.hits[self.at..].partition_point(|h| h.0 < target);
926            self.doc()
927        }
928        fn size_hint(&self) -> u32 {
929            (self.hits.len() - self.at) as u32
930        }
931    }
932    impl Scorer for ChunkHits {
933        fn score(&self) -> Score {
934            self.hits.get(self.at).map_or(0.0, |h| h.1)
935        }
936    }
937
938    fn test_chunk_map(owners: &[(u32, u16)]) -> crate::segment::chunk_map::ChunkMap {
939        use crate::segment::chunk_map::{ChunkMapBuilder, read_chunk_maps, write_chunk_maps};
940        let mut builder = ChunkMapBuilder::default();
941        for &(doc, ordinal) in owners {
942            builder.push(doc, ordinal, 10).unwrap();
943        }
944        let mut bytes = Vec::new();
945        write_chunk_maps(&mut bytes, &[(0, &builder)], &[]).unwrap();
946        read_chunk_maps(crate::directories::OwnedBytes::new(bytes))
947            .unwrap()
948            .chunk_maps
949            .remove(&0)
950            .unwrap()
951    }
952
953    fn chunk_hits(hits: Vec<(u32, f32)>) -> ChunkHits {
954        ChunkHits {
955            hits,
956            at: 0,
957            advances: Arc::default(),
958        }
959    }
960
961    #[test]
962    fn lazy_phrase_fold_matches_stable_eager_oracle_including_reordered_ordinals() {
963        for owners in [
964            vec![],
965            vec![(0, 0)],
966            vec![(2, 2), (2, 0), (2, 1), (5, 1), (5, 0), (9, 0)],
967            vec![(5, 1), (2, 2), (9, 0), (2, 0), (5, 0), (2, 1)],
968        ] {
969            let map = test_chunk_map(&owners);
970            assert_eq!(
971                map.is_doc_ordered(),
972                owners.windows(2).all(|p| p[0].0 <= p[1].0)
973            );
974            for stride in [1, 2, 3] {
975                let hits: Vec<_> = (0..owners.len() as u32)
976                    .step_by(stride)
977                    .map(|vid| (vid, (vid % 3) as f32 * 0.5))
978                    .collect();
979                let raw: Vec<_> = hits
980                    .iter()
981                    .map(|&(vid, score)| {
982                        let (doc, ord) = map.resolve(vid);
983                        (doc, ord, score)
984                    })
985                    .collect();
986                let expected = crate::segment::combine_ordinal_results(
987                    raw,
988                    super::super::MultiValueCombiner::Max,
989                    usize::MAX,
990                );
991                let mut expected = super::super::vector::VectorResultScorer::new(expected, 7);
992                let mut actual = fold_chunked_phrase_scorer(chunk_hits(hits), map.clone(), 7, None);
993                while expected.doc() != TERMINATED {
994                    assert_eq!(actual.doc(), expected.doc());
995                    assert_eq!(actual.score().to_bits(), expected.score().to_bits());
996                    let signature = |s: &dyn Scorer| {
997                        s.matched_positions()
998                            .unwrap()
999                            .into_iter()
1000                            .map(|(field, positions)| {
1001                                (
1002                                    field,
1003                                    positions
1004                                        .into_iter()
1005                                        .map(|p| (p.position, p.score.to_bits()))
1006                                        .collect::<Vec<_>>(),
1007                                )
1008                            })
1009                            .collect::<Vec<_>>()
1010                    };
1011                    assert_eq!(signature(actual.as_ref()), signature(&expected));
1012                    actual.advance();
1013                    expected.advance();
1014                }
1015                assert_eq!(actual.doc(), TERMINATED);
1016                assert_eq!(actual.advance(), TERMINATED);
1017                assert_eq!(actual.score(), 0.0);
1018            }
1019        }
1020    }
1021
1022    #[test]
1023    fn lazy_phrase_fold_only_consumes_one_document_and_can_skip_to_late_matches() {
1024        let owners: Vec<_> = (0..100)
1025            .flat_map(|doc| [(doc * 2, 0), (doc * 2, 1)])
1026            .collect();
1027        let inner = chunk_hits((0..200).map(|vid| (vid, vid as f32)).collect());
1028        let advances = inner.advances.clone();
1029        let mut scorer = fold_chunked_phrase_scorer(inner, test_chunk_map(&owners), 0, None);
1030        assert_eq!(advances.load(std::sync::atomic::Ordering::Relaxed), 2);
1031        assert_eq!(scorer.doc(), 0);
1032        assert_eq!(scorer.seek(179), 180);
1033        assert_eq!(advances.load(std::sync::atomic::Ordering::Relaxed), 4);
1034        assert_eq!(scorer.score(), 181.0);
1035        assert_eq!(scorer.seek(179), 180);
1036        assert_eq!(scorer.seek(199), TERMINATED);
1037        assert!(scorer.matched_positions().is_none());
1038        assert_eq!(scorer.advance(), TERMINATED);
1039    }
1040
1041    #[test]
1042    fn lazy_phrase_fold_discards_current_result_at_budget_boundary() {
1043        let inner = chunk_hits(vec![(0, 1.0), (1, 2.0), (2, 3.0)]);
1044        let advances = inner.advances.clone();
1045        let mut scorer = ChunkedPhraseScorer {
1046            inner,
1047            chunk_map: test_chunk_map(&[(0, 0), (1, 0), (1, 1)]),
1048            field_id: 0,
1049            budget: None,
1050            current_doc: TERMINATED,
1051            score: 0.0,
1052            ordinals: crate::segment::VectorOrdinals::new(),
1053        };
1054        assert_eq!(scorer.fold_next_document(), 0);
1055        assert_eq!(scorer.score(), 1.0);
1056        let budget = super::super::SharedThreshold::for_limit(1)
1057            .with_deadline(Some(std::time::Instant::now()));
1058        scorer.budget = Some(budget.clone());
1059        assert_eq!(scorer.advance(), TERMINATED);
1060        assert_eq!(scorer.score(), 0.0);
1061        assert!(scorer.matched_positions().is_none());
1062        assert!(budget.truncated());
1063        assert_eq!(advances.load(std::sync::atomic::Ordering::Relaxed), 1);
1064    }
1065
1066    #[test]
1067    fn phrase_stops_when_budget_expires_after_first_match() {
1068        use super::super::docset::DocSet;
1069        use crate::structures::{PositionStreamEncoder, PostingList};
1070        let mut lists = Vec::new();
1071        let mut positions = Vec::new();
1072        for term in 0..2 {
1073            let mut list = PostingList::new();
1074            let mut bytes = Vec::new();
1075            let mut encoder = PositionStreamEncoder::new(&mut bytes);
1076            for doc in 0..1000 {
1077                list.push(doc, 1);
1078                encoder
1079                    .push_doc(&mut [if doc == 0 { term } else { term * 10 }])
1080                    .unwrap();
1081            }
1082            encoder.finish().unwrap();
1083            lists.push(BlockPostingList::from_posting_list_with(&list, true, None).unwrap());
1084            positions
1085                .push(TermPositions::open(crate::directories::OwnedBytes::new(bytes)).unwrap());
1086        }
1087        let mut scorer = PhraseScorer::new(lists, positions, &[0, 1], 0, 1.0, 2.0, None);
1088        assert_eq!(scorer.doc(), 0);
1089        let budget = super::super::SharedThreshold::for_limit(10)
1090            .with_deadline(Some(std::time::Instant::now()));
1091        scorer.budget = Some(budget.clone());
1092        assert_eq!(scorer.advance(), TERMINATED);
1093        assert_eq!(scorer.score(), 0.0);
1094        assert!(budget.truncated());
1095        assert_eq!(
1096            scorer.position_bufs,
1097            vec![vec![0], vec![1]],
1098            "no further positions decoded"
1099        );
1100    }
1101
1102    #[test]
1103    fn monotone_phrase_frequency_matches_naive_offsets_slop_and_repeated_terms() {
1104        for seed in 0..100u32 {
1105            let a: Vec<_> = (0..200).filter(|i| (i * 17 + seed) % 11 < 5).collect();
1106            let b: Vec<_> = (0..200).filter(|i| (i * 13 + seed) % 19 < 4).collect();
1107            for bufs in [
1108                vec![a.clone(), b.clone()],
1109                vec![a.clone(), b.clone(), a.clone()],
1110            ] {
1111                for slop in [0, 1, 3, 100] {
1112                    let deltas = [0, 2, 7];
1113                    let expected = bufs[0]
1114                        .iter()
1115                        .filter(|&&start| {
1116                            bufs.iter().enumerate().skip(1).all(|(i, positions)| {
1117                                positions
1118                                    .iter()
1119                                    .any(|&p| p.abs_diff(start + deltas[i]) <= slop)
1120                            })
1121                        })
1122                        .count() as u32;
1123                    assert_eq!(
1124                        count_phrase_matches(&bufs, &deltas, slop, &mut [0; 3]),
1125                        expected
1126                    );
1127                }
1128            }
1129        }
1130        assert_eq!(
1131            count_phrase_matches(&[vec![0, 0, 5], vec![1, 6]], &[0, 1], 0, &mut [0; 2]),
1132            3
1133        );
1134        assert_eq!(
1135            count_phrase_matches(&[vec![u32::MAX], vec![0]], &[0, 1], 0, &mut [0; 2]),
1136            0
1137        );
1138        assert_eq!(
1139            count_phrase_matches(&[vec![1], vec![]], &[0, 1], 3, &mut [0; 2]),
1140            0
1141        );
1142    }
1143}