1use 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#[derive(Clone)]
18pub struct PhraseQuery {
19 pub field: Field,
20 pub terms: Vec<Vec<u8>>,
22 pub offsets: Vec<u32>,
28 pub slop: u32,
30 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 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 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 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 fn is_adjacent(&self) -> bool {
120 self.offsets.windows(2).all(|pair| pair[1] == pair[0] + 1)
121 }
122
123 pub fn with_slop(mut self, slop: u32) -> Self {
125 self.slop = slop;
126 self
127 }
128
129 pub fn with_global_stats(mut self, stats: Arc<GlobalStats>) -> Self {
131 self.global_stats = Some(stats);
132 self
133 }
134}
135
136fn 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
183fn 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 return Box::new(EmptyScorer);
222 }
223 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 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
335fn 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
370macro_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 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 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 #[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 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 #[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 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 Ok((min_count / 10).max(1))
614 })
615 }
616}
617
618enum 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
634struct PhraseScorer {
636 budget: Option<super::SharedThreshold>,
637 posting_iters: Vec<BlockPostingIterator<'static>>,
639 position_lists: Vec<crate::structures::postings::TermPositionCursor>,
641 position_indices: Vec<usize>,
643 deltas: Vec<u32>,
646 slop: u32,
648 current_doc: DocId,
650 current_matches: u32,
653 idf: f32,
655 params: super::Bm25Params,
657 avg_field_len: f32,
659 lengths: Option<Lengths>,
662 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 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 fn with_lengths(mut self, lengths: Lengths) -> Self {
714 self.lengths = Some(lengths);
715 self
716 }
717
718 fn with_params(mut self, params: super::Bm25Params) -> Self {
720 self.params = params;
721 self
722 }
723
724 fn find_next_phrase_match(&mut self) {
726 loop {
727 let doc = self.find_next_and_match();
729 if doc == TERMINATED {
730 self.current_doc = TERMINATED;
731 return;
732 }
733
734 if self.check_phrase_positions(doc) {
736 self.current_doc = doc;
737 return;
738 }
739
740 self.posting_iters[0].advance();
742 }
743 }
744
745 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 fn check_phrase_positions(&mut self, doc_id: DocId) -> bool {
784 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 self.current_matches = self.count_phrase_matches_in_bufs();
798 self.current_matches > 0
799 }
800
801 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
812fn 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 let tf = self.current_matches.max(1) as f32;
887
888 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}