From 98be1a542317a2c9c811b54712e36d43a5c82fbd Mon Sep 17 00:00:00 2001 From: Paul Masurel Date: Tue, 30 Dec 2025 17:26:15 +0100 Subject: [PATCH] Added seek_doc to intersections. tantivy requires Scorer to be positioned on a DocId at all time. This decision is not performance neutral. When we have an intersection of a heavy DocSet with a lighter one forcing the positioning of the first doc is needlessly expensive. This PR fixes this by introducing a seek_doc parameter in the weight function. Weights may skip over documents when they create the Scorer. --- src/aggregation/bucket/filter.rs | 2 +- src/docset.rs | 14 +++ src/fieldnorm/mod.rs | 4 +- src/index/inverted_index_reader.rs | 48 +++++++--- src/indexer/delete_queue.rs | 7 +- src/indexer/merger.rs | 7 +- src/postings/block_segment_postings.rs | 58 ++++++++---- src/postings/mod.rs | 1 + src/postings/segment_postings.rs | 23 ++--- src/query/all_query.rs | 15 ++- src/query/automaton_weight.rs | 17 +++- src/query/boolean_query/boolean_weight.rs | 64 ++++++++++--- src/query/boolean_query/mod.rs | 24 ++--- src/query/boost_query.rs | 13 ++- src/query/const_score_query.rs | 15 ++- src/query/empty_query.rs | 13 ++- src/query/exist_query.rs | 9 +- .../phrase_prefix_weight.rs | 42 ++++++--- src/query/phrase_query/mod.rs | 41 +++++++- src/query/phrase_query/phrase_weight.rs | 36 ++++--- src/query/phrase_query/regex_phrase_weight.rs | 14 ++- .../range_query/fast_field_range_doc_set.rs | 8 +- src/query/range_query/range_query.rs | 15 ++- .../range_query/range_query_fastfield.rs | 94 ++++++++++++++++--- src/query/term_query/mod.rs | 57 ++++++++++- src/query/term_query/term_weight.rs | 37 ++++++-- src/query/weight.rs | 29 +++++- 27 files changed, 549 insertions(+), 158 deletions(-) diff --git a/src/aggregation/bucket/filter.rs b/src/aggregation/bucket/filter.rs index 18f2a692a..4b88aa195 100644 --- a/src/aggregation/bucket/filter.rs +++ b/src/aggregation/bucket/filter.rs @@ -446,7 +446,7 @@ impl DocumentQueryEvaluator { let weight = query.weight(EnableScoring::disabled_from_schema(&schema))?; // Get a scorer that iterates over matching documents - let mut scorer = weight.scorer(segment_reader, 1.0)?; + let mut scorer = weight.scorer(segment_reader, 1.0, 0)?; // Create a BitSet to hold all matching documents let mut bitset = BitSet::with_max_value(max_doc); diff --git a/src/docset.rs b/src/docset.rs index 01ea1125a..ce420ac1d 100644 --- a/src/docset.rs +++ b/src/docset.rs @@ -42,6 +42,7 @@ pub trait DocSet: Send { /// Calling `seek(TERMINATED)` is also legal and is the normal way to consume a `DocSet`. /// /// `target` has to be larger or equal to `.doc()` when calling `seek`. + /// If `target` is equal to `.doc()` then the DocSet should not advance. fn seek(&mut self, target: DocId) -> DocId { let mut doc = self.doc(); debug_assert!(doc <= target); @@ -166,6 +167,19 @@ pub trait DocSet: Send { } } +/// Consumes the `DocSet` and returns a Vec with all of the docs in the DocSet +/// including the current doc. +#[cfg(test)] +pub fn docset_to_doc_vec(mut doc_set: Box) -> Vec { + let mut output = Vec::new(); + let mut doc = doc_set.doc(); + while doc != TERMINATED { + output.push(doc); + doc = doc_set.advance(); + } + output +} + impl DocSet for &mut dyn DocSet { fn advance(&mut self) -> u32 { (**self).advance() diff --git a/src/fieldnorm/mod.rs b/src/fieldnorm/mod.rs index 32c957d1e..97490e6ab 100644 --- a/src/fieldnorm/mod.rs +++ b/src/fieldnorm/mod.rs @@ -113,7 +113,7 @@ mod tests { IndexRecordOption::WithFreqs, ); let weight = query.weight(EnableScoring::enabled_from_searcher(&searcher))?; - let mut scorer = weight.scorer(searcher.segment_reader(0), 1.0f32)?; + let mut scorer = weight.scorer(searcher.segment_reader(0), 1.0f32, 0)?; assert_eq!(scorer.doc(), 0); assert!((scorer.score() - 0.22920431).abs() < 0.001f32); assert_eq!(scorer.advance(), 1); @@ -142,7 +142,7 @@ mod tests { IndexRecordOption::WithFreqs, ); let weight = query.weight(EnableScoring::enabled_from_searcher(&searcher))?; - let mut scorer = weight.scorer(searcher.segment_reader(0), 1.0f32)?; + let mut scorer = weight.scorer(searcher.segment_reader(0), 1.0f32, 0)?; assert_eq!(scorer.doc(), 0); assert!((scorer.score() - 0.22920431).abs() < 0.001f32); assert_eq!(scorer.advance(), 1); diff --git a/src/index/inverted_index_reader.rs b/src/index/inverted_index_reader.rs index 7314f8741..3ce4b555f 100644 --- a/src/index/inverted_index_reader.rs +++ b/src/index/inverted_index_reader.rs @@ -11,9 +11,12 @@ use tantivy_fst::automaton::{AlwaysMatch, Automaton}; use crate::directory::FileSlice; use crate::positions::PositionReader; -use crate::postings::{BlockSegmentPostings, SegmentPostings, TermInfo}; +use crate::postings::{ + BlockSegmentPostings, BlockSegmentPostingsNotLoaded, SegmentPostings, TermInfo, +}; use crate::schema::{IndexRecordOption, Term, Type}; use crate::termdict::TermDictionary; +use crate::DocId; /// The inverted index reader is in charge of accessing /// the inverted index associated with a specific field. @@ -192,20 +195,23 @@ impl InvertedIndexReader { term: &Term, option: IndexRecordOption, ) -> io::Result> { - self.get_term_info(term)? - .map(move |term_info| self.read_block_postings_from_terminfo(&term_info, option)) - .transpose() + let Some(term_info) = self.get_term_info(term)? else { + return Ok(None); + }; + let block_postings_not_loaded = + self.read_block_postings_from_terminfo(&term_info, option)?; + Ok(Some(block_postings_not_loaded)) } /// Returns a block postings given a `term_info`. /// This method is for an advanced usage only. /// /// Most users should prefer using [`Self::read_postings()`] instead. - pub fn read_block_postings_from_terminfo( + pub(crate) fn read_block_postings_from_terminfo_not_loaded( &self, term_info: &TermInfo, requested_option: IndexRecordOption, - ) -> io::Result { + ) -> io::Result { let postings_data = self .postings_file_slice .slice(term_info.postings_range.clone()); @@ -217,6 +223,21 @@ impl InvertedIndexReader { ) } + /// Returns a block postings given a `term_info`. + /// This method is for an advanced usage only. + /// + /// Most users should prefer using [`Self::read_postings()`] instead. + pub fn read_block_postings_from_terminfo( + &self, + term_info: &TermInfo, + requested_option: IndexRecordOption, + ) -> io::Result { + let block_segment_postings_not_loaded = self + .read_block_postings_from_terminfo_not_loaded(term_info, requested_option)? + .load_at_start(); + Ok(block_segment_postings_not_loaded) + } + /// Returns a posting object given a `term_info`. /// This method is for an advanced usage only. /// @@ -224,13 +245,13 @@ impl InvertedIndexReader { pub fn read_postings_from_terminfo( &self, term_info: &TermInfo, - option: IndexRecordOption, + record_option: IndexRecordOption, + seek_doc: DocId, ) -> io::Result { - let option = option.downgrade(self.record_option); - - let block_postings = self.read_block_postings_from_terminfo(term_info, option)?; + let block_segment_postings_not_loaded = + self.read_block_postings_from_terminfo_not_loaded(term_info, record_option)?; let position_reader = { - if option.has_positions() { + if record_option.has_positions() { let positions_data = self .positions_file_slice .read_bytes_slice(term_info.positions_range.clone())?; @@ -241,8 +262,9 @@ impl InvertedIndexReader { } }; Ok(SegmentPostings::from_block_postings( - block_postings, + block_segment_postings_not_loaded, position_reader, + seek_doc, )) } @@ -268,7 +290,7 @@ impl InvertedIndexReader { option: IndexRecordOption, ) -> io::Result> { self.get_term_info(term)? - .map(move |term_info| self.read_postings_from_terminfo(&term_info, option)) + .map(move |term_info| self.read_postings_from_terminfo(&term_info, option, 0u32)) .transpose() } diff --git a/src/indexer/delete_queue.rs b/src/indexer/delete_queue.rs index 7cc485e1d..29e91b67d 100644 --- a/src/indexer/delete_queue.rs +++ b/src/indexer/delete_queue.rs @@ -249,7 +249,12 @@ mod tests { struct DummyWeight; impl Weight for DummyWeight { - fn scorer(&self, _reader: &SegmentReader, _boost: Score) -> crate::Result> { + fn scorer( + &self, + _reader: &SegmentReader, + _boost: Score, + _seek_doc: DocId, + ) -> crate::Result> { Err(crate::TantivyError::InternalError("dummy impl".to_owned())) } diff --git a/src/indexer/merger.rs b/src/indexer/merger.rs index 47ac5a55b..a5e1a86a6 100644 --- a/src/indexer/merger.rs +++ b/src/indexer/merger.rs @@ -367,8 +367,11 @@ impl IndexMerger { for (segment_ord, term_info) in merged_terms.current_segment_ords_and_term_infos() { let segment_reader = &self.readers[segment_ord]; let inverted_index: &InvertedIndexReader = &field_readers[segment_ord]; - let segment_postings = inverted_index - .read_postings_from_terminfo(&term_info, segment_postings_option)?; + let segment_postings = inverted_index.read_postings_from_terminfo( + &term_info, + segment_postings_option, + 0u32, + )?; let alive_bitset_opt = segment_reader.alive_bitset(); let doc_freq = if let Some(alive_bitset) = alive_bitset_opt { segment_postings.doc_freq_given_deletes(alive_bitset) diff --git a/src/postings/block_segment_postings.rs b/src/postings/block_segment_postings.rs index cc8894294..f29031fbe 100644 --- a/src/postings/block_segment_postings.rs +++ b/src/postings/block_segment_postings.rs @@ -87,6 +87,31 @@ fn split_into_skips_and_postings( Ok((Some(skip_data), postings_data)) } +/// A block segment postings for which the first block has not been loaded yet. +/// +/// You can either call `load_at_start` to load it its first block, +/// or skip a few blocks by calling `seek_and_load`. +pub(crate) struct BlockSegmentPostingsNotLoaded(BlockSegmentPostings); + +impl BlockSegmentPostingsNotLoaded { + /// Seek into the block segment postings directly, possibly avoiding loading its first block. + pub fn seek_and_load(self, seek_doc: DocId) -> (BlockSegmentPostings, usize) { + let BlockSegmentPostingsNotLoaded(mut block_segment_postings) = self; + block_segment_postings.load_block(); + let inner_pos = if seek_doc == 0 { + 0 + } else { + block_segment_postings.seek(seek_doc) + }; + (block_segment_postings, inner_pos) + } + + /// Load the first block of segment postings. + pub fn load_at_start(self) -> BlockSegmentPostings { + self.seek_and_load(0u32).0 + } +} + impl BlockSegmentPostings { /// Opens a `BlockSegmentPostings`. /// `doc_freq` is the number of documents in the posting list. @@ -99,7 +124,7 @@ impl BlockSegmentPostings { data: FileSlice, mut record_option: IndexRecordOption, requested_option: IndexRecordOption, - ) -> io::Result { + ) -> io::Result { let bytes = data.read_bytes()?; let (skip_data_opt, postings_data) = split_into_skips_and_postings(doc_freq, bytes)?; let skip_reader = match skip_data_opt { @@ -125,7 +150,7 @@ impl BlockSegmentPostings { (_, _) => FreqReadingOption::ReadFreq, }; - let mut block_segment_postings = BlockSegmentPostings { + Ok(BlockSegmentPostingsNotLoaded(BlockSegmentPostings { doc_decoder: BlockDecoder::with_val(TERMINATED), block_loaded: false, freq_decoder: BlockDecoder::with_val(1), @@ -134,9 +159,7 @@ impl BlockSegmentPostings { doc_freq, data: postings_data, skip_reader, - }; - block_segment_postings.load_block(); - Ok(block_segment_postings) + })) } /// Returns the block_max_score for the current block. @@ -258,7 +281,9 @@ impl BlockSegmentPostings { self.doc_decoder.output_len } - /// Position on a block that may contains `target_doc`. + /// Position on a block that may contains `target_doc`, and returns the + /// position of the first document greater than or equal to `target_doc` + /// within that block. /// /// If all docs are smaller than target, the block loaded may be empty, /// or be the last an incomplete VInt block. @@ -388,7 +413,7 @@ mod tests { use crate::index::Index; use crate::postings::compression::COMPRESSION_BLOCK_SIZE; use crate::postings::postings::Postings; - use crate::postings::SegmentPostings; + use crate::postings::{BlockSegmentPostingsNotLoaded, SegmentPostings}; use crate::schema::{IndexRecordOption, Schema, Term, INDEXED}; use crate::DocId; @@ -427,7 +452,8 @@ mod tests { #[test] fn test_block_segment_postings() -> crate::Result<()> { - let mut block_segments = build_block_postings(&(0..100_000).collect::>())?; + let mut block_segments = + build_block_postings(&(0..100_000).collect::>())?.load_at_start(); let mut offset: u32 = 0u32; // checking that the `doc_freq` is correct assert_eq!(block_segments.doc_freq(), 100_000); @@ -453,7 +479,7 @@ mod tests { doc_ids.push(130); { let block_segments = build_block_postings(&doc_ids)?; - let mut docset = SegmentPostings::from_block_postings(block_segments, None); + let mut docset = SegmentPostings::from_block_postings(block_segments, None, 0); assert_eq!(docset.seek(128), 129); assert_eq!(docset.doc(), 129); assert_eq!(docset.advance(), 130); @@ -462,7 +488,7 @@ mod tests { } { let block_segments = build_block_postings(&doc_ids).unwrap(); - let mut docset = SegmentPostings::from_block_postings(block_segments, None); + let mut docset = SegmentPostings::from_block_postings(block_segments, None, 0); assert_eq!(docset.seek(129), 129); assert_eq!(docset.doc(), 129); assert_eq!(docset.advance(), 130); @@ -471,7 +497,7 @@ mod tests { } { let block_segments = build_block_postings(&doc_ids)?; - let mut docset = SegmentPostings::from_block_postings(block_segments, None); + let mut docset = SegmentPostings::from_block_postings(block_segments, None, 0); assert_eq!(docset.doc(), 0); assert_eq!(docset.seek(131), TERMINATED); assert_eq!(docset.doc(), TERMINATED); @@ -479,7 +505,7 @@ mod tests { Ok(()) } - fn build_block_postings(docs: &[DocId]) -> crate::Result { + fn build_block_postings(docs: &[DocId]) -> crate::Result { let mut schema_builder = Schema::builder(); let int_field = schema_builder.add_u64_field("id", INDEXED); let schema = schema_builder.build(); @@ -499,9 +525,9 @@ mod tests { let inverted_index = segment_reader.inverted_index(int_field).unwrap(); let term = Term::from_field_u64(int_field, 0u64); let term_info = inverted_index.get_term_info(&term)?.unwrap(); - let block_postings = inverted_index - .read_block_postings_from_terminfo(&term_info, IndexRecordOption::Basic)?; - Ok(block_postings) + let block_postings_not_loaded = inverted_index + .read_block_postings_from_terminfo_not_loaded(&term_info, IndexRecordOption::Basic)?; + Ok(block_postings_not_loaded) } #[test] @@ -510,7 +536,7 @@ mod tests { for i in 0..1300 { docs.push((i * i / 100) + i); } - let mut block_postings = build_block_postings(&docs[..])?; + let mut block_postings = build_block_postings(&docs[..])?.load_at_start(); for i in &[0, 424, 10000] { block_postings.seek(*i); let docs = block_postings.docs(); diff --git a/src/postings/mod.rs b/src/postings/mod.rs index efc0e069d..e5dd49d1f 100644 --- a/src/postings/mod.rs +++ b/src/postings/mod.rs @@ -22,6 +22,7 @@ pub(crate) use loaded_postings::LoadedPostings; pub(crate) use stacker::compute_table_memory_size; pub use self::block_segment_postings::BlockSegmentPostings; +pub(crate) use self::block_segment_postings::BlockSegmentPostingsNotLoaded; pub(crate) use self::indexing_context::IndexingContext; pub(crate) use self::per_field_postings_writer::PerFieldPostingsWriter; pub use self::postings::Postings; diff --git a/src/postings/segment_postings.rs b/src/postings/segment_postings.rs index d9ba33eb2..40eaeab4f 100644 --- a/src/postings/segment_postings.rs +++ b/src/postings/segment_postings.rs @@ -4,7 +4,7 @@ use crate::docset::DocSet; use crate::fastfield::AliveBitSet; use crate::positions::PositionReader; use crate::postings::compression::COMPRESSION_BLOCK_SIZE; -use crate::postings::{BlockSegmentPostings, Postings}; +use crate::postings::{BlockSegmentPostings, BlockSegmentPostingsNotLoaded, Postings}; use crate::{DocId, TERMINATED}; /// `SegmentPostings` represents the inverted list or postings associated with @@ -86,7 +86,7 @@ impl SegmentPostings { IndexRecordOption::Basic, ) .unwrap(); - SegmentPostings::from_block_postings(block_segment_postings, None) + SegmentPostings::from_block_postings(block_segment_postings, None, 0) } /// Helper functions to create `SegmentPostings` for tests. @@ -134,21 +134,22 @@ impl SegmentPostings { IndexRecordOption::WithFreqs, ) .unwrap(); - SegmentPostings::from_block_postings(block_segment_postings, None) + SegmentPostings::from_block_postings(block_segment_postings, None, 0) } - /// Reads a Segment postings from an &[u8] - /// - /// * `len` - number of document in the posting lists. - /// * `data` - data array. The complete data is not necessarily used. - /// * `freq_handler` - the freq handler is in charge of decoding frequencies and/or positions + /// Creates a Segment Postings from a + /// - `BlockSegmentPostings`, + /// - a position reader + /// - a target document to seek to pub(crate) fn from_block_postings( - segment_block_postings: BlockSegmentPostings, + segment_block_postings: BlockSegmentPostingsNotLoaded, position_reader: Option, + seek_doc: DocId, ) -> SegmentPostings { + let (block_cursor, cur) = segment_block_postings.seek_and_load(seek_doc); SegmentPostings { - block_cursor: segment_block_postings, - cur: 0, // cursor within the block + block_cursor, + cur, position_reader, } } diff --git a/src/query/all_query.rs b/src/query/all_query.rs index 612d7408c..2bc342019 100644 --- a/src/query/all_query.rs +++ b/src/query/all_query.rs @@ -21,7 +21,12 @@ impl Query for AllQuery { pub struct AllWeight; impl Weight for AllWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + _seek_doc: DocId, + ) -> crate::Result> { let all_scorer = AllScorer::new(reader.max_doc()); if boost != 1.0 { Ok(Box::new(BoostScorer::new(all_scorer, boost))) @@ -140,7 +145,7 @@ mod tests { let weight = AllQuery.weight(EnableScoring::disabled_from_schema(&index.schema()))?; { let reader = searcher.segment_reader(0); - let mut scorer = weight.scorer(reader, 1.0)?; + let mut scorer = weight.scorer(reader, 1.0, 0)?; assert_eq!(scorer.doc(), 0u32); assert_eq!(scorer.advance(), 1u32); assert_eq!(scorer.doc(), 1u32); @@ -148,7 +153,7 @@ mod tests { } { let reader = searcher.segment_reader(1); - let mut scorer = weight.scorer(reader, 1.0)?; + let mut scorer = weight.scorer(reader, 1.0, 0)?; assert_eq!(scorer.doc(), 0u32); assert_eq!(scorer.advance(), TERMINATED); } @@ -163,12 +168,12 @@ mod tests { let weight = AllQuery.weight(EnableScoring::disabled_from_schema(searcher.schema()))?; let reader = searcher.segment_reader(0); { - let mut scorer = weight.scorer(reader, 2.0)?; + let mut scorer = weight.scorer(reader, 2.0, 0)?; assert_eq!(scorer.doc(), 0u32); assert_eq!(scorer.score(), 2.0); } { - let mut scorer = weight.scorer(reader, 1.5)?; + let mut scorer = weight.scorer(reader, 1.5, 0)?; assert_eq!(scorer.doc(), 0u32); assert_eq!(scorer.score(), 1.5); } diff --git a/src/query/automaton_weight.rs b/src/query/automaton_weight.rs index 5f1053fb6..1763abe78 100644 --- a/src/query/automaton_weight.rs +++ b/src/query/automaton_weight.rs @@ -84,7 +84,12 @@ where A: Automaton + Send + Sync + 'static, A::State: Clone, { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { let max_doc = reader.max_doc(); let mut doc_bitset = BitSet::with_max_value(max_doc); let inverted_index = reader.inverted_index(self.field)?; @@ -93,7 +98,9 @@ where while term_stream.advance() { let term_info = term_stream.value(); let mut block_segment_postings = inverted_index - .read_block_postings_from_terminfo(term_info, IndexRecordOption::Basic)?; + .read_block_postings_from_terminfo_not_loaded(term_info, IndexRecordOption::Basic)? + .seek_and_load(seek_doc) + .0; loop { let docs = block_segment_postings.docs(); if docs.is_empty() { @@ -111,7 +118,7 @@ where } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if scorer.seek(doc) == doc { Ok(Explanation::new("AutomatonScorer", 1.0)) } else { @@ -186,7 +193,7 @@ mod tests { let automaton_weight = AutomatonWeight::new(field, PrefixedByA); let reader = index.reader()?; let searcher = reader.searcher(); - let mut scorer = automaton_weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let mut scorer = automaton_weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert_eq!(scorer.doc(), 0u32); assert_eq!(scorer.score(), 1.0); assert_eq!(scorer.advance(), 2u32); @@ -203,7 +210,7 @@ mod tests { let automaton_weight = AutomatonWeight::new(field, PrefixedByA); let reader = index.reader()?; let searcher = reader.searcher(); - let mut scorer = automaton_weight.scorer(searcher.segment_reader(0u32), 1.32)?; + let mut scorer = automaton_weight.scorer(searcher.segment_reader(0u32), 1.32, 0)?; assert_eq!(scorer.doc(), 0u32); assert_eq!(scorer.score(), 1.32); Ok(()) diff --git a/src/query/boolean_query/boolean_weight.rs b/src/query/boolean_query/boolean_weight.rs index c46e9b0b1..d95eacc29 100644 --- a/src/query/boolean_query/boolean_weight.rs +++ b/src/query/boolean_query/boolean_weight.rs @@ -12,7 +12,7 @@ use crate::query::{ intersect_scorers, AllScorer, BufferedUnionScorer, EmptyScorer, Exclude, Explanation, Occur, RequiredOptionalScorer, Scorer, Weight, }; -use crate::{DocId, Score}; +use crate::{DocId, Score, TERMINATED}; enum SpecializedScorer { TermUnion(Vec), @@ -156,6 +156,19 @@ fn effective_should_scorer_for_union( } } +fn create_scorer( + weight: &dyn Weight, + reader: &SegmentReader, + boost: Score, + target_doc: DocId, +) -> crate::Result> { + if target_doc >= reader.max_doc() { + Ok(Box::new(EmptyScorer)) + } else { + weight.scorer(reader, boost, target_doc) + } +} + enum ShouldScorersCombinationMethod { // Should scorers are irrelevant. Ignored, @@ -207,10 +220,29 @@ impl BooleanWeight { &self, reader: &SegmentReader, boost: Score, + mut seek_first_doc: DocId, ) -> crate::Result>>> { let mut per_occur_scorers: HashMap>> = HashMap::new(); - for (occur, subweight) in &self.weights { - let sub_scorer: Box = subweight.scorer(reader, boost)?; + let (mut must_weights, other_weights): (Vec<(Occur, _)>, Vec<(Occur, _)>) = self + .weights + .iter() + .map(|(occur, weight)| (*occur, weight)) + .partition(|(occur, _weight)| *occur == Occur::Must); + // We start by must weights in order to get the best "seek_first_doc" so that we + // can skip the first few documents of the other scorers. + must_weights.sort_by_key(|weight| weight.1.intersection_priority()); + for (_, must_sub_weight) in must_weights { + let sub_scorer: Box = + create_scorer(must_sub_weight.as_ref(), reader, boost, seek_first_doc)?; + seek_first_doc = seek_first_doc.max(sub_scorer.doc()); + per_occur_scorers + .entry(Occur::Must) + .or_default() + .push(sub_scorer); + } + for (occur, sub_weight) in &other_weights { + let sub_scorer: Box = + create_scorer(sub_weight.as_ref(), reader, boost, seek_first_doc)?; per_occur_scorers .entry(*occur) .or_default() @@ -224,9 +256,10 @@ impl BooleanWeight { reader: &SegmentReader, boost: Score, score_combiner_fn: impl Fn() -> TComplexScoreCombiner, + seek_doc: u32, ) -> crate::Result { let num_docs = reader.num_docs(); - let mut per_occur_scorers = self.per_occur_scorers(reader, boost)?; + let mut per_occur_scorers = self.per_occur_scorers(reader, boost, seek_doc)?; // Indicate how should clauses are combined with must clauses. let mut must_scorers: Vec> = @@ -407,7 +440,7 @@ fn remove_and_count_all_and_empty_scorers( if scorer.is::() { counts.num_all_scorers += 1; false - } else if scorer.is::() { + } else if scorer.doc() == TERMINATED { counts.num_empty_scorers += 1; false } else { @@ -418,7 +451,12 @@ fn remove_and_count_all_and_empty_scorers( } impl Weight for BooleanWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { let num_docs = reader.num_docs(); if self.weights.is_empty() { Ok(Box::new(EmptyScorer)) @@ -427,15 +465,15 @@ impl Weight for BooleanWeight Weight for BooleanWeight crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if scorer.seek(doc) != doc { return Err(does_not_match(doc)); } @@ -467,7 +505,7 @@ impl Weight for BooleanWeight crate::Result<()> { - let scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn)?; + let scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn, 0)?; match scorer { SpecializedScorer::TermUnion(term_scorers) => { let mut union_scorer = BufferedUnionScorer::build( @@ -489,7 +527,7 @@ impl Weight for BooleanWeight crate::Result<()> { - let scorer = self.complex_scorer(reader, 1.0, || DoNothingCombiner)?; + let scorer = self.complex_scorer(reader, 1.0, || DoNothingCombiner, 0u32)?; let mut buffer = [0u32; COLLECT_BLOCK_BUFFER_LEN]; match scorer { @@ -524,7 +562,7 @@ impl Weight for BooleanWeight Score, ) -> crate::Result<()> { - let scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn)?; + let scorer = self.complex_scorer(reader, 1.0, &self.score_combiner_fn, 0u32)?; match scorer { SpecializedScorer::TermUnion(term_scorers) => { super::block_wand(term_scorers, threshold, callback); diff --git a/src/query/boolean_query/mod.rs b/src/query/boolean_query/mod.rs index 681881c11..fd57374d2 100644 --- a/src/query/boolean_query/mod.rs +++ b/src/query/boolean_query/mod.rs @@ -57,7 +57,7 @@ mod tests { let query = query_parser.parse_query("+a")?; let searcher = index.reader()?.searcher(); let weight = query.weight(EnableScoring::enabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert!(scorer.is::()); Ok(()) } @@ -70,13 +70,13 @@ mod tests { { let query = query_parser.parse_query("+a +b +c")?; let weight = query.weight(EnableScoring::enabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert!(scorer.is::>()); } { let query = query_parser.parse_query("+a +(b c)")?; let weight = query.weight(EnableScoring::enabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert!(scorer.is::>>()); } Ok(()) @@ -90,14 +90,14 @@ mod tests { { let query = query_parser.parse_query("+a b")?; let weight = query.weight(EnableScoring::enabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert!(scorer .is::, Box, SumCombiner>>()); } { let query = query_parser.parse_query("+a b")?; let weight = query.weight(EnableScoring::disabled_from_schema(searcher.schema()))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert!(scorer.is::()); } Ok(()) @@ -244,12 +244,14 @@ mod tests { .weight(EnableScoring::enabled_from_searcher(&searcher)) .unwrap(); { - let mut boolean_scorer = boolean_weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let mut boolean_scorer = + boolean_weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert_eq!(boolean_scorer.doc(), 0u32); assert_nearly_equals!(boolean_scorer.score(), 0.84163445); } { - let mut boolean_scorer = boolean_weight.scorer(searcher.segment_reader(0u32), 2.0)?; + let mut boolean_scorer = + boolean_weight.scorer(searcher.segment_reader(0u32), 2.0, 0)?; assert_eq!(boolean_scorer.doc(), 0u32); assert_nearly_equals!(boolean_scorer.score(), 1.6832689); } @@ -343,7 +345,7 @@ mod tests { (Occur::Must, term_match_some.box_clone()), ]); let weight = query.weight(EnableScoring::disabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32, 0)?; assert!(scorer.is::()); } { @@ -353,7 +355,7 @@ mod tests { (Occur::Must, term_match_none.box_clone()), ]); let weight = query.weight(EnableScoring::disabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32, 0)?; assert!(scorer.is::()); } { @@ -362,7 +364,7 @@ mod tests { (Occur::Should, term_match_none.box_clone()), ]); let weight = query.weight(EnableScoring::disabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32, 0)?; assert!(scorer.is::()); } { @@ -371,7 +373,7 @@ mod tests { (Occur::Should, term_match_none.box_clone()), ]); let weight = query.weight(EnableScoring::disabled_from_searcher(&searcher))?; - let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32)?; + let scorer = weight.scorer(searcher.segment_reader(0u32), 1.0f32, 0)?; assert!(scorer.is::()); } Ok(()) diff --git a/src/query/boost_query.rs b/src/query/boost_query.rs index ecbf3d8d6..0c06a829b 100644 --- a/src/query/boost_query.rs +++ b/src/query/boost_query.rs @@ -67,8 +67,13 @@ impl BoostWeight { } impl Weight for BoostWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { - self.weight.scorer(reader, boost * self.boost) + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { + self.weight.scorer(reader, boost * self.boost, seek_doc) } fn explain(&self, reader: &SegmentReader, doc: u32) -> crate::Result { @@ -83,6 +88,10 @@ impl Weight for BoostWeight { fn count(&self, reader: &SegmentReader) -> crate::Result { self.weight.count(reader) } + + fn intersection_priority(&self) -> u32 { + self.weight.intersection_priority() + } } pub(crate) struct BoostScorer { diff --git a/src/query/const_score_query.rs b/src/query/const_score_query.rs index 570c7feca..6202b7043 100644 --- a/src/query/const_score_query.rs +++ b/src/query/const_score_query.rs @@ -63,13 +63,18 @@ impl ConstWeight { } impl Weight for ConstWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { - let inner_scorer = self.weight.scorer(reader, boost)?; + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { + let inner_scorer = self.weight.scorer(reader, boost, seek_doc)?; Ok(Box::new(ConstScorer::new(inner_scorer, boost * self.score))) } fn explain(&self, reader: &SegmentReader, doc: u32) -> crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if scorer.seek(doc) != doc { return Err(TantivyError::InvalidArgument(format!( "Document #({doc}) does not match" @@ -84,6 +89,10 @@ impl Weight for ConstWeight { fn count(&self, reader: &SegmentReader) -> crate::Result { self.weight.count(reader) } + + fn intersection_priority(&self) -> u32 { + self.weight.intersection_priority() + } } /// Wraps a `DocSet` and simply returns a constant `Scorer`. diff --git a/src/query/empty_query.rs b/src/query/empty_query.rs index 86ff84c08..d46878081 100644 --- a/src/query/empty_query.rs +++ b/src/query/empty_query.rs @@ -26,13 +26,24 @@ impl Query for EmptyQuery { /// It is useful for tests and handling edge cases. pub struct EmptyWeight; impl Weight for EmptyWeight { - fn scorer(&self, _reader: &SegmentReader, _boost: Score) -> crate::Result> { + fn scorer( + &self, + _reader: &SegmentReader, + _boost: Score, + _seek_doc: DocId, + ) -> crate::Result> { Ok(Box::new(EmptyScorer)) } fn explain(&self, _reader: &SegmentReader, doc: DocId) -> crate::Result { Err(does_not_match(doc)) } + + /// Returns a priority number used to sort weights when running an + /// intersection. + fn intersection_priority(&self) -> u32 { + 0u32 + } } /// `EmptyScorer` is a dummy `Scorer` in which no document matches. diff --git a/src/query/exist_query.rs b/src/query/exist_query.rs index 7eb09722c..645c9a054 100644 --- a/src/query/exist_query.rs +++ b/src/query/exist_query.rs @@ -98,7 +98,12 @@ pub struct ExistsWeight { } impl Weight for ExistsWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + _seek_doc: DocId, + ) -> crate::Result> { let fast_field_reader = reader.fast_fields(); let mut column_handles = fast_field_reader.dynamic_column_handles(&self.field_name)?; if self.field_type == Type::Json && self.json_subpaths { @@ -166,7 +171,7 @@ impl Weight for ExistsWeight { } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if scorer.seek(doc) != doc { return Err(does_not_match(doc)); } diff --git a/src/query/phrase_prefix_query/phrase_prefix_weight.rs b/src/query/phrase_prefix_query/phrase_prefix_weight.rs index 546eb89e8..8c2de4c59 100644 --- a/src/query/phrase_prefix_query/phrase_prefix_weight.rs +++ b/src/query/phrase_prefix_query/phrase_prefix_weight.rs @@ -42,10 +42,11 @@ impl PhrasePrefixWeight { Ok(FieldNormReader::constant(reader.max_doc(), 1)) } - pub(crate) fn phrase_scorer( + pub(crate) fn prefix_phrase_scorer( &self, reader: &SegmentReader, boost: Score, + seek_doc: DocId, ) -> crate::Result>> { let similarity_weight_opt = self .similarity_weight_opt @@ -54,14 +55,16 @@ impl PhrasePrefixWeight { let fieldnorm_reader = self.fieldnorm_reader(reader)?; let mut term_postings_list = Vec::new(); for &(offset, ref term) in &self.phrase_terms { - if let Some(postings) = reader - .inverted_index(term.field())? - .read_postings(term, IndexRecordOption::WithFreqsAndPositions)? - { - term_postings_list.push((offset, postings)); - } else { + let inverted_index = reader.inverted_index(term.field())?; + let Some(term_info) = inverted_index.get_term_info(term)? else { return Ok(None); - } + }; + let postings = inverted_index.read_postings_from_terminfo( + &term_info, + IndexRecordOption::WithFreqsAndPositions, + seek_doc, + )?; + term_postings_list.push((offset, postings)); } let inv_index = reader.inverted_index(self.prefix.1.field())?; @@ -114,8 +117,13 @@ impl PhrasePrefixWeight { } impl Weight for PhrasePrefixWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { - if let Some(scorer) = self.phrase_scorer(reader, boost)? { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { + if let Some(scorer) = self.prefix_phrase_scorer(reader, boost, seek_doc)? { Ok(Box::new(scorer)) } else { Ok(Box::new(EmptyScorer)) @@ -123,7 +131,7 @@ impl Weight for PhrasePrefixWeight { } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - let scorer_opt = self.phrase_scorer(reader, 1.0)?; + let scorer_opt = self.prefix_phrase_scorer(reader, 1.0, doc)?; if scorer_opt.is_none() { return Err(does_not_match(doc)); } @@ -140,6 +148,10 @@ impl Weight for PhrasePrefixWeight { } Ok(explanation) } + + fn intersection_priority(&self) -> u32 { + 50u32 + } } #[cfg(test)] @@ -187,7 +199,7 @@ mod tests { .unwrap() .unwrap(); let mut phrase_scorer = phrase_weight - .phrase_scorer(searcher.segment_reader(0u32), 1.0)? + .prefix_phrase_scorer(searcher.segment_reader(0u32), 1.0, 0u32)? .unwrap(); assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.phrase_count(), 2); @@ -214,7 +226,7 @@ mod tests { .unwrap() .unwrap(); let mut phrase_scorer = phrase_weight - .phrase_scorer(searcher.segment_reader(0u32), 1.0)? + .prefix_phrase_scorer(searcher.segment_reader(0u32), 1.0, 0u32)? .unwrap(); assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.phrase_count(), 2); @@ -238,7 +250,7 @@ mod tests { .unwrap() .is_none()); let weight = phrase_query.weight(enable_scoring).unwrap(); - let mut phrase_scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let mut phrase_scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.advance(), 2); assert_eq!(phrase_scorer.doc(), 2); @@ -259,7 +271,7 @@ mod tests { ]); let enable_scoring = EnableScoring::enabled_from_searcher(&searcher); let weight = phrase_query.weight(enable_scoring).unwrap(); - let mut phrase_scorer = weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let mut phrase_scorer = weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert_eq!(phrase_scorer.advance(), TERMINATED); Ok(()) } diff --git a/src/query/phrase_query/mod.rs b/src/query/phrase_query/mod.rs index 938e34442..a230afc5d 100644 --- a/src/query/phrase_query/mod.rs +++ b/src/query/phrase_query/mod.rs @@ -84,7 +84,7 @@ pub(crate) mod tests { let phrase_query = PhraseQuery::new(terms); let phrase_weight = phrase_query.phrase_weight(EnableScoring::disabled_from_schema(searcher.schema()))?; - let mut phrase_scorer = phrase_weight.scorer(searcher.segment_reader(0), 1.0)?; + let mut phrase_scorer = phrase_weight.scorer(searcher.segment_reader(0), 1.0, 0)?; assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.advance(), TERMINATED); Ok(()) @@ -343,6 +343,43 @@ pub(crate) mod tests { Ok(()) } + #[test] + pub fn test_phrase_weight_seek_doc() -> crate::Result<()> { + // Create an index with documents where the phrase "a b" appears in some of them. + // Documents: 0: "c d", 1: "a b", 2: "e f", 3: "a b c", 4: "g h", 5: "a b", 6: "i j" + let index = create_index(&["c d", "a b", "e f", "a b c", "g h", "a b", "i j"])?; + let text_field = index.schema().get_field("text").unwrap(); + let searcher = index.reader()?.searcher(); + let segment_reader = searcher.segment_reader(0); + + let phrase_query = PhraseQuery::new(vec![ + Term::from_field_text(text_field, "a"), + Term::from_field_text(text_field, "b"), + ]); + let phrase_weight = + phrase_query.phrase_weight(EnableScoring::disabled_from_schema(searcher.schema()))?; + + // Helper function to collect all docs from a scorer created with a given seek_doc + let docs_when_seeking_from = |seek_from: DocId| { + let scorer = phrase_weight + .scorer(segment_reader, 1.0f32, seek_from) + .unwrap(); + crate::docset::docset_to_doc_vec(scorer) + }; + + // Documents with "a b": 1, 3, 5 + assert_eq!(docs_when_seeking_from(0), vec![1, 3, 5]); + assert_eq!(docs_when_seeking_from(1), vec![1, 3, 5]); + assert_eq!(docs_when_seeking_from(2), vec![3, 5]); + assert_eq!(docs_when_seeking_from(3), vec![3, 5]); + assert_eq!(docs_when_seeking_from(4), vec![5]); + assert_eq!(docs_when_seeking_from(5), vec![5]); + assert_eq!(docs_when_seeking_from(6), Vec::::new()); + assert_eq!(docs_when_seeking_from(7), Vec::::new()); + + Ok(()) + } + #[test] pub fn test_phrase_query_on_json() -> crate::Result<()> { let mut schema_builder = Schema::builder(); @@ -373,7 +410,7 @@ pub(crate) mod tests { .weight(EnableScoring::disabled_from_schema(searcher.schema())) .unwrap(); let mut phrase_scorer = phrase_weight - .scorer(searcher.segment_reader(0), 1.0f32) + .scorer(searcher.segment_reader(0), 1.0f32, 0) .unwrap(); let mut docs = Vec::new(); loop { diff --git a/src/query/phrase_query/phrase_weight.rs b/src/query/phrase_query/phrase_weight.rs index 4118f79f6..1d0ef9413 100644 --- a/src/query/phrase_query/phrase_weight.rs +++ b/src/query/phrase_query/phrase_weight.rs @@ -43,6 +43,7 @@ impl PhraseWeight { &self, reader: &SegmentReader, boost: Score, + seek_doc: DocId, ) -> crate::Result>> { let similarity_weight_opt = self .similarity_weight_opt @@ -51,14 +52,16 @@ impl PhraseWeight { let fieldnorm_reader = self.fieldnorm_reader(reader)?; let mut term_postings_list = Vec::new(); for &(offset, ref term) in &self.phrase_terms { - if let Some(postings) = reader - .inverted_index(term.field())? - .read_postings(term, IndexRecordOption::WithFreqsAndPositions)? - { - term_postings_list.push((offset, postings)); - } else { + let inverted_index = reader.inverted_index(term.field())?; + let Some(term_info) = inverted_index.get_term_info(term)? else { return Ok(None); - } + }; + let postings = inverted_index.read_postings_from_terminfo( + &term_info, + IndexRecordOption::WithFreqsAndPositions, + seek_doc, + )?; + term_postings_list.push((offset, postings)); } Ok(Some(PhraseScorer::new( term_postings_list, @@ -74,8 +77,13 @@ impl PhraseWeight { } impl Weight for PhraseWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { - if let Some(scorer) = self.phrase_scorer(reader, boost)? { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { + if let Some(scorer) = self.phrase_scorer(reader, boost, seek_doc)? { Ok(Box::new(scorer)) } else { Ok(Box::new(EmptyScorer)) @@ -83,12 +91,12 @@ impl Weight for PhraseWeight { } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - let scorer_opt = self.phrase_scorer(reader, 1.0)?; + let scorer_opt = self.phrase_scorer(reader, 1.0, doc)?; if scorer_opt.is_none() { return Err(does_not_match(doc)); } let mut scorer = scorer_opt.unwrap(); - if scorer.seek(doc) != doc { + if scorer.doc() != doc { return Err(does_not_match(doc)); } let fieldnorm_reader = self.fieldnorm_reader(reader)?; @@ -100,6 +108,10 @@ impl Weight for PhraseWeight { } Ok(explanation) } + + fn intersection_priority(&self) -> u32 { + 40u32 + } } #[cfg(test)] @@ -122,7 +134,7 @@ mod tests { let enable_scoring = EnableScoring::enabled_from_searcher(&searcher); let phrase_weight = phrase_query.phrase_weight(enable_scoring).unwrap(); let mut phrase_scorer = phrase_weight - .phrase_scorer(searcher.segment_reader(0u32), 1.0)? + .phrase_scorer(searcher.segment_reader(0u32), 1.0, 0)? .unwrap(); assert_eq!(phrase_scorer.doc(), 1); assert_eq!(phrase_scorer.phrase_count(), 2); diff --git a/src/query/phrase_query/regex_phrase_weight.rs b/src/query/phrase_query/regex_phrase_weight.rs index 4e850d2e2..93119a84b 100644 --- a/src/query/phrase_query/regex_phrase_weight.rs +++ b/src/query/phrase_query/regex_phrase_weight.rs @@ -195,8 +195,11 @@ impl RegexPhraseWeight { const SPARSE_TERM_DOC_THRESHOLD: u32 = 100; for term_info in term_infos { - let mut term_posting = inverted_index - .read_postings_from_terminfo(term_info, IndexRecordOption::WithFreqsAndPositions)?; + let mut term_posting = inverted_index.read_postings_from_terminfo( + term_info, + IndexRecordOption::WithFreqsAndPositions, + 0u32, + )?; let num_docs = term_posting.doc_freq(); if num_docs < SPARSE_TERM_DOC_THRESHOLD { @@ -269,7 +272,12 @@ impl RegexPhraseWeight { } impl Weight for RegexPhraseWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + _seek_doc: DocId, + ) -> crate::Result> { if let Some(scorer) = self.phrase_scorer(reader, boost)? { Ok(Box::new(scorer)) } else { diff --git a/src/query/range_query/fast_field_range_doc_set.rs b/src/query/range_query/fast_field_range_doc_set.rs index 24d2b1fe3..32e054476 100644 --- a/src/query/range_query/fast_field_range_doc_set.rs +++ b/src/query/range_query/fast_field_range_doc_set.rs @@ -61,7 +61,11 @@ pub(crate) struct RangeDocSet { const DEFAULT_FETCH_HORIZON: u32 = 128; impl RangeDocSet { - pub(crate) fn new(value_range: RangeInclusive, column: Column) -> Self { + pub(crate) fn new( + value_range: RangeInclusive, + column: Column, + seek_first_doc: DocId, + ) -> Self { if *value_range.start() > column.max_value() || *value_range.end() < column.min_value() { return Self { value_range, @@ -77,7 +81,7 @@ impl RangeDocSet { value_range, column, loaded_docs: VecCursor::new(), - next_fetch_start: 0, + next_fetch_start: seek_first_doc, fetch_horizon: DEFAULT_FETCH_HORIZON, last_seek_pos_opt: None, }; diff --git a/src/query/range_query/range_query.rs b/src/query/range_query/range_query.rs index 1893a06a5..ba80c8398 100644 --- a/src/query/range_query/range_query.rs +++ b/src/query/range_query/range_query.rs @@ -212,7 +212,12 @@ impl InvertedIndexRangeWeight { } impl Weight for InvertedIndexRangeWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { let max_doc = reader.max_doc(); let mut doc_bitset = BitSet::with_max_value(max_doc); @@ -229,7 +234,9 @@ impl Weight for InvertedIndexRangeWeight { processed_count += 1; let term_info = term_range.value(); let mut block_segment_postings = inverted_index - .read_block_postings_from_terminfo(term_info, IndexRecordOption::Basic)?; + .read_block_postings_from_terminfo_not_loaded(term_info, IndexRecordOption::Basic)? + .seek_and_load(seek_doc) + .0; loop { let docs = block_segment_postings.docs(); if docs.is_empty() { @@ -246,7 +253,7 @@ impl Weight for InvertedIndexRangeWeight { } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if scorer.seek(doc) != doc { return Err(does_not_match(doc)); } @@ -686,7 +693,7 @@ mod tests { .weight(EnableScoring::disabled_from_schema(&schema)) .unwrap(); let range_scorer = range_weight - .scorer(&searcher.segment_readers()[0], 1.0f32) + .scorer(&searcher.segment_readers()[0], 1.0f32, 0) .unwrap(); range_scorer }; diff --git a/src/query/range_query/range_query_fastfield.rs b/src/query/range_query/range_query_fastfield.rs index e379e108e..dad0d66fc 100644 --- a/src/query/range_query/range_query_fastfield.rs +++ b/src/query/range_query/range_query_fastfield.rs @@ -52,7 +52,12 @@ impl FastFieldRangeWeight { } impl Weight for FastFieldRangeWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { // Check if both bounds are Bound::Unbounded if self.bounds.is_unbounded() { return Ok(Box::new(AllScorer::new(reader.max_doc()))); @@ -109,11 +114,21 @@ impl Weight for FastFieldRangeWeight { else { return Ok(Box::new(EmptyScorer)); }; - search_on_u64_ff(column, boost, BoundsRange::new(lower_bound, upper_bound)) - } - Type::U64 | Type::I64 | Type::F64 => { - search_on_json_numerical_field(reader, &field_name, typ, bounds, boost) + search_on_u64_ff( + column, + boost, + BoundsRange::new(lower_bound, upper_bound), + seek_doc, + ) } + Type::U64 | Type::I64 | Type::F64 => search_on_json_numerical_field( + reader, + &field_name, + typ, + bounds, + boost, + seek_doc, + ), Type::Date => { let fast_field_reader = reader.fast_fields(); let Some((column, _col_type)) = fast_field_reader @@ -126,6 +141,7 @@ impl Weight for FastFieldRangeWeight { column, boost, BoundsRange::new(bounds.lower_bound, bounds.upper_bound), + seek_doc, ) } Type::Bool | Type::Facet | Type::Bytes | Type::Json | Type::IpAddr => { @@ -154,7 +170,7 @@ impl Weight for FastFieldRangeWeight { ip_addr_column.min_value(), ip_addr_column.max_value(), ); - let docset = RangeDocSet::new(value_range, ip_addr_column); + let docset = RangeDocSet::new(value_range, ip_addr_column, seek_doc); Ok(Box::new(ConstScorer::new(docset, boost))) } else if field_type.is_str() { let Some(str_dict_column): Option = reader.fast_fields().str(&field_name)? @@ -173,7 +189,12 @@ impl Weight for FastFieldRangeWeight { else { return Ok(Box::new(EmptyScorer)); }; - search_on_u64_ff(column, boost, BoundsRange::new(lower_bound, upper_bound)) + search_on_u64_ff( + column, + boost, + BoundsRange::new(lower_bound, upper_bound), + seek_doc, + ) } else { assert!( maps_to_u64_fastfield(field_type.value_type()), @@ -215,12 +236,13 @@ impl Weight for FastFieldRangeWeight { column, boost, BoundsRange::new(bounds.lower_bound, bounds.upper_bound), + seek_doc, ) } } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if scorer.seek(doc) != doc { return Err(TantivyError::InvalidArgument(format!( "Document #({doc}) does not match" @@ -230,6 +252,10 @@ impl Weight for FastFieldRangeWeight { Ok(explanation) } + + fn intersection_priority(&self) -> u32 { + 30u32 + } } /// On numerical fields the column type may not match the user provided one. @@ -241,6 +267,7 @@ fn search_on_json_numerical_field( typ: Type, bounds: BoundsRange>>, boost: Score, + seek_doc: DocId, ) -> crate::Result> { // Since we don't know which type was interpolated for the internal column we // have to check for all numeric types (only one exists) @@ -318,6 +345,7 @@ fn search_on_json_numerical_field( column, boost, BoundsRange::new(bounds.lower_bound, bounds.upper_bound), + seek_doc, ) } @@ -396,6 +424,7 @@ fn search_on_u64_ff( column: Column, boost: Score, bounds: BoundsRange, + seek_doc: DocId, ) -> crate::Result> { let col_min_value = column.min_value(); let col_max_value = column.max_value(); @@ -426,8 +455,8 @@ fn search_on_u64_ff( } } - let docset = RangeDocSet::new(value_range, column); - Ok(Box::new(ConstScorer::new(docset, boost))) + let doc_set = RangeDocSet::new(value_range, column, seek_doc); + Ok(Box::new(ConstScorer::new(doc_set, boost))) } /// Returns true if the type maps to a u64 fast field @@ -504,7 +533,7 @@ mod tests { DateOptions, Field, NumericOptions, Schema, SchemaBuilder, FAST, INDEXED, STORED, STRING, TEXT, }; - use crate::{Index, IndexWriter, TantivyDocument, Term, TERMINATED}; + use crate::{DocId, Index, IndexWriter, TantivyDocument, Term, TERMINATED}; #[test] fn test_text_field_ff_range_query() -> crate::Result<()> { @@ -1142,11 +1171,52 @@ mod tests { Bound::Included(Term::from_field_u64(field, 50_002)), )); let scorer = range_query - .scorer(searcher.segment_reader(0), 1.0f32) + .scorer(searcher.segment_reader(0), 1.0f32, 0) .unwrap(); assert_eq!(scorer.doc(), TERMINATED); } + #[test] + fn test_fastfield_range_weight_seek_doc() { + let mut schema_builder = SchemaBuilder::new(); + let field = schema_builder.add_u64_field("value", FAST); + let schema = schema_builder.build(); + let index = Index::create_in_ram(schema); + let mut writer: IndexWriter = index.writer_for_tests().unwrap(); + + // Create 20 documents with values + // 0, 10, 20, ..., 90 + // and then 50 again. + for i in 0..10 { + writer.add_document(doc!(field => (i * 10) as u64)).unwrap(); + } + writer.add_document(doc!(field => 50u64)).unwrap(); + writer.commit().unwrap(); + + let searcher = index.reader().unwrap().searcher(); + let segment_reader = searcher.segment_reader(0); + + let range_weight = FastFieldRangeWeight::new(BoundsRange::new( + Bound::Included(Term::from_field_u64(field, 30)), + Bound::Included(Term::from_field_u64(field, 70)), + )); + + let doc_when_seeking_from = |seek_from: DocId| { + let doc_set = range_weight + .scorer(segment_reader, 1.0f32, seek_from) + .unwrap(); + crate::docset::docset_to_doc_vec(doc_set) + }; + + assert_eq!(doc_when_seeking_from(0), vec![3, 4, 5, 6, 7, 10]); + assert_eq!(doc_when_seeking_from(1), vec![3, 4, 5, 6, 7, 10]); + assert_eq!(doc_when_seeking_from(3), vec![3, 4, 5, 6, 7, 10]); + assert_eq!(doc_when_seeking_from(7), vec![7, 10]); + assert_eq!(doc_when_seeking_from(8), vec![10]); + assert_eq!(doc_when_seeking_from(10), vec![10]); + assert_eq!(doc_when_seeking_from(11), Vec::::new()); + } + #[test] fn range_regression3_test() { let ops = vec![doc_from_id_1(1), doc_from_id_1(2), doc_from_id_1(3)]; diff --git a/src/query/term_query/mod.rs b/src/query/term_query/mod.rs index 0811725be..65d338c2f 100644 --- a/src/query/term_query/mod.rs +++ b/src/query/term_query/mod.rs @@ -37,7 +37,7 @@ mod tests { ); let term_weight = term_query.weight(EnableScoring::enabled_from_searcher(&searcher))?; let segment_reader = searcher.segment_reader(0); - let mut term_scorer = term_weight.scorer(segment_reader, 1.0)?; + let mut term_scorer = term_weight.scorer(segment_reader, 1.0, 0)?; assert_eq!(term_scorer.doc(), 0); assert_nearly_equals!(term_scorer.score(), 0.28768212); Ok(()) @@ -65,7 +65,7 @@ mod tests { ); let term_weight = term_query.weight(EnableScoring::enabled_from_searcher(&searcher))?; let segment_reader = searcher.segment_reader(0); - let mut term_scorer = term_weight.scorer(segment_reader, 1.0)?; + let mut term_scorer = term_weight.scorer(segment_reader, 1.0, 0)?; for i in 0u32..COMPRESSION_BLOCK_SIZE as u32 { assert_eq!(term_scorer.doc(), i); if i == COMPRESSION_BLOCK_SIZE as u32 - 1u32 { @@ -162,7 +162,7 @@ mod tests { let searcher = index.reader()?.searcher(); let term_weight = term_query.weight(EnableScoring::disabled_from_schema(searcher.schema()))?; - let mut term_scorer = term_weight.scorer(searcher.segment_reader(0u32), 1.0)?; + let mut term_scorer = term_weight.scorer(searcher.segment_reader(0u32), 1.0, 0)?; assert_eq!(term_scorer.doc(), 0u32); term_scorer.seek(1u32); assert_eq!(term_scorer.doc(), 1u32); @@ -470,7 +470,7 @@ mod tests { .weight(EnableScoring::disabled_from_schema(&schema)) .unwrap(); term_weight - .scorer(searcher.segment_reader(0u32), 1.0f32) + .scorer(searcher.segment_reader(0u32), 1.0f32, 0) .unwrap() }; // Should be an allscorer @@ -484,6 +484,53 @@ mod tests { assert!(empty_scorer.is::()); } + #[test] + fn test_term_weight_seek_doc() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let text_field = schema_builder.add_text_field("text", TEXT); + let schema = schema_builder.build(); + let index = Index::create_in_ram(schema); + let mut index_writer: IndexWriter = index.writer_for_tests()?; + + // Create 11 documents where docs 3, 4, 5, 6, 7, and 10 contain "target" + // (similar pattern to test_fastfield_range_weight_seek_doc) + for i in 0..11 { + if i == 3 || i == 4 || i == 5 || i == 6 || i == 7 || i == 10 { + index_writer.add_document(doc!(text_field => "target"))?; + } else { + index_writer.add_document(doc!(text_field => "other"))?; + } + } + index_writer.commit()?; + + let searcher = index.reader()?.searcher(); + let segment_reader = searcher.segment_reader(0); + + let term_query = TermQuery::new( + Term::from_field_text(text_field, "target"), + IndexRecordOption::Basic, + ); + let term_weight = + term_query.weight(EnableScoring::disabled_from_schema(searcher.schema()))?; + + let doc_when_seeking_from = |seek_from: crate::DocId| { + let scorer = term_weight + .scorer(segment_reader, 1.0f32, seek_from) + .unwrap(); + crate::docset::docset_to_doc_vec(scorer) + }; + + assert_eq!(doc_when_seeking_from(0), vec![3, 4, 5, 6, 7, 10]); + assert_eq!(doc_when_seeking_from(1), vec![3, 4, 5, 6, 7, 10]); + assert_eq!(doc_when_seeking_from(3), vec![3, 4, 5, 6, 7, 10]); + assert_eq!(doc_when_seeking_from(7), vec![7, 10]); + assert_eq!(doc_when_seeking_from(8), vec![10]); + assert_eq!(doc_when_seeking_from(10), vec![10]); + assert_eq!(doc_when_seeking_from(11), Vec::::new()); + + Ok(()) + } + #[test] fn test_term_weight_all_query_optimization_disable_when_scoring_enabled() { let mut schema_builder = Schema::builder(); @@ -509,7 +556,7 @@ mod tests { .weight(EnableScoring::enabled_from_searcher(&searcher)) .unwrap(); term_weight - .scorer(searcher.segment_reader(0u32), 1.0f32) + .scorer(searcher.segment_reader(0u32), 1.0f32, 0) .unwrap() }; // Should be an allscorer diff --git a/src/query/term_query/term_weight.rs b/src/query/term_query/term_weight.rs index 89b527cca..82ab443be 100644 --- a/src/query/term_query/term_weight.rs +++ b/src/query/term_query/term_weight.rs @@ -34,12 +34,19 @@ impl TermOrEmptyOrAllScorer { } impl Weight for TermWeight { - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result> { - Ok(self.specialized_scorer(reader, boost)?.into_boxed_scorer()) + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result> { + Ok(self + .specialized_scorer(reader, boost, seek_doc)? + .into_boxed_scorer()) } fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result { - match self.specialized_scorer(reader, 1.0)? { + match self.specialized_scorer(reader, 1.0, doc)? { TermOrEmptyOrAllScorer::TermScorer(mut term_scorer) => { if term_scorer.doc() > doc || term_scorer.seek(doc) != doc { return Err(does_not_match(doc)); @@ -55,7 +62,7 @@ impl Weight for TermWeight { fn count(&self, reader: &SegmentReader) -> crate::Result { if let Some(alive_bitset) = reader.alive_bitset() { - Ok(self.scorer(reader, 1.0)?.count(alive_bitset)) + Ok(self.scorer(reader, 1.0, 0)?.count(alive_bitset)) } else { let field = self.term.field(); let inv_index = reader.inverted_index(field)?; @@ -71,7 +78,7 @@ impl Weight for TermWeight { reader: &SegmentReader, callback: &mut dyn FnMut(DocId, Score), ) -> crate::Result<()> { - match self.specialized_scorer(reader, 1.0)? { + match self.specialized_scorer(reader, 1.0, 0u32)? { TermOrEmptyOrAllScorer::TermScorer(mut term_scorer) => { for_each_scorer(&mut *term_scorer, callback); } @@ -90,7 +97,7 @@ impl Weight for TermWeight { reader: &SegmentReader, callback: &mut dyn FnMut(&[DocId]), ) -> crate::Result<()> { - match self.specialized_scorer(reader, 1.0)? { + match self.specialized_scorer(reader, 1.0, 0u32)? { TermOrEmptyOrAllScorer::TermScorer(mut term_scorer) => { let mut buffer = [0u32; COLLECT_BLOCK_BUFFER_LEN]; for_each_docset_buffered(&mut term_scorer, &mut buffer, callback); @@ -121,7 +128,7 @@ impl Weight for TermWeight { reader: &SegmentReader, callback: &mut dyn FnMut(DocId, Score) -> Score, ) -> crate::Result<()> { - let specialized_scorer = self.specialized_scorer(reader, 1.0)?; + let specialized_scorer = self.specialized_scorer(reader, 1.0, 0u32)?; match specialized_scorer { TermOrEmptyOrAllScorer::TermScorer(term_scorer) => { crate::query::boolean_query::block_wand_single_scorer( @@ -139,6 +146,12 @@ impl Weight for TermWeight { } Ok(()) } + + /// Returns a priority number used to sort weights when running an + /// intersection. + fn intersection_priority(&self) -> u32 { + 10u32 + } } impl TermWeight { @@ -169,7 +182,7 @@ impl TermWeight { reader: &SegmentReader, boost: Score, ) -> crate::Result> { - let scorer = self.specialized_scorer(reader, boost)?; + let scorer = self.specialized_scorer(reader, boost, 0u32)?; Ok(match scorer { TermOrEmptyOrAllScorer::TermScorer(scorer) => Some(*scorer), _ => None, @@ -180,6 +193,7 @@ impl TermWeight { &self, reader: &SegmentReader, boost: Score, + seek_doc: DocId, ) -> crate::Result { let field = self.term.field(); let inverted_index = reader.inverted_index(field)?; @@ -196,8 +210,11 @@ impl TermWeight { ))); } - let segment_postings: SegmentPostings = - inverted_index.read_postings_from_terminfo(&term_info, self.index_record_option)?; + let segment_postings: SegmentPostings = inverted_index.read_postings_from_terminfo( + &term_info, + self.index_record_option, + seek_doc, + )?; let fieldnorm_reader = self.fieldnorm_reader(reader)?; let similarity_weight = self.similarity_weight.boost_by(boost); diff --git a/src/query/weight.rs b/src/query/weight.rs index 23ff55c04..c6ec9f3d3 100644 --- a/src/query/weight.rs +++ b/src/query/weight.rs @@ -68,15 +68,28 @@ pub trait Weight: Send + Sync + 'static { /// /// `boost` is a multiplier to apply to the score. /// + /// As an optimization, the scorer can be positioned on any document below `seek_doc` + /// matching the request. + /// If there are no such document, it should match the first document matching the request; + /// (or TERMINATED if no documents match). + /// + /// Entirely ignoring that parameter and positionning the Scorer on the first document + /// is always correct. + /// /// See [`Query`](crate::query::Query). - fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result>; + fn scorer( + &self, + reader: &SegmentReader, + boost: Score, + seek_doc: DocId, + ) -> crate::Result>; /// Returns an [`Explanation`] for the given document. fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result; /// Returns the number documents within the given [`SegmentReader`]. fn count(&self, reader: &SegmentReader) -> crate::Result { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; if let Some(alive_bitset) = reader.alive_bitset() { Ok(scorer.count(alive_bitset)) } else { @@ -91,7 +104,7 @@ pub trait Weight: Send + Sync + 'static { reader: &SegmentReader, callback: &mut dyn FnMut(DocId, Score), ) -> crate::Result<()> { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; for_each_scorer(scorer.as_mut(), callback); Ok(()) } @@ -103,7 +116,7 @@ pub trait Weight: Send + Sync + 'static { reader: &SegmentReader, callback: &mut dyn FnMut(&[DocId]), ) -> crate::Result<()> { - let mut docset = self.scorer(reader, 1.0)?; + let mut docset = self.scorer(reader, 1.0, 0)?; let mut buffer = [0u32; COLLECT_BLOCK_BUFFER_LEN]; for_each_docset_buffered(&mut docset, &mut buffer, callback); @@ -126,8 +139,14 @@ pub trait Weight: Send + Sync + 'static { reader: &SegmentReader, callback: &mut dyn FnMut(DocId, Score) -> Score, ) -> crate::Result<()> { - let mut scorer = self.scorer(reader, 1.0)?; + let mut scorer = self.scorer(reader, 1.0, 0)?; for_each_pruning_scorer(scorer.as_mut(), threshold, callback); Ok(()) } + + /// Returns a priority number used to sort weights when running an + /// intersection. + fn intersection_priority(&self) -> u32 { + 20u32 + } }