[AI] use self_cell to make posting list evaluation lazy

prompt:

These recent changes precollect all bytes that can be used during the
query before actually running the
  query. This is fine, and desired, but the PR takes it a step further
and precollects the evaluation of
  the data. This can potentially be expensive since we need to process
the entirety of the data.

  It is done like this so that lifetime tracking gets easier. But I have
a challenge for you. Can you make
  it so that the evaluation is preserved lazily? Meaning, to keep using
iterators as outputs.

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<dyn Iterator<Item = PointOffsetType> + '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::Item> {
+        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()))
+    }
+}

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<Vec<PointOffsetType>> {
+    /// 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<LazyPostingsIter<'_>> {
         // 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<V: ZerocopyPostingValue>(
-            postings: &UniversalPostings<V, MmapFile>,
+        fn intersection<'a, V: ZerocopyPostingValue>(
+            postings: &'a UniversalPostings<V, MmapFile>,
             tokens: TokenSet,
-            filter: impl Fn(PointOffsetType) -> bool,
-        ) -> OperationResult<Vec<PointOffsetType>> {
-            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<LazyPostingsIter<'a>> {
+            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::<V>())
+                    .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<Vec<PointOffsetType>> {
+    /// Iterate lazily over point ids whose documents contain **at least one** of the given tokens.
+    fn filter_has_any(&self, tokens: TokenSet) -> OperationResult<LazyPostingsIter<'_>> {
         // in case of immutable index, deleted documents are still in the postings
         let is_active = move |idx| self.is_active(idx);
 
-        fn merge<V: ZerocopyPostingValue>(
-            postings: &UniversalPostings<V, MmapFile>,
+        fn merge<'a, V: ZerocopyPostingValue>(
+            postings: &'a UniversalPostings<V, MmapFile>,
             tokens: TokenSet,
-            is_active: impl Fn(PointOffsetType) -> bool,
-        ) -> OperationResult<Vec<PointOffsetType>> {
-            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<LazyPostingsIter<'a>> {
+            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::<V>())
+                    .try_collect()?;
+                Ok(Box::new(merge_compressed_postings_iterator(
+                    views, is_active,
+                )))
             })
         }
 
@@ -285,13 +297,15 @@ impl MmapInvertedIndex {
             tokens: &TokenSet,
             point_id: PointOffsetType,
         ) -> OperationResult<bool> {
-            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::<V>()?.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<bool> {
-            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::<V>()?.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<Vec<PointOffsetType>> {
+    /// 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<LazyPostingsIter<'_>> {
         // 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::<Positions>().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::<Positions>().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<Box<dyn Iterator<Item = PointOffsetType> + '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(

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<V: ZerocopyPostingValue>(&'a self) -> OperationResult<PostingListView<'a, V>> {
+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<V: ZerocopyPostingValue>(&self) -> OperationResult<PostingListView<'_, V>> {
         let (last_doc_id, bytes) = PointOffsetType::read_from_prefix(self.bytes.as_ref())?;
 
         let (chunks, bytes) = <[PostingChunk<SizedTypeFor<V>>]>::ref_from_prefix_with_elems(

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<V: ZerocopyPostingValue, S: UniversalRead<u8>> UniversalPostings<V, S> {
     }
 
     /// 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<T>(
-        &self,
-        header_iter: Box<dyn Iterator<Item = HeaderResult> + '_>,
+    /// 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<dyn Iterator<Item = HeaderResult> + 'a>,
         expected_capacity: usize,
-        callback: impl FnOnce(Vec<(TokenId, PostingListView<'_, V>)>) -> OperationResult<T>,
-    ) -> OperationResult<T> {
+    ) -> OperationResult<Vec<(TokenId, RawPostingList<'a>)>> {
         let header_err: Cell<Option<UniversalIoError>> = Cell::new(None);
 
         let range_iter = header_iter.filter_map(|header_res| match header_res {
@@ -197,7 +196,7 @@ impl<V: ZerocopyPostingValue, S: UniversalRead<u8>> UniversalPostings<V, S> {
             }
         });
 
-        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::<Sequential, _>(range_iter)? {
@@ -209,21 +208,16 @@ impl<V: ZerocopyPostingValue, S: UniversalRead<u8>> UniversalPostings<V, S> {
             return Err(err.into());
         }
 
-        let views = raw_postings
-            .iter()
-            .map(|(token_id, raw)| raw.as_view::<V>().map(|view| (*token_id, view)))
-            .collect::<OperationResult<Vec<_>>>()?;
-
-        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<T>(
-        &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<T>,
-    ) -> OperationResult<Option<T>> {
+    ) -> OperationResult<Option<Vec<(TokenId, RawPostingList<'a>)>>> {
         let HeadersBatch {
             iter: header_iter,
             missing,
@@ -233,36 +227,30 @@ impl<V: ZerocopyPostingValue, S: UniversalRead<u8>> UniversalPostings<V, S> {
             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<T>(
-        &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<T>,
-    ) -> OperationResult<T> {
+    ) -> OperationResult<Vec<(TokenId, RawPostingList<'a>)>> {
         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<Vec<PostingList<V>>> {
-        let mut result = Vec::new();
         let all_tokens = (0..self.header.posting_count as TokenId).collect::<Vec<_>>();
+        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::<V>().map(|view| view.to_owned()))
+            .collect()
     }
 }