refactor(fts): add exact posting load policies

rust/lance-index/src/scalar/inverted/index/partition.rs

@@ -2,6 +2,105 @@
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
 use super::*;
+use smallvec::SmallVec;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct PositionMatchSummary {
+    exact_scoring_required: bool,
+    every_position_matched: bool,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(in super::super) struct PostingLoadOptions {
+    force_global_scorer: bool,
+    read_policy: PostingReadPolicy,
+}
+
+impl PostingLoadOptions {
+    const fn read_ahead(force_global_scorer: bool) -> Self {
+        Self {
+            force_global_scorer,
+            read_policy: PostingReadPolicy::ReadAhead,
+        }
+    }
+
+    #[cfg(test)]
+    pub(in super::super) const fn cache_aware_exact(force_global_scorer: bool) -> Self {
+        Self {
+            force_global_scorer,
+            read_policy: PostingReadPolicy::CacheAwareExact,
+        }
+    }
+}
+
+fn summarize_position_matches(mut positions: SmallVec<[(u32, bool); 8]>) -> PositionMatchSummary {
+    positions.sort_unstable_by_key(|(position, _)| *position);
+
+    let mut exact_scoring_required = false;
+    let mut every_position_matched = true;
+    let mut group_start = 0;
+    while group_start < positions.len() {
+        let position = positions[group_start].0;
+        let mut group_end = group_start + 1;
+        let mut group_matched = positions[group_start].1;
+        while group_end < positions.len() && positions[group_end].0 == position {
+            exact_scoring_required = true;
+            group_matched |= positions[group_end].1;
+            group_end += 1;
+        }
+        every_position_matched &= group_matched;
+        group_start = group_end;
+    }
+
+    PositionMatchSummary {
+        exact_scoring_required,
+        every_position_matched,
+    }
+}
+
+fn posting_group_demand_counts(
+    inverted_list: &PostingListReader,
+    token_ids: &[(u32, String, u32)],
+) -> HashMap<(u32, u32), usize> {
+    let mut demanded_token_ids = token_ids
+        .iter()
+        .map(|(token_id, _, _)| *token_id)
+        .collect::<Vec<_>>();
+    demanded_token_ids.sort_unstable();
+    demanded_token_ids.dedup();
+
+    let mut counts = HashMap::new();
+    for token_id in demanded_token_ids {
+        if let Some(group) = inverted_list.group_range_for_token(token_id) {
+            *counts.entry(group).or_default() += 1;
+        }
+    }
+    counts
+}
+
+fn effective_posting_read_policy(
+    inverted_list: &PostingListReader,
+    requested_policy: PostingReadPolicy,
+    group_demand_counts: &HashMap<(u32, u32), usize>,
+    token_id: u32,
+) -> PostingReadPolicy {
+    if requested_policy == PostingReadPolicy::ReadAhead {
+        return PostingReadPolicy::ReadAhead;
+    }
+    let Some(group) = inverted_list.group_range_for_token(token_id) else {
+        return PostingReadPolicy::CacheAwareExact;
+    };
+    let demand_count = group_demand_counts.get(&group).copied();
+    debug_assert!(
+        demand_count.is_some(),
+        "posting group {group:?} must have a demand count for token {token_id}"
+    );
+    if demand_count == Some(1) {
+        PostingReadPolicy::CacheAwareExact
+    } else {
+        PostingReadPolicy::ReadAhead
+    }
+}
 
 /// Query-level inputs for a grouped-term score upper bound.
 ///
@@ -510,7 +609,6 @@ impl InvertedPartition {
     // bounds must share corpus-level statistics before the global collector
     // can safely propagate its threshold. Old posting formats without impacts
     // fall back to a scorer-derived global upper bound in that mode.
-    #[instrument(level = "debug", skip_all)]
     pub(in super::super) async fn load_posting_lists(
         &self,
         tokens: &Tokens,
@@ -520,59 +618,96 @@ impl InvertedPartition {
         metrics: &dyn MetricsCollector,
         force_global_scorer: bool,
     ) -> Result<LoadedPostings> {
+        self.load_posting_lists_with_policy(
+            tokens,
+            params,
+            operator,
+            impact_scorer,
+            metrics,
+            PostingLoadOptions::read_ahead(force_global_scorer),
+        )
+        .await
+    }
+
+    #[instrument(name = "load_posting_lists", level = "debug", skip_all)]
+    pub(in super::super) async fn load_posting_lists_with_policy(
+        &self,
+        tokens: &Tokens,
+        params: &FtsSearchParams,
+        operator: Operator,
+        impact_scorer: &MemBM25Scorer,
+        metrics: &dyn MetricsCollector,
+        options: PostingLoadOptions,
+    ) -> Result<LoadedPostings> {
+        let PostingLoadOptions {
+            force_global_scorer,
+            read_policy: requested_read_policy,
+        } = options;
         let is_phrase_query = params.phrase_slop.is_some();
         let is_and_query = operator == Operator::And;
-        let required_positions = (is_and_query || is_phrase_query).then(|| {
-            (0..tokens.len())
-                .map(|index| tokens.position(index))
-                .collect::<HashSet<_>>()
-        });
         // Fuzzy expansion already ran once at the index level (see
         // `InvertedIndex::bm25_search`) under the global `max_expansions`
         // budget. Positions identify alternatives that must share one posting
         // iterator, including code identifier subwords and fuzzy expansions.
-        let tokens = tokens.clone();
-        let token_positions = (0..tokens.len())
-            .map(|index| tokens.position(index))
-            .collect::<Vec<_>>();
-        let mut seen_positions = HashSet::with_capacity(token_positions.len());
-        let exact_scoring_required = token_positions
-            .iter()
-            .any(|position| !seen_positions.insert(*position));
         let mut token_ids = Vec::with_capacity(tokens.len());
-        let mut matched_positions = required_positions.as_ref().map(|_| HashSet::new());
-        for (index, token) in tokens.into_iter().enumerate() {
-            let token_id = self.map(&token);
+        let mut position_matches = SmallVec::<[(u32, bool); 8]>::new();
+        for index in 0..tokens.len() {
+            let token = tokens.get_token(index);
+            let position = tokens.position(index);
+            let token_id = self.map(token);
+            position_matches.push((position, token_id.is_some()));
             if let Some(token_id) = token_id {
-                let position = token_positions[index];
-                if let Some(matched_positions) = matched_positions.as_mut() {
-                    matched_positions.insert(position);
-                }
-                token_ids.push((token_id, token, position));
+                token_ids.push((token_id, token.to_owned(), position));
             }
         }
+        let position_summary = summarize_position_matches(position_matches);
+        let exact_scoring_required = position_summary.exact_scoring_required;
         if token_ids.is_empty() {
             return Ok(LoadedPostings::empty());
         }
-        if let Some(required_positions) = required_positions.as_ref()
-            && let Some(matched_positions) = matched_positions.as_ref()
-            && !required_positions.is_subset(matched_positions)
-        {
+        if (is_and_query || is_phrase_query) && !position_summary.every_position_matched {
             return Ok(LoadedPostings::empty());
         }
 
         token_ids.sort_unstable_by_key(|(token_id, _, position)| (*position, *token_id));
         token_ids.dedup_by(|lhs, rhs| lhs.0 == rhs.0 && lhs.2 == rhs.2);
 
+        let group_demand_counts = if requested_read_policy == PostingReadPolicy::CacheAwareExact {
+            posting_group_demand_counts(self.inverted_list.as_ref(), &token_ids)
+        } else {
+            HashMap::new()
+        };
+
         let num_docs = self.docs.len();
         let loaded_postings = stream::iter(token_ids)
-            .map(|(token_id, token, position)| async move {
-                let posting = self
-                    .inverted_list
-                    .posting_list(token_id, is_phrase_query, metrics)
-                    .await?;
+            .map(|(token_id, token, position)| {
+                let read_policy = effective_posting_read_policy(
+                    self.inverted_list.as_ref(),
+                    requested_read_policy,
+                    &group_demand_counts,
+                    token_id,
+                );
+                async move {
+                    let posting = match read_policy {
+                        PostingReadPolicy::ReadAhead => {
+                            self.inverted_list
+                                .posting_list(token_id, is_phrase_query, metrics)
+                                .await?
+                        }
+                        PostingReadPolicy::CacheAwareExact => {
+                            self.inverted_list
+                                .posting_list_with_policy(
+                                    token_id,
+                                    is_phrase_query,
+                                    metrics,
+                                    read_policy,
+                                )
+                                .await?
+                        }
+                    };
 
-                Result::Ok((token_id, token, position, posting))
+                    Result::Ok((token_id, token, position, posting))
+                }
             })
             .buffered(self.store.io_parallelism())
             .try_collect::<Vec<_>>()
@@ -857,6 +992,49 @@ mod tests {
 
     use super::*;
 
+    fn position_summary(entries: &[(u32, bool)]) -> PositionMatchSummary {
+        summarize_position_matches(entries.iter().copied().collect())
+    }
+
+    #[test]
+    fn position_summary_marks_or_duplicates_for_exact_scoring() {
+        let summary = position_summary(&[(0, true), (0, false), (1, false)]);
+
+        assert!(summary.exact_scoring_required);
+        assert!(!summary.every_position_matched);
+    }
+
+    #[test]
+    fn position_summary_requires_a_match_in_every_and_group() {
+        let complete = position_summary(&[(0, false), (0, true), (1, true)]);
+        let incomplete = position_summary(&[(0, true), (1, false), (1, false)]);
+
+        assert!(complete.exact_scoring_required);
+        assert!(complete.every_position_matched);
+        assert!(incomplete.exact_scoring_required);
+        assert!(!incomplete.every_position_matched);
+    }
+
+    #[test]
+    fn position_summary_groups_nonadjacent_positions() {
+        let summary = position_summary(&[(9, false), (1, true), (4, true), (9, true)]);
+
+        assert!(summary.exact_scoring_required);
+        assert!(summary.every_position_matched);
+    }
+
+    #[test]
+    fn position_summary_spills_past_eight_tokens_without_losing_exactness() {
+        let mut positions = SmallVec::<[(u32, bool); 8]>::new();
+        positions.extend((0..10).map(|position| (position, true)));
+        positions.push((3, false));
+        assert!(positions.spilled());
+
+        let summary = summarize_position_matches(positions);
+        assert!(summary.exact_scoring_required);
+        assert!(summary.every_position_matched);
+    }
+
     #[rstest]
     #[case::plain(false)]
     #[case::v3_compressed(true)]