From 4dfcac762881dedfc8e86fdd69c1f29b075c03c6 Mon Sep 17 00:00:00 2001 From: Pascal Seitz Date: Mon, 5 Oct 2026 16:33:35 +0200 Subject: [PATCH] Return matching term infos from automaton warmup --- src/index/inverted_index_reader.rs | 221 ++++++++++++++++++++++++----- src/lib.rs | 2 +- 2 files changed, 188 insertions(+), 35 deletions(-) diff --git a/src/index/inverted_index_reader.rs b/src/index/inverted_index_reader.rs index 5c0f9fde6..a4bd944ee 100644 --- a/src/index/inverted_index_reader.rs +++ b/src/index/inverted_index_reader.rs @@ -15,6 +15,30 @@ use crate::postings::{BlockSegmentPostings, SegmentPostings, TermInfo}; use crate::schema::{IndexRecordOption, Term, Type}; use crate::termdict::TermDictionary; +#[cfg(feature = "quickwit")] +// This is how many bytes we can hope to receive during a TTFB from S3 (~80MiB/s, 50ms). +const MERGE_HOLES_UNDER_BYTES: usize = (80 * 1024 * 1024 * 50) / 1000; + +#[cfg(feature = "quickwit")] +fn send_coalesced_posting_ranges( + posting_ranges: impl Iterator>, + sender: futures_channel::mpsc::UnboundedSender>, +) -> io::Result<()> { + let merged_posting_ranges = posting_ranges.coalesce(|range1, range2| { + if range1.end + MERGE_HOLES_UNDER_BYTES >= range2.start { + Ok(range1.start..range2.end) + } else { + Err((range1, range2)) + } + }); + for posting_range in merged_posting_ranges { + sender + .unbounded_send(posting_range) + .map_err(|_| io::Error::other("failed to send posting range back"))?; + } + Ok(()) +} + /// The inverted index reader is in charge of accessing /// the inverted index associated with a specific field. /// @@ -393,6 +417,22 @@ impl InvertedIndexReader { Ok(true) } + async fn download_posting_ranges( + &self, + posting_ranges: futures_channel::mpsc::UnboundedReceiver>, + ) -> io::Result { + let slices_downloaded = posting_ranges + .map(|posting_range| { + self.postings_file_slice + .read_bytes_slice_async(posting_range) + .map(|result| result.map(|_slice| ())) + }) + .buffer_unordered(5) + .try_collect::>() + .await?; + Ok(!slices_downloaded.is_empty()) + } + /// Warmup a block postings given a range of `Term`s. /// This method is for an advanced usage only. /// @@ -411,10 +451,7 @@ impl InvertedIndexReader { where A::State: Clone, { - // merge holes under 4MiB, that's how many bytes we can hope to receive during a TTFB from - // S3 (~80MiB/s, and 50ms latency) - const MERGE_HOLES_UNDER_BYTES: usize = (80 * 1024 * 1024 * 50) / 1000; - // we build a first iterator to download everything. Simply calling the function already + // We build a first iterator to download everything. Simply calling the function already // download everything we need from the sstable, but doesn't start iterating over it. let _term_info_iter = self .get_term_range_async(.., automaton.clone(), None, MERGE_HOLES_UNDER_BYTES) @@ -433,45 +470,75 @@ impl InvertedIndexReader { // more leaky abstraction-wise, but a lot better than the alternative let mut stream = termdict.search(automaton).without_keys().into_stream()?; - // we could do without an iterator, but this allows us access to coalesce which simplify - // things - let posting_ranges_iter = std::iter::from_fn(move || { + let posting_ranges = std::iter::from_fn(move || { stream .next_without_key() .map(|value| value.postings_range.clone()) }); - - let merged_posting_ranges_iter = posting_ranges_iter.coalesce(|range1, range2| { - if range1.end + MERGE_HOLES_UNDER_BYTES >= range2.start { - Ok(range1.start..range2.end) - } else { - Err((range1, range2)) - } - }); - - for posting_range in merged_posting_ranges_iter { - if sender.unbounded_send(posting_range).is_err() { - // this should happen only when search is cancelled - return Err(io::Error::other("failed to send posting range back")); - } - } - Ok(()) + send_coalesced_posting_ranges(posting_ranges, sender) }; let task_handle = executor(Box::new(cpu_bound_task)); - let posting_downloader = posting_ranges_to_load_stream - .map(|posting_slice| { - self.postings_file_slice - .read_bytes_slice_async(posting_slice) - .map(|result| result.map(|_slice| ())) - }) - .buffer_unordered(5) - .try_collect::>(); + let (_, postings_found) = futures_util::future::try_join( + task_handle, + self.download_posting_ranges(posting_ranges_to_load_stream), + ) + .await?; + Ok(postings_found) + } - let (_, slices_downloaded) = - futures_util::future::try_join(task_handle, posting_downloader).await?; + /// Warms postings for one automaton and returns the matching term infos. + /// + /// The term infos belong to this reader and allow reading the warmed postings without + /// searching the dictionary again. The directory must cache asynchronous reads for + /// subsequent synchronous reads. The executor must run independently of this future and + /// complete only after its task finishes. That task enumerates terms and queues downloads; + /// it does not wait for I/O or decode postings. + pub async fn warm_postings_automaton_with_term_infos< + A: Automaton + Send + 'static, + E: FnOnce(Box io::Result<()> + Send>) -> F, + F: std::future::Future>, + >( + &self, + automaton: A, + executor: E, + ) -> io::Result> + where + A::State: Clone, + { + let term_info_stream = self + .termdict + .search(&automaton) + .into_stream_async_merging_holes(MERGE_HOLES_UNDER_BYTES) + .await?; + drop(term_info_stream); - Ok(!slices_downloaded.is_empty()) + let (posting_range_sender, posting_range_receiver) = futures_channel::mpsc::unbounded(); + let (term_infos_sender, term_infos_receiver) = futures_channel::oneshot::channel(); + let termdict = self.termdict.clone(); + let cpu_bound_task = move || { + let mut stream = termdict.search(automaton).into_stream()?; + let matching_terms: Vec = + std::iter::from_fn(|| stream.next().map(|(_, info)| info.clone())).collect(); + send_coalesced_posting_ranges( + matching_terms + .iter() + .map(|info| info.postings_range.clone()), + posting_range_sender, + )?; + term_infos_sender + .send(matching_terms) + .map_err(|_| io::Error::other("failed to send automaton term infos"))?; + Ok(()) + }; + futures_util::future::try_join( + executor(Box::new(cpu_bound_task)), + self.download_posting_ranges(posting_range_receiver), + ) + .await?; + term_infos_receiver + .await + .map_err(|_| io::Error::other("automaton term info task stopped unexpectedly")) } /// Warmup the block postings for all terms. @@ -496,3 +563,89 @@ impl InvertedIndexReader { .unwrap_or(0u32)) } } + +#[cfg(all(test, feature = "quickwit"))] +mod tests { + use std::io; + + use futures::channel::oneshot; + use tantivy_fst::Regex; + + use crate::schema::{Schema, STRING}; + use crate::{DocSet, Index, IndexWriter, TERMINATED}; + + fn execute_on_thread( + task: Box io::Result<()> + Send>, + ) -> impl std::future::Future> { + let (sender, receiver) = oneshot::channel(); + std::thread::spawn(move || { + let _ = sender.send(task()); + }); + async move { + receiver + .await + .map_err(|_| io::Error::other("executor task panicked"))? + } + } + + #[test] + fn test_warm_postings_automaton_with_term_infos() -> crate::Result<()> { + let mut schema_builder = Schema::builder(); + let field = schema_builder.add_text_field("field", STRING); + let index = Index::create_in_ram(schema_builder.build()); + let mut writer: IndexWriter = index.writer_for_tests()?; + writer.add_document(doc!(field => "apple", field => "banana"))?; + writer.add_document(doc!(field => "apricot"))?; + writer.add_document(doc!(field => "banana", field => "berry"))?; + writer.add_document(doc!(field => "carrot"))?; + // Exercise full postings blocks as well as short lists and duplicate matches. + for _ in 0..260 { + writer.add_document(doc!(field => "apple", field => "apricot"))?; + } + writer.commit()?; + let searcher = index.reader()?.searcher(); + let segment_reader = searcher.segment_reader(0); + let inverted_index = segment_reader.inverted_index(field)?; + for (pattern, expected) in [ + ("b.*", vec![0, 2]), + ("z.*", vec![]), + ("a.*", [vec![0, 1], (4..264).collect()].concat()), + (".*", (0..264).collect()), + ] { + let term_infos = futures::executor::block_on( + inverted_index.warm_postings_automaton_with_term_infos( + Regex::new(pattern).unwrap(), + execute_on_thread, + ), + )?; + let mut docs = Vec::new(); + for term_info in term_infos { + let mut postings = inverted_index.read_postings_from_terminfo( + &term_info, + crate::schema::IndexRecordOption::Basic, + )?; + while postings.doc() != TERMINATED { + docs.push(postings.doc()); + postings.advance(); + } + } + docs.sort_unstable(); + docs.dedup(); + assert_eq!(docs, expected); + } + Ok(()) + } + + #[test] + fn test_warm_postings_automaton_with_term_infos_empty_index() -> crate::Result<()> { + let inverted_index = + super::InvertedIndexReader::empty(crate::schema::IndexRecordOption::Basic); + let term_infos = + futures::executor::block_on(inverted_index.warm_postings_automaton_with_term_infos( + Regex::new(".*").unwrap(), + execute_on_thread, + ))?; + assert!(term_infos.is_empty()); + Ok(()) + } +} diff --git a/src/lib.rs b/src/lib.rs index 8d1912690..8ffcab1f6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -170,7 +170,7 @@ mod future_result; // Re-exports pub use columnar; -pub use common::{ByteCount, DateTime}; +pub use common::{BitSet, ByteCount, DateTime}; pub use query_grammar; pub use time;