diff --git a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/lazy_iter.rs b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/lazy_iter.rs new file mode 100644 index 0000000000..3d717207b7 --- /dev/null +++ b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/lazy_iter.rs @@ -0,0 +1,41 @@ +use common::types::PointOffsetType; +use self_cell::self_cell; + +use crate::index::field_index::full_text_index::inverted_index::TokenId; +use crate::index::field_index::full_text_index::inverted_index::mmap_inverted_index::raw_posting_list::RawPostingList; + +// Alias is required by `self_cell`: the dependent's lifetime parameter must be +// expressed via a single-lifetime alias. +type BoxedPointOffsetIter<'a> = Box + 'a>; + +self_cell! { + /// Lazy iterator over `PointOffsetType` that keeps the underlying raw + /// posting list bytes alive. The owner is a `Vec<(TokenId, RawPostingList)>` + /// produced by `UniversalPostings::fetch_*`, and the dependent is a boxed + /// iterator built on top of views into that vector. + /// + /// Construction happens up-front (I/O is done in `fetch_*`), but the + /// actual query evaluation — posting-list traversal, intersection, phrase + /// matching — is driven on-demand from `next()`. + pub struct LazyPostingsIter<'a> { + owner: Vec<(TokenId, RawPostingList<'a>)>, + #[covariant] + dependent: BoxedPointOffsetIter, + } +} + +impl Iterator for LazyPostingsIter<'_> { + type Item = PointOffsetType; + + fn next(&mut self) -> Option { + self.with_dependent_mut(|_owner, dependent| dependent.next()) + } +} + +impl<'a> LazyPostingsIter<'a> { + /// A `LazyPostingsIter` that yields nothing. Used when the query cannot + /// possibly match (empty posting set, missing tokens, etc.). + pub fn empty() -> Self { + Self::new(Vec::new(), |_| Box::new(std::iter::empty())) + } +} diff --git a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/mod.rs b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/mod.rs index fcfca80ee7..39c3039607 100644 --- a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/mod.rs +++ b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/mod.rs @@ -7,6 +7,7 @@ use common::mmap::{self, Advice, AdviceSetting, MmapSlice, create_and_ensure_len use common::mmap_hashmap::{MmapHashMap, READ_ENTRY_OVERHEAD}; use common::types::PointOffsetType; use common::universal_io::{MmapFile, OpenOptions}; +use itertools::Itertools; use types::ZerocopyPostingValue; use uio_postings::UniversalPostings; @@ -29,11 +30,14 @@ use crate::index::field_index::full_text_index::inverted_index::postings_iterato }; mod create_postings; +mod lazy_iter; pub mod mmap_postings_enum; mod raw_posting_list; pub(in crate::index::field_index::full_text_index) mod types; mod uio_postings; +use self::lazy_iter::LazyPostingsIter; + const POSTINGS_FILE: &str = "postings.dat"; const VOCAB_FILE: &str = "vocab.dat"; const POINT_TO_TOKENS_COUNT_FILE: &str = "point_to_tokens_count.dat"; @@ -202,60 +206,68 @@ impl MmapInvertedIndex { !is_deleted } - /// Iterate over point ids whose documents contain all given tokens. + /// Iterate lazily over point ids whose documents contain **all** given tokens. /// - /// Pre-collected upfront because [`UniversalPostings`] exposes posting - /// views via a `FnOnce` callback. Acceptable since this index lives - /// on disk. - pub fn filter_has_all(&self, tokens: TokenSet) -> OperationResult> { + /// I/O (the header reads and posting-list byte reads) is performed eagerly + /// inside `fetch_all_or_none`, so the returned iterator does not issue new + /// reads. The actual intersection work is performed on-demand as the + /// caller pulls items. + pub fn filter_has_all(&self, tokens: TokenSet) -> OperationResult> { // in case of mmap immutable index, deleted points are still in the postings - let filter = move |idx| self.is_active(idx); + let is_active = move |idx| self.is_active(idx); - fn intersection( - postings: &UniversalPostings, + fn intersection<'a, V: ZerocopyPostingValue>( + postings: &'a UniversalPostings, tokens: TokenSet, - filter: impl Fn(PointOffsetType) -> bool, - ) -> OperationResult> { - let result = - postings.with_all_or_none_postings(tokens.tokens(), |posting_readers| { - if posting_readers.is_empty() { - return Ok(Vec::new()); - } - let posting_readers = posting_readers - .into_iter() - .map(|(_token_id, posting_list_view)| posting_list_view) - .collect(); - Ok(intersect_compressed_postings_iterator(posting_readers, filter).collect()) - })?; - // Some token has no posting list -> no matches - Ok(result.unwrap_or_default()) + is_active: impl Fn(PointOffsetType) -> bool + 'a, + ) -> OperationResult> { + let Some(raw_postings) = postings.fetch_all_or_none(tokens.tokens())? else { + return Ok(LazyPostingsIter::empty()); + }; + if raw_postings.is_empty() { + return Ok(LazyPostingsIter::empty()); + } + + LazyPostingsIter::try_new(raw_postings, move |raw_postings| { + let views = raw_postings + .iter() + .map(|(_token_id, raw)| raw.as_view::()) + .try_collect()?; + Ok(Box::new(intersect_compressed_postings_iterator( + views, is_active, + ))) + }) } match &self.storage.postings { - MmapPostingsEnum::Ids(postings) => intersection(postings, tokens, filter), - MmapPostingsEnum::WithPositions(postings) => intersection(postings, tokens, filter), + MmapPostingsEnum::Ids(postings) => intersection(postings, tokens, is_active), + MmapPostingsEnum::WithPositions(postings) => intersection(postings, tokens, is_active), } } - /// Iterate over point ids whose documents contain at least one of the given tokens - fn filter_has_any(&self, tokens: TokenSet) -> OperationResult> { + /// Iterate lazily over point ids whose documents contain **at least one** of the given tokens. + fn filter_has_any(&self, tokens: TokenSet) -> OperationResult> { // in case of immutable index, deleted documents are still in the postings let is_active = move |idx| self.is_active(idx); - fn merge( - postings: &UniversalPostings, + fn merge<'a, V: ZerocopyPostingValue>( + postings: &'a UniversalPostings, tokens: TokenSet, - is_active: impl Fn(PointOffsetType) -> bool, - ) -> OperationResult> { - postings.with_existing_postings(tokens.tokens(), |posting_readers| { - if posting_readers.is_empty() { - return Ok(Vec::new()); - } - let posting_readers = posting_readers - .into_iter() - .map(|(_token_id, posting_list_view)| posting_list_view) - .collect(); - Ok(merge_compressed_postings_iterator(posting_readers, is_active).collect()) + is_active: impl Fn(PointOffsetType) -> bool + 'a, + ) -> OperationResult> { + let raw_postings = postings.fetch_existing(tokens.tokens())?; + if raw_postings.is_empty() { + return Ok(LazyPostingsIter::empty()); + } + + LazyPostingsIter::try_new(raw_postings, move |raw_postings| { + let views = raw_postings + .iter() + .map(|(_token_id, raw)| raw.as_view::()) + .try_collect()?; + Ok(Box::new(merge_compressed_postings_iterator( + views, is_active, + ))) }) } @@ -285,13 +297,15 @@ impl MmapInvertedIndex { tokens: &TokenSet, point_id: PointOffsetType, ) -> OperationResult { - let result = postings.with_all_or_none_postings(tokens.tokens(), |all_postings| { - Ok(all_postings - .into_iter() - .all(|(_token_id, posting)| posting.visitor().contains(point_id))) - })?; - // Some token has no posting list -> no match - Ok(result.unwrap_or(false)) + let Some(raw_postings) = postings.fetch_all_or_none(tokens.tokens())? else { + return Ok(false); + }; + for (_token_id, raw) in &raw_postings { + if !raw.as_view::()?.visitor().contains(point_id) { + return Ok(false); + } + } + Ok(true) } match &self.storage.postings { @@ -317,11 +331,13 @@ impl MmapInvertedIndex { tokens: &TokenSet, point_id: PointOffsetType, ) -> OperationResult { - postings.with_existing_postings(tokens.tokens(), |all_postings| { - Ok(all_postings - .into_iter() - .any(|(_token_id, posting)| posting.visitor().contains(point_id))) - }) + let raw_postings = postings.fetch_existing(tokens.tokens())?; + for (_token_id, raw) in &raw_postings { + if raw.as_view::()?.visitor().contains(point_id) { + return Ok(true); + } + } + Ok(false) } match &self.storage.postings { @@ -330,8 +346,9 @@ impl MmapInvertedIndex { } } - /// Iterate over point ids whose documents contain all given tokens in the same order they are provided - pub fn filter_has_phrase(&self, phrase: Document) -> OperationResult> { + /// Iterate lazily over point ids whose documents contain all given tokens + /// in the same order they are provided. + pub fn filter_has_phrase(&self, phrase: Document) -> OperationResult> { // in case of mmap immutable index, deleted points are still in the postings let is_active = move |idx| self.is_active(idx); @@ -341,22 +358,27 @@ impl MmapInvertedIndex { // not fetch the same posting list twice, otherwise positions get // added twice in `phrase_in_all_postings`. let unique_tokens = phrase.to_token_set(); - let result = postings.with_all_or_none_postings( - unique_tokens.tokens(), - |selected_postings| { - Ok(intersect_compressed_postings_phrase_iterator( - phrase, - selected_postings, - is_active, - ) - .collect()) - }, - )?; - // Some token has no posting list -> no matches - Ok(result.unwrap_or_default()) + let Some(raw_postings) = postings.fetch_all_or_none(unique_tokens.tokens())? else { + return Ok(LazyPostingsIter::empty()); + }; + if raw_postings.is_empty() { + return Ok(LazyPostingsIter::empty()); + } + + LazyPostingsIter::try_new(raw_postings, move |raw_postings| { + let views = raw_postings + .iter() + .map(|(token_id, raw)| { + raw.as_view::().map(|view| (*token_id, view)) + }) + .try_collect()?; + Ok(Box::new(intersect_compressed_postings_phrase_iterator( + phrase, views, is_active, + ))) + }) } // cannot do phrase matching if there's no positional information - MmapPostingsEnum::Ids(_postings) => Ok(Vec::new()), + MmapPostingsEnum::Ids(_postings) => Ok(LazyPostingsIter::empty()), } } @@ -373,18 +395,19 @@ impl MmapInvertedIndex { match &self.storage.postings { MmapPostingsEnum::WithPositions(postings) => { let unique_tokens = phrase.to_token_set(); - let result = postings.with_all_or_none_postings( - unique_tokens.tokens(), - |selected_postings| { - Ok(check_compressed_postings_phrase( - phrase, - point_id, - selected_postings, - )) - }, - )?; - // Some token has no posting list -> no match - Ok(result.unwrap_or(false)) + let Some(raw_postings) = postings.fetch_all_or_none(unique_tokens.tokens())? else { + return Ok(false); + }; + let views = { + let raw_postings = &raw_postings; + raw_postings + .iter() + .map(|(token_id, raw)| { + raw.as_view::().map(|view| (*token_id, view)) + }) + .try_collect() + }?; + Ok(check_compressed_postings_phrase(phrase, point_id, views)) } // cannot do phrase matching if there's no positional information MmapPostingsEnum::Ids(_postings) => Ok(false), @@ -496,12 +519,11 @@ impl InvertedIndex for MmapInvertedIndex { query: ParsedQuery, _hw_counter: &HardwareCounterCell, ) -> OperationResult + 'a>> { - let ids = match query { - ParsedQuery::AllTokens(tokens) => self.filter_has_all(tokens)?, - ParsedQuery::Phrase(phrase) => self.filter_has_phrase(phrase)?, - ParsedQuery::AnyTokens(tokens) => self.filter_has_any(tokens)?, - }; - Ok(Box::new(ids.into_iter())) + match query { + ParsedQuery::AllTokens(tokens) => Ok(Box::new(self.filter_has_all(tokens)?)), + ParsedQuery::Phrase(phrase) => Ok(Box::new(self.filter_has_phrase(phrase)?)), + ParsedQuery::AnyTokens(tokens) => Ok(Box::new(self.filter_has_any(tokens)?)), + } } fn get_posting_len( diff --git a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/raw_posting_list.rs b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/raw_posting_list.rs index dc336067f6..fdb8ed8c13 100644 --- a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/raw_posting_list.rs +++ b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/raw_posting_list.rs @@ -21,8 +21,12 @@ impl<'a> RawPostingList<'a> { } } -impl<'a> RawPostingList<'a> { - pub fn as_view(&'a self) -> OperationResult> { +impl RawPostingList<'_> { + /// Decode the raw bytes into a [`PostingListView`]. The view borrows from + /// `self` (not from the original `'a`), which allows it to be produced from + /// contexts that hold `RawPostingList` by value (e.g., inside a `self_cell` + /// builder where the outer `'a` is not reachable). + pub fn as_view(&self) -> OperationResult> { let (last_doc_id, bytes) = PointOffsetType::read_from_prefix(self.bytes.as_ref())?; let (chunks, bytes) = <[PostingChunk>]>::ref_from_prefix_with_elems( diff --git a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/uio_postings.rs b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/uio_postings.rs index 3195bbb1b6..b0f8425844 100644 --- a/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/uio_postings.rs +++ b/lib/segment/src/index/field_index/full_text_index/inverted_index/mmap_inverted_index/uio_postings.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use common::generic_consts::{Random, Sequential}; use common::universal_io::{OpenOptions, ReadRange, UniversalIoError, UniversalRead}; -use posting_list::{PostingList, PostingListView}; +use posting_list::PostingList; use zerocopy::FromBytes; use crate::common::operation_error::OperationResult; @@ -172,15 +172,14 @@ impl> UniversalPostings { } /// Read the posting lists for every header yielded by `header_iter` and - /// hand the resulting views to `callback`. Header iteration is pipelined - /// into the posting reads - submissions are scheduled lazily as headers - /// arrive. - fn with_posting_views( - &self, - header_iter: Box + '_>, + /// return them as `Vec<(TokenId, RawPostingList)>`. Header iteration is + /// pipelined into the posting reads — submissions are scheduled lazily as + /// headers arrive. + fn fetch_raw_postings<'a>( + &'a self, + header_iter: Box + 'a>, expected_capacity: usize, - callback: impl FnOnce(Vec<(TokenId, PostingListView<'_, V>)>) -> OperationResult, - ) -> OperationResult { + ) -> OperationResult)>> { let header_err: Cell> = Cell::new(None); let range_iter = header_iter.filter_map(|header_res| match header_res { @@ -197,7 +196,7 @@ impl> UniversalPostings { } }); - let mut raw_postings: Vec<(TokenId, RawPostingList<'_>)> = + let mut raw_postings: Vec<(TokenId, RawPostingList<'a>)> = Vec::with_capacity(expected_capacity); for entry in self.storage.read_iter::(range_iter)? { @@ -209,21 +208,16 @@ impl> UniversalPostings { return Err(err.into()); } - let views = raw_postings - .iter() - .map(|(token_id, raw)| raw.as_view::().map(|view| (*token_id, view))) - .collect::>>()?; - - callback(views) + Ok(raw_postings) } - /// Fetch posting lists for the given token ids and hand them to `callback`. - /// If any of posting lists is not found - return `None`. - pub fn with_all_or_none_postings( - &self, + /// Fetch the raw bytes of posting lists for the given token ids. + /// Returns `None` if any posting list is missing; otherwise returns + /// the full set in the order they were read from storage. + pub fn fetch_all_or_none<'a>( + &'a self, token_ids: &[TokenId], - callback: impl FnOnce(Vec<(TokenId, PostingListView<'_, V>)>) -> OperationResult, - ) -> OperationResult> { + ) -> OperationResult)>>> { let HeadersBatch { iter: header_iter, missing, @@ -233,36 +227,30 @@ impl> UniversalPostings { return Ok(None); } - self.with_posting_views(header_iter, token_ids.len(), callback) + self.fetch_raw_postings(header_iter, token_ids.len()) .map(Some) } - /// Fetch all existing posting lists for the given token ids and hand them - /// to `callback`. Token ids without a posting list are silently ignored. - pub fn with_existing_postings( - &self, + /// Fetch the raw bytes of every existing posting list among `token_ids`. + /// Token ids without a posting list are silently ignored. + pub fn fetch_existing<'a>( + &'a self, token_ids: &[TokenId], - callback: impl FnOnce(Vec<(TokenId, PostingListView<'_, V>)>) -> OperationResult, - ) -> OperationResult { + ) -> OperationResult)>> { let HeadersBatch { iter: header_iter, .. } = self.headers_iter(token_ids)?; - self.with_posting_views(header_iter, token_ids.len(), callback) + self.fetch_raw_postings(header_iter, token_ids.len()) } pub fn all_postings(&self) -> OperationResult>> { - let mut result = Vec::new(); let all_tokens = (0..self.header.posting_count as TokenId).collect::>(); + let raw_postings = self.fetch_existing(&all_tokens)?; - self.with_existing_postings(&all_tokens, |views| { - for (_, view) in views { - let posting_list = view.to_owned(); - result.push(posting_list); - } - Ok(()) - })?; - - Ok(result) + raw_postings + .iter() + .map(|(_token_id, raw)| raw.as_view::().map(|view| view.to_owned())) + .collect() } }