lance-index 13.0.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

use super::partition::FuzzyAutomaton;
use super::*;

const LANCE_FTS_REUSE_PREPARED_SCORER_ENV: &str = "LANCE_FTS_REUSE_PREPARED_SCORER";

/// Segment BM25 statistics read without I/O, with the dictionary lookups they
/// were derived from.
#[derive(Debug, PartialEq)]
pub(in crate::scalar::inverted) struct LoadedTermStats {
    /// `(total_tokens, num_docs, per_term_doc_freq)`
    pub(in crate::scalar::inverted) stats: (u64, usize, Vec<usize>),
    /// Partition-major token ids: `token_ids[partition * terms.len() + term]`.
    pub(in crate::scalar::inverted) token_ids: Box<[Option<u32>]>,
}

fn reuse_prepared_scorer_enabled_from_value(value: Option<&str>) -> bool {
    !value.is_some_and(|value| {
        let value = value.trim();
        value == "0" || value.eq_ignore_ascii_case("off")
    })
}

static LANCE_FTS_REUSE_PREPARED_SCORER_ENABLED: LazyLock<bool> = LazyLock::new(|| {
    reuse_prepared_scorer_enabled_from_value(
        std::env::var(LANCE_FTS_REUSE_PREPARED_SCORER_ENV)
            .ok()
            .as_deref(),
    )
});

fn select_impact_scorer(
    scorer: &MemBM25Scorer,
    prepared_scorer: Option<&Arc<MemBM25Scorer>>,
    is_legacy: impl FnOnce() -> bool,
    reuse_prepared_scorer_enabled: impl FnOnce() -> bool,
) -> Arc<MemBM25Scorer> {
    let Some(prepared_scorer) = prepared_scorer else {
        return Arc::new(scorer.clone());
    };
    if is_legacy() || !reuse_prepared_scorer_enabled() {
        return Arc::new(scorer.clone());
    }
    debug_assert!(
        std::ptr::eq(scorer, prepared_scorer.as_ref()),
        "prepared BM25 scorer must be the canonical final scorer"
    );
    Arc::clone(prepared_scorer)
}

impl InvertedIndex {
    /// Add this segment's lexicographically smallest fuzzy candidates for one
    /// compiled query-token automaton to a caller-owned cross-segment merge
    /// set.
    ///
    /// The set is trimmed after every partition merge, so it always contains
    /// at most `limit` terms. Dropping the current largest term is lossless:
    /// no later merge can make it one of the globally smallest `limit` terms.
    pub(in crate::scalar::inverted) fn collect_fuzzy_candidates_with_automaton(
        &self,
        automaton: &FuzzyAutomaton,
        limit: usize,
        candidates: &mut BTreeSet<String>,
    ) -> Result<()> {
        // The caller owns one compiled automaton for this source token. Reuse
        // it across every physical partition instead of rebuilding a DFA per
        // dictionary.
        while candidates.len() > limit {
            candidates.pop_last();
        }
        for partition in &self.partitions {
            partition.collect_fuzzy_candidates_with_automaton(automaton, limit, candidates)?;
            while candidates.len() > limit {
                candidates.pop_last();
            }
            debug_assert!(candidates.len() <= limit);
        }
        Ok(())
    }

    /// Build a single-segment [`MemBM25Scorer`] whose per-term IDF table
    /// covers every token that the per-partition scoring loop will look
    /// up. For fuzzy queries that means the union of Levenshtein
    /// expansions, not just the raw query tokens — otherwise
    /// `query_weight(expanded_token)` returns 0 and the BM25 contribution
    /// of every expanded match is discarded.
    pub async fn bm25_base_scorer(
        &self,
        query_tokens: &Tokens,
        params: &FtsSearchParams,
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<MemBM25Scorer> {
        if uses_fuzzy_expansion(params.fuzziness) {
            let expanded = self.expand_fuzzy_tokens(query_tokens, params)?;
            self.bm25_scorer_for_final_tokens(&expanded, metrics).await
        } else {
            self.bm25_scorer_for_final_tokens(query_tokens, metrics)
                .await
        }
    }

    /// Scorer for a token list that needs no further fuzzy expansion: dedup
    /// the terms and pull their document frequencies. `bm25_search` calls
    /// this with the tokens it already expanded, so the expansion runs once
    /// per query rather than once for the scorer and once per partition.
    async fn bm25_scorer_for_final_tokens(
        &self,
        tokens: &Tokens,
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<MemBM25Scorer> {
        let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?;
        let mut terms: Vec<String> = Vec::new();
        let mut seen = HashSet::new();
        for token in tokens {
            if seen.insert(token.to_string()) {
                terms.push(token.to_string());
            }
        }
        let mut token_docs = HashMap::with_capacity(terms.len());
        for term in &terms {
            let df = self.df_for_term(term, metrics).await?;
            token_docs.insert(term.clone(), df);
        }
        Ok(MemBM25Scorer::new(total_tokens, num_docs, token_docs))
    }

    /// Corpus statistics at document granularity: `(total_tokens, docCount,
    /// per_term_docFreq)`, where a document is one entry of a partition's
    /// [`DocSet`]. This is what the single-column scoring path pairs with, since
    /// wand scores one posting per document.
    pub async fn bm25_stats_for_terms(
        &self,
        terms: &[String],
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<(u64, usize, Vec<usize>)> {
        let (total_tokens, num_docs) = self.aggregate_corpus_stats().await?;
        let token_docs =
            futures::future::try_join_all(terms.iter().map(|term| self.df_for_term(term, metrics)))
                .await?;
        Ok((total_tokens, num_docs, token_docs))
    }

    /// Return BM25 statistics synchronously when every required value is loaded.
    ///
    /// This is an all-or-nothing probe for the full-prewarm query path. It never
    /// starts I/O: an absent segment corpus statistic or one absent modern
    /// posting-length table returns `None` for the entire segment.
    pub(in crate::scalar::inverted) fn bm25_stats_for_terms_if_loaded(
        &self,
        terms: &[String],
    ) -> Result<Option<LoadedTermStats>> {
        // Keep the legacy reader on its frozen asynchronous compatibility
        // path; this optimization targets current partitioned formats.
        if self.is_legacy() {
            return Ok(None);
        }
        let Some(&(total_tokens, num_docs)) = self.corpus_stats.get() else {
            return Ok(None);
        };
        if self
            .partitions
            .iter()
            .any(|partition| !partition.inverted_list.posting_lengths_loaded())
        {
            return Ok(None);
        }
        let lookups = self
            .partitions
            .iter()
            .flat_map(|partition| {
                terms
                    .iter()
                    .map(|term| (partition.tokens.as_ref(), term.as_str()))
            })
            .collect::<Vec<_>>();
        let token_ids = TokenSet::get_many(&lookups);
        let mut token_docs = vec![0_usize; terms.len()];
        for (partition, partition_token_ids) in self
            .partitions
            .iter()
            .zip(token_ids.chunks(terms.len().max(1)))
        {
            for ((term, term_docs), &token_id) in terms
                .iter()
                .zip(token_docs.iter_mut())
                .zip(partition_token_ids)
            {
                let Some(token_id) = token_id else {
                    continue;
                };
                let posting_len = partition
                    .inverted_list
                    .loaded_posting_len(token_id)
                    .ok_or_else(|| {
                        Error::index(format!(
                            "FTS token '{term}' maps to invalid posting token id {token_id} in partition {}",
                            partition.id()
                        ))
                    })?;
                *term_docs = term_docs.checked_add(posting_len).ok_or_else(|| {
                    Error::index(format!(
                        "FTS document frequency for term '{term}' overflows usize"
                    ))
                })?;
            }
        }
        Ok(Some(LoadedTermStats {
            stats: (total_tokens, num_docs, token_docs),
            token_ids: token_ids.into_boxed_slice(),
        }))
    }

    /// Corpus statistics at row granularity: `(total_tokens, docCount,
    /// per_term_docFreq)` counted over distinct row ids.
    ///
    /// Cross-field BM25F scores rows, not documents: `combined_fields_search`
    /// sums all of a row's postings into one `tf'` and all of its document
    /// lengths into one `dl'`. Released V1/V2 indexes indexed each
    /// `List<String>` element as its own document, so pairing that row-granular
    /// `tf'`/`dl'` with [`Self::bm25_stats_for_terms`] would measure it against a
    /// corpus of a different size: `idf'` and `avgdl'` would describe elements
    /// while the frequencies describe rows, shifting an old index's top-k relative
    /// to the same data reindexed on a current build.
    ///
    /// V3 indexes always map one row to one document, so they delegate to
    /// [`Self::bm25_stats_for_terms`] and keep its IO profile: no `DocSet` load,
    /// one posting-metadata row per term and partition. `total_tokens` is
    /// granularity-independent, so it comes from the shared
    /// `aggregate_corpus_stats` cache either way.
    ///
    /// `metrics`, when provided, receives the posting-metadata cache lookups these
    /// statistics trigger, so both granularities report comparable
    /// `index_cache_hits`/`index_cache_misses`.
    pub async fn bm25_row_stats_for_terms(
        &self,
        terms: &[String],
        stale_rows: Option<&RowAddrTreeMap>,
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<(u64, usize, Vec<usize>)> {
        let base = self.bm25_row_stats_for_terms_unrestricted(terms, metrics);
        let Some(stale_rows) = stale_rows.filter(|rows| !rows.is_empty()) else {
            return base.await;
        };
        let (total_tokens, num_rows, mut row_freqs) = base.await?;
        let (stale_tokens, stale_num_rows, stale_freqs) = self
            .stale_row_stats_for_terms(terms, stale_rows, metrics)
            .await?;
        for (total, stale) in row_freqs.iter_mut().zip(stale_freqs) {
            *total = total.saturating_sub(stale);
        }
        Ok((
            total_tokens.saturating_sub(stale_tokens),
            num_rows.saturating_sub(stale_num_rows),
            row_freqs,
        ))
    }

    /// What the rows an overlay made stale still contribute to this index's
    /// statistics, for the caller to subtract.
    ///
    /// The index holds their pre-overlay values until it is optimized, and the flat
    /// scan folds in their current ones. Counting both inflates `docCount` and
    /// `docFreq`, which is enough to reorder results after an overlay that changed
    /// no searchable value.
    async fn stale_row_stats_for_terms(
        &self,
        terms: &[String],
        stale_rows: &RowAddrTreeMap,
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<(u64, usize, Vec<usize>)> {
        let io_parallelism = self.store.io_parallelism();
        let futures = self
            .partitions
            .iter()
            .map(|partition| {
                let partition = partition.clone();
                async move {
                    partition
                        .stale_row_stats_for_terms(terms, stale_rows, metrics)
                        .await
                }
            })
            .collect::<Vec<_>>();
        let per_partition: Vec<(u64, usize, Vec<usize>)> = stream::iter(futures)
            .buffer_unordered(io_parallelism)
            .try_collect()
            .await?;

        let mut total_tokens = 0u64;
        let mut num_rows = 0usize;
        let mut row_freqs = vec![0usize; terms.len()];
        // Partitions hold disjoint row sets, so a row lands in at most one of these.
        for (partition_tokens, partition_rows, partition_freqs) in per_partition {
            total_tokens += partition_tokens;
            num_rows += partition_rows;
            for (total, partition_freq) in row_freqs.iter_mut().zip(partition_freqs) {
                *total += partition_freq;
            }
        }
        Ok((total_tokens, num_rows, row_freqs))
    }

    /// [`Self::bm25_row_stats_for_terms`] with nothing subtracted.
    async fn bm25_row_stats_for_terms_unrestricted(
        &self,
        terms: &[String],
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<(u64, usize, Vec<usize>)> {
        if !matches!(
            self.format_version,
            InvertedListFormatVersion::V1 | InvertedListFormatVersion::V2
        ) {
            return self.bm25_stats_for_terms(terms, metrics).await;
        }

        let (total_tokens, _) = self.aggregate_corpus_stats().await?;
        let io_parallelism = self.store.io_parallelism();
        let futures = self
            .partitions
            .iter()
            .map(|partition| {
                let partition = partition.clone();
                async move { partition.row_stats_for_terms(terms, metrics).await }
            })
            .collect::<Vec<_>>();
        let per_partition: Vec<(usize, Vec<usize>)> = stream::iter(futures)
            .buffer_unordered(io_parallelism)
            .try_collect()
            .await?;

        let mut num_rows = 0usize;
        let mut row_freqs = vec![0usize; terms.len()];
        for (partition_rows, partition_freqs) in per_partition {
            // Partitions hold disjoint row sets, so per-partition distinct counts
            // add up to the index's distinct count.
            num_rows += partition_rows;
            for (total, partition_freq) in row_freqs.iter_mut().zip(partition_freqs) {
                *total += partition_freq;
            }
        }
        Ok((total_tokens, num_rows, row_freqs))
    }

    /// Aggregate immutable per-partition corpus statistics.  New modern files
    /// read both values from the already-opened docs footer; older partitioned
    /// files scan `_num_tokens` once as a compatibility fallback.
    pub(super) async fn aggregate_corpus_stats(&self) -> Result<(u64, usize)> {
        self.corpus_stats
            .get_or_try_init(|| async {
                let io_parallelism = self.store.io_parallelism();
                let futures = self
                    .partitions
                    .iter()
                    .map(|p| {
                        let part = p.clone();
                        async move { part.docs.stats().await }
                    })
                    .collect::<Vec<_>>();
                let stats = stream::iter(futures)
                    .buffer_unordered(io_parallelism)
                    .try_collect::<Vec<_>>()
                    .await?;
                let mut total_tokens = 0_u64;
                let mut num_docs = 0_usize;
                for stat in stats {
                    total_tokens = total_tokens
                        .checked_add(stat.total_tokens)
                        .ok_or_else(|| Error::index("FTS corpus token count overflows u64"))?;
                    num_docs = num_docs
                        .checked_add(stat.num_docs)
                        .ok_or_else(|| Error::index("FTS corpus document count overflows usize"))?;
                }
                Ok((total_tokens, num_docs))
            })
            .await
            .copied()
    }

    /// Sum the posting-list length for `term` across this index's partitions
    /// via single-row reads, with partition lookups bounded by the store's
    /// `io_parallelism()`.
    async fn df_for_term(
        &self,
        term: &str,
        metrics: Option<&dyn MetricsCollector>,
    ) -> Result<usize> {
        let io_parallelism = self.store.io_parallelism();
        let futures = self
            .partitions
            .iter()
            .map(|part| {
                let part = part.clone();
                async move {
                    match part.tokens.get(term) {
                        Some(token_id) => {
                            part.inverted_list
                                .posting_len_for_token(token_id, metrics)
                                .await
                        }
                        None => Ok(0),
                    }
                }
            })
            .collect::<Vec<_>>();
        let dfs: Vec<usize> = stream::iter(futures)
            .buffer_unordered(io_parallelism)
            .try_collect()
            .await?;
        Ok(dfs.into_iter().sum())
    }

    /// Expand fuzzy query tokens against all partitions in this segment.
    ///
    /// `params.max_expansions` caps the whole query's expansion, not any
    /// single partition's: source terms at the same query position and their
    /// per-partition candidates merge into one lexicographically ordered set,
    /// and the remaining budget takes a prefix of it. The selected terms are
    /// a pure function of the segment's vocabulary, so changing source-token
    /// order or splitting the same corpus into more partitions cannot change
    /// which terms a fuzzy query matches.
    pub fn expand_fuzzy_tokens(&self, tokens: &Tokens, params: &FtsSearchParams) -> Result<Tokens> {
        let initial_capacity = tokens.len().min(params.max_expansions);
        let mut expanded_tokens = Vec::with_capacity(initial_capacity);
        let mut expanded_positions = Vec::with_capacity(initial_capacity);
        let mut seen = HashSet::new();
        let mut source_terms_by_position = BTreeMap::<u32, Vec<&str>>::new();
        for token_idx in 0..tokens.len() {
            source_terms_by_position
                .entry(tokens.position(token_idx))
                .or_default()
                .push(tokens.get_token(token_idx));
        }
        for (position, source_terms) in source_terms_by_position {
            let remaining = params.max_expansions.saturating_sub(expanded_tokens.len());
            if remaining == 0 {
                break;
            }
            // Each partition contributes at most its `remaining`
            // lexicographically smallest candidates, so the global
            // lex-smallest `remaining` selection below is unaffected by the
            // per-partition truncation.
            let mut candidates = BTreeSet::new();
            let mut seen_source_terms = HashSet::new();
            for source_term in source_terms {
                if !seen_source_terms.insert(source_term) {
                    continue;
                }
                let automaton = FuzzyAutomaton::new(source_term, tokens.token_type(), params)?;
                self.collect_fuzzy_candidates_with_automaton(
                    &automaton,
                    remaining,
                    &mut candidates,
                )?;
            }
            for candidate in candidates {
                if expanded_tokens.len() >= params.max_expansions {
                    break;
                }
                if seen.insert((candidate.clone(), position)) {
                    expanded_tokens.push(candidate);
                    expanded_positions.push(position);
                }
            }
        }
        Ok(Tokens::with_positions(
            expanded_tokens,
            expanded_positions,
            tokens.token_type().clone(),
        ))
    }

    /// Search documents that match the query and return row ids sorted by BM25 score.
    ///
    /// When `base_scorer` is provided for an exact query, search uses those
    /// corpus-level BM25 statistics instead of deriving them from this segment
    /// alone. Fuzzy queries must use [`Self::bm25_search_prepared`], because a
    /// scorer alone does not identify the canonical capped vocabulary.
    #[instrument(level = "debug", skip_all)]
    pub async fn bm25_search(
        &self,
        tokens: Arc<Tokens>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        base_scorer: Option<&MemBM25Scorer>,
    ) -> Result<(Vec<u64>, Vec<f32>)> {
        let documents = self
            .bm25_search_documents(tokens, params, operator, prefilter, metrics, base_scorer)
            .await?;
        Ok(documents
            .into_iter()
            .map(|document| (document.row_id, document.score.0))
            .unzip())
    }

    /// Search logical FTS documents, retaining element coordinates when present.
    ///
    /// A scorer-only override is valid for exact queries. Fuzzy callers must
    /// use [`Self::bm25_search_prepared_documents`] so the canonical vocabulary
    /// and its statistics cannot diverge.
    #[instrument(level = "debug", skip_all)]
    pub async fn bm25_search_documents(
        &self,
        tokens: Arc<Tokens>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        base_scorer: Option<&MemBM25Scorer>,
    ) -> Result<Vec<ScoredDoc>> {
        self.bm25_search_documents_impl(
            tokens,
            params,
            operator,
            prefilter,
            metrics,
            base_scorer,
            None,
        )
        .await
    }

    /// Search logical FTS documents with an exclusive initial raw-score floor.
    ///
    /// This is an internal optimization hook for a repeated bounded WAND pass.
    /// Callers must round the floor down far enough to retain every candidate
    /// equal to their inclusive logical threshold.
    #[doc(hidden)]
    #[instrument(level = "debug", skip_all)]
    #[allow(clippy::too_many_arguments)]
    pub async fn bm25_search_documents_with_score_floor(
        &self,
        tokens: Arc<Tokens>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        base_scorer: Option<&MemBM25Scorer>,
        initial_score_floor: f32,
    ) -> Result<Vec<ScoredDoc>> {
        if self.is_legacy() {
            return Err(Error::invalid_input(
                "an initial Match WAND score floor requires a modern FTS index",
            ));
        }
        if !initial_score_floor.is_finite() {
            return Err(Error::invalid_input(format!(
                "initial Match WAND score floor must be finite, got {initial_score_floor}"
            )));
        }
        self.bm25_search_documents_impl(
            tokens,
            params,
            operator,
            prefilter,
            metrics,
            base_scorer,
            Some(initial_score_floor),
        )
        .await
    }

    #[allow(clippy::too_many_arguments)]
    async fn bm25_search_documents_impl(
        &self,
        tokens: Arc<Tokens>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        base_scorer: Option<&MemBM25Scorer>,
        initial_score_floor: Option<f32>,
    ) -> Result<Vec<ScoredDoc>> {
        if base_scorer.is_some() && uses_fuzzy_expansion(params.fuzziness) {
            return Err(Error::invalid_input(
                "fuzzy BM25 search cannot use an injected scorer without its prepared vocabulary; use bm25_search_prepared or bm25_search_prepared_documents",
            ));
        }
        // Fuzzy expansion runs once here, with the global `max_expansions`
        // budget, instead of once per partition: partitions receive the
        // final token list, so the matched terms cannot depend on how the
        // corpus happens to be partitioned.
        let tokens = if uses_fuzzy_expansion(params.fuzziness) {
            let expanded = Arc::new(self.expand_fuzzy_tokens(tokens.as_ref(), params.as_ref())?);
            if operator == Operator::And || params.phrase_slop.is_some() {
                // AND/phrase semantics require every original token position
                // to keep at least one expansion; a position that expands to
                // nothing anywhere in the segment can never be matched.
                let surviving = (0..expanded.len())
                    .map(|idx| expanded.position(idx))
                    .collect::<HashSet<_>>();
                if (0..tokens.len()).any(|idx| !surviving.contains(&tokens.position(idx))) {
                    return Ok(Vec::new());
                }
            }
            expanded
        } else {
            tokens
        };

        let local_scorer;
        let scorer: &MemBM25Scorer = if let Some(base_scorer) = base_scorer {
            base_scorer
        } else {
            local_scorer = self
                .bm25_scorer_for_final_tokens(tokens.as_ref(), Some(metrics.as_ref()))
                .await?;
            &local_scorer
        };
        self.bm25_search_final_documents(
            tokens,
            params,
            operator,
            prefilter,
            metrics,
            scorer,
            None,
            initial_score_floor,
            None,
        )
        .await
    }

    /// Search with a vocabulary/scorer pair prepared once across every
    /// physical segment. No fuzzy expansion or local scorer construction is
    /// permitted below this boundary.
    #[doc(hidden)]
    pub async fn bm25_search_prepared(
        &self,
        prepared: Arc<crate::scalar::inverted::PreparedBm25Query>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
    ) -> Result<(Vec<u64>, Vec<f32>)> {
        let documents = self
            .bm25_search_prepared_documents(prepared, params, operator, prefilter, metrics)
            .await?;
        Ok(documents
            .into_iter()
            .map(|document| (document.row_id, document.score.0))
            .unzip())
    }

    /// Search logical FTS documents with a vocabulary/scorer pair prepared
    /// once across every physical segment.
    #[doc(hidden)]
    pub async fn bm25_search_prepared_documents(
        &self,
        prepared: Arc<crate::scalar::inverted::PreparedBm25Query>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
    ) -> Result<Vec<ScoredDoc>> {
        self.bm25_search_prepared_documents_impl(
            prepared, params, operator, prefilter, metrics, None,
        )
        .await
    }

    /// Search logical FTS documents with a prepared vocabulary/scorer pair and
    /// an exclusive initial raw-score floor.
    #[doc(hidden)]
    #[allow(clippy::too_many_arguments)]
    pub async fn bm25_search_prepared_documents_with_score_floor(
        &self,
        prepared: Arc<crate::scalar::inverted::PreparedBm25Query>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        initial_score_floor: f32,
    ) -> Result<Vec<ScoredDoc>> {
        if self.is_legacy() {
            return Err(Error::invalid_input(
                "an initial Match WAND score floor requires a modern FTS index",
            ));
        }
        if !initial_score_floor.is_finite() {
            return Err(Error::invalid_input(format!(
                "initial Match WAND score floor must be finite, got {initial_score_floor}"
            )));
        }
        self.bm25_search_prepared_documents_impl(
            prepared,
            params,
            operator,
            prefilter,
            metrics,
            Some(initial_score_floor),
        )
        .await
    }

    #[allow(clippy::too_many_arguments)]
    async fn bm25_search_prepared_documents_impl(
        &self,
        prepared: Arc<crate::scalar::inverted::PreparedBm25Query>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        initial_score_floor: Option<f32>,
    ) -> Result<Vec<ScoredDoc>> {
        if (operator == Operator::And || params.phrase_slop.is_some())
            && !prepared.has_all_query_positions()
        {
            return Ok(Vec::new());
        }
        let scorer = prepared.scorer();
        let reusable_scorer = prepared.reusable_scorer();
        let term_ids = prepared
            .term_ids()
            .and_then(|term_ids| term_ids.for_segment(self, prepared.tokens().len()));
        self.bm25_search_final_documents(
            prepared.tokens().clone(),
            params,
            operator,
            prefilter,
            metrics,
            scorer.as_ref(),
            reusable_scorer,
            initial_score_floor,
            term_ids,
        )
        .await
    }

    // This is the boundary between query preparation and final document search;
    // each argument is an independent prepared input consumed by both legacy and
    // modern implementations.
    #[allow(clippy::too_many_arguments)]
    async fn bm25_search_final_documents(
        &self,
        tokens: Arc<Tokens>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        prefilter: Arc<dyn PreFilter>,
        metrics: Arc<dyn MetricsCollector>,
        scorer: &MemBM25Scorer,
        prepared_scorer: Option<&Arc<MemBM25Scorer>>,
        initial_score_floor: Option<f32>,
        term_ids: Option<SegmentTermIds<'_>>,
    ) -> Result<Vec<ScoredDoc>> {
        // The wand only consults `scorer.doc_weight`, which is metadata-free.
        // Candidate rescoring also consults `query_weight`. The impact scorer
        // is the final scorer or a clone of it, so either may rescore; pairing
        // final tokens with its precomputed per-term IDFs avoids the v2 bulk
        // metadata pull and keeps scoring aligned with the rewrite.
        let impact_scorer = select_impact_scorer(
            scorer,
            prepared_scorer,
            || self.is_legacy(),
            || *LANCE_FTS_REUSE_PREPARED_SCORER_ENABLED,
        );

        let limit = params.limit.unwrap_or(usize::MAX);
        if limit == 0 {
            return Ok(Vec::new());
        }
        let mask = prefilter.mask();
        if self.is_legacy() {
            let (row_ids, scores) = self
                .bm25_search_legacy(
                    tokens,
                    params,
                    operator,
                    mask,
                    metrics,
                    scorer,
                    impact_scorer,
                    limit,
                )
                .await?;
            Ok(row_ids
                .into_iter()
                .zip(scores)
                .map(|(row_id, score)| ScoredDoc::new(row_id, score))
                .collect())
        } else {
            self.bm25_search_modern(ModernSearchRequest {
                tokens,
                params,
                operator,
                mask,
                metrics,
                impact_scorer,
                limit,
                initial_score_floor,
                term_ids,
            })
            .await
        }
    }

    #[allow(clippy::too_many_arguments)]
    pub(super) async fn bm25_search_legacy(
        &self,
        tokens: Arc<Tokens>,
        params: Arc<FtsSearchParams>,
        operator: Operator,
        mask: Arc<RowAddrMask>,
        metrics: Arc<dyn MetricsCollector>,
        scorer: &MemBM25Scorer,
        impact_scorer: Arc<MemBM25Scorer>,
        limit: usize,
    ) -> Result<(Vec<u64>, Vec<f32>)> {
        let impact_shared_threshold = Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits()));
        let io_parallelism = self.store.io_parallelism();
        let parts = self
            .partitions
            .chunks(fts_search_chunk())
            .map(|chunk| {
                let chunk = chunk.to_vec();
                let tokens = tokens.clone();
                let params = params.clone();
                let mask = mask.clone();
                let metrics = metrics.clone();
                let impact_scorer = impact_scorer.clone();
                let impact_shared_threshold = impact_shared_threshold.clone();
                async move {
                    let loads = chunk.into_iter().map(|part| {
                        let tokens = tokens.clone();
                        let params = params.clone();
                        let metrics = metrics.clone();
                        let impact_scorer = impact_scorer.clone();
                        let impact_shared_threshold = impact_shared_threshold.clone();
                        async move {
                            let LoadedPostings {
                                postings,
                                grouped_expansions,
                                impact_safe,
                                exact_scoring_required,
                                ..
                            } = part
                                .load_posting_lists(
                                    tokens.as_ref(),
                                    params.as_ref(),
                                    operator,
                                    impact_scorer.as_ref(),
                                    metrics.as_ref(),
                                    false,
                                )
                                .await?;
                            if postings.is_empty() {
                                return Result::Ok(None);
                            }
                            let docs = part.docs.legacy().cloned().ok_or_else(|| {
                                Error::internal("legacy index contains modern partition documents")
                            })?;
                            let use_global_scorer = impact_safe || exact_scoring_required;
                            let threshold = if use_global_scorer {
                                impact_shared_threshold
                            } else {
                                Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits()))
                            };
                            let wand_scorer = use_global_scorer.then(|| impact_scorer.clone());
                            Result::Ok(Some((
                                part,
                                docs,
                                postings,
                                wand_scorer,
                                threshold,
                                grouped_expansions,
                            )))
                        }
                    });
                    let loaded = stream::iter(loads)
                        .buffer_unordered(io_parallelism)
                        .try_collect::<Vec<_>>()
                        .await?
                        .into_iter()
                        .flatten()
                        .collect::<Vec<_>>();
                    if loaded.is_empty() {
                        return Result::Ok(Vec::new());
                    }

                    let results = spawn_cpu(move || {
                        let mut results = Vec::with_capacity(loaded.len());
                        for (part, docs, postings, wand_scorer, threshold, grouped_expansions) in
                            loaded
                        {
                            let idf_by_position =
                                idf_by_position(&postings, impact_scorer.as_ref());
                            let candidates = part.bm25_search_legacy(
                                docs.as_ref(),
                                params.as_ref(),
                                operator,
                                mask.as_ref(),
                                postings,
                                wand_scorer,
                                metrics.as_ref(),
                                threshold,
                            )?;
                            results.push(PartitionCandidates {
                                idf_by_position,
                                grouped_expansions,
                                candidates,
                            });
                        }
                        Result::Ok(results)
                    })
                    .await?;
                    Result::Ok(results)
                }
            })
            .collect::<Vec<_>>();

        let mut ranked = BinaryHeap::new();
        let mut parts = stream::iter(parts)
            .buffer_unordered(get_num_compute_intensive_cpus().min(32))
            .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok)))
            .try_flatten();
        while let Some(partition) = parts.try_next().await? {
            for (row_id, score) in rescore_partition_candidates(partition, scorer) {
                push_scored_key(&mut ranked, limit, row_id, score);
            }
        }
        Ok(ranked
            .into_sorted_vec()
            .into_iter()
            .map(|Reverse(doc)| (doc.row_id, doc.score.0))
            .unzip())
    }

    pub(super) async fn bm25_search_modern(
        &self,
        request: ModernSearchRequest<'_>,
    ) -> Result<Vec<ScoredDoc>> {
        // Select a concrete completion path before candidate search.  The
        // fully resident future never builds deferred address-read state, while
        // a cold query keeps DocIds until its final bounded I/O phase.
        if self.has_resident_document_projections() {
            self.bm25_search_modern_resident(request).await
        } else {
            self.bm25_search_modern_deferred(request).await
        }
    }

    pub(super) fn has_resident_document_projections(&self) -> bool {
        if self.document_projections_resident.load(Ordering::Acquire) {
            return true;
        }
        let resident = self.document_projections_resident_now();
        if resident {
            self.document_projections_resident
                .store(true, Ordering::Release);
        }
        resident
    }

    pub(super) fn document_projections_resident_now(&self) -> bool {
        self.partitions.iter().all(|partition| {
            partition
                .docs
                .modern()
                .is_some_and(|documents| documents.projection_resident())
        })
    }

    async fn bm25_search_modern_resident(
        &self,
        request: ModernSearchRequest<'_>,
    ) -> Result<Vec<ScoredDoc>> {
        let ranked = self.bm25_search_modern_candidates(request).await?;
        if let Some(result) = self.resolve_resident_modern_candidates(&ranked)? {
            return Ok(result);
        }
        self.document_projections_resident
            .store(false, Ordering::Release);
        self.resolve_deferred_modern_candidates(ranked).await
    }

    async fn bm25_search_modern_deferred(
        &self,
        request: ModernSearchRequest<'_>,
    ) -> Result<Vec<ScoredDoc>> {
        // Old partitioned files without persisted stats populate their
        // fallback stats before deferred candidate orchestration.  A resident
        // search can skip this full-index synchronization: standard prewarm
        // has already initialized it, while any independently resident
        // partition loads its lengths before constructing a local scorer.
        if self.corpus_stats.get().is_none() {
            self.aggregate_corpus_stats().await?;
        }
        // For new-format indexes, aggregate_corpus_stats reads corpus stats from
        // persisted schema metadata (O(1)) without loading doc lengths as a side
        // effect. Pre-load lengths in parallel now only for partitions whose
        // dictionaries can satisfy the query leaf, so required-token misses avoid
        // reading the full length column. The scoring phase then gets cache hits
        // instead of issuing sequential per-partition IO.
        let io_parallelism = self.store.io_parallelism();
        let uncached_lengths = self
            .partitions
            .iter()
            .filter_map(|part| {
                let docs = part.docs.modern()?.clone();
                if docs.cached_lengths().is_some() {
                    return None;
                }
                part.may_match_tokens(
                    request.tokens.as_ref(),
                    request.operator,
                    request.params.phrase_slop.is_some(),
                )
                .then_some(async move { docs.lengths().await.map(|_| ()) })
            })
            .collect::<Vec<_>>();
        if !uncached_lengths.is_empty() {
            stream::iter(uncached_lengths)
                .buffer_unordered(io_parallelism)
                .try_collect::<Vec<_>>()
                .await?;
        }
        let ranked = self.bm25_search_modern_candidates(request).await?;
        self.resolve_deferred_modern_candidates(ranked).await
    }

    async fn bm25_search_modern_candidates(
        &self,
        request: ModernSearchRequest<'_>,
    ) -> Result<Vec<Reverse<ScoredPartitionDoc>>> {
        let ModernSearchRequest {
            tokens,
            params,
            operator,
            mask,
            metrics,
            impact_scorer,
            limit,
            initial_score_floor,
            term_ids,
        } = request;
        if self.partitions.len() > u32::MAX as usize {
            return Err(Error::index(format!(
                "FTS partition count {} exceeds candidate identity capacity",
                self.partitions.len()
            )));
        }
        let impact_shared_threshold = Arc::new(AtomicU32::new(
            initial_score_floor.unwrap_or(f32::NEG_INFINITY).to_bits(),
        ));
        let shared_norm_addends = SharedNormAddends::default();
        let io_parallelism = self.store.io_parallelism();
        let is_phrase_query = params.phrase_slop.is_some();
        // With dictionary lookups already resolved, drop partitions that
        // cannot satisfy the leaf before building any per-partition state.
        let searched_partitions = self
            .partitions
            .iter()
            .enumerate()
            .filter(|(partition_ordinal, _)| {
                term_ids.is_none_or(|term_ids| {
                    token_ids_may_match(
                        tokens.as_ref(),
                        term_ids.partition(*partition_ordinal),
                        operator,
                        is_phrase_query,
                    )
                })
            })
            .collect::<Vec<_>>();
        let parts = searched_partitions
            .chunks(fts_search_chunk())
            .map(|chunk| {
                let chunk = chunk
                    .iter()
                    .map(|&(partition_ordinal, part)| (partition_ordinal, Arc::clone(part)))
                    .collect::<Vec<_>>();
                let tokens = tokens.clone();
                let params = params.clone();
                let mask = mask.clone();
                let metrics = metrics.clone();
                let impact_scorer = impact_scorer.clone();
                let impact_shared_threshold = impact_shared_threshold.clone();
                let shared_norm_addends = shared_norm_addends.clone();
                async move {
                    let loads = chunk.into_iter().map(|(partition_ordinal, part)| {
                        let tokens = tokens.clone();
                        let params = params.clone();
                        let mask = mask.clone();
                        let metrics = metrics.clone();
                        let partition_term_ids =
                            term_ids.map(|term_ids| term_ids.partition(partition_ordinal));
                        async move {
                            let Some(postings) = part
                                .fetch_posting_lists(
                                    tokens.as_ref(),
                                    params.as_ref(),
                                    operator,
                                    metrics.as_ref(),
                                    PostingLoadOptions::read_ahead(false),
                                    partition_term_ids,
                                )
                                .await?
                            else {
                                return Result::Ok(None);
                            };
                            let documents = part.docs.modern().ok_or_else(|| {
                                Error::internal("modern index contains legacy partition documents")
                            })?;
                            let materialize_selected = operator == Operator::Or
                                && mask.max_len().is_some_and(|selected| {
                                    u128::from(selected).saturating_mul(100)
                                        <= u128::from(*FLAT_SEARCH_PERCENT_THRESHOLD)
                                            .saturating_mul(documents.len() as u128)
                                });
                            let visibility = match documents
                                .immediate_visibility(mask.clone(), materialize_selected)
                            {
                                Some(visibility) => visibility,
                                None => {
                                    documents
                                        .visibility(mask.clone(), materialize_selected)
                                        .await?
                                }
                            };
                            if visibility.is_empty() {
                                return Result::Ok(None);
                            }
                            let lengths = match documents.cached_lengths() {
                                Some(lengths) => lengths,
                                None => documents.lengths().await?,
                            };
                            Result::Ok(Some((
                                partition_ordinal,
                                part,
                                lengths,
                                visibility,
                                postings,
                            )))
                        }
                    });
                    let loaded = stream::iter(loads)
                        .buffer_unordered(io_parallelism)
                        .try_collect::<Vec<_>>()
                        .await?
                        .into_iter()
                        .flatten()
                        .collect::<Vec<_>>();
                    if loaded.is_empty() {
                        return Result::Ok(Vec::new());
                    }

                    // Iterators are built here rather than during the load so
                    // their decode buffers are allocated and freed by the same
                    // thread that runs WAND.
                    let results = spawn_cpu(move || {
                        let mut results = Vec::with_capacity(loaded.len());
                        for (partition_ordinal, part, lengths, visibility, postings) in loaded {
                            let LoadedPostings {
                                postings,
                                grouped_expansions,
                                impact_safe,
                                exact_scoring_required,
                                ..
                            } = postings.into_loaded(tokens.as_ref(), impact_scorer.as_ref())?;
                            if postings.is_empty() {
                                continue;
                            }
                            let use_global_scorer = impact_safe || exact_scoring_required;
                            let threshold = if use_global_scorer {
                                impact_shared_threshold.clone()
                            } else {
                                Arc::new(AtomicU32::new(f32::NEG_INFINITY.to_bits()))
                            };
                            let wand_scorer = use_global_scorer.then(|| impact_scorer.clone());
                            let idf_by_position =
                                idf_by_position(&postings, impact_scorer.as_ref());
                            let candidates = part.bm25_search_modern(
                                lengths.as_ref(),
                                &visibility,
                                params.as_ref(),
                                operator,
                                postings,
                                wand_scorer,
                                shared_norm_addends.clone(),
                                metrics.as_ref(),
                                threshold,
                            )?;
                            // Rescore here too, so each candidate's frequency
                            // list is freed by the thread that allocated it.
                            let scored = rescore_partition_candidates(
                                PartitionCandidates {
                                    idf_by_position,
                                    grouped_expansions,
                                    candidates,
                                },
                                impact_scorer.as_ref(),
                            );
                            if !scored.is_empty() {
                                results.push((partition_ordinal, scored));
                            }
                        }
                        Result::Ok(results)
                    })
                    .await?;
                    Result::Ok(results)
                }
            })
            .collect::<Vec<_>>();

        let mut ranked = BinaryHeap::new();
        let mut parts = stream::iter(parts)
            .buffer_unordered(get_num_compute_intensive_cpus().min(32))
            .map_ok(|results| stream::iter(results.into_iter().map(Result::Ok)))
            .try_flatten();
        while let Some((partition_ordinal, scored)) = parts.try_next().await? {
            for (doc_id, score) in scored {
                push_scored_partition_doc(
                    &mut ranked,
                    limit,
                    PartitionDocId::try_new(partition_ordinal, doc_id)?,
                    score,
                );
            }
        }

        Ok(ranked.into_sorted_vec())
    }

    fn resolve_resident_modern_candidates(
        &self,
        ranked: &[Reverse<ScoredPartitionDoc>],
    ) -> Result<Option<Vec<ScoredDoc>>> {
        if self.partitions.iter().any(|partition| {
            partition
                .docs
                .modern()
                .is_some_and(|documents| documents.coordinate_rank() > 0)
        }) {
            return Ok(None);
        }
        let mut resolved_documents = ranked
            .iter()
            .map(|Reverse(candidate)| ScoredDoc::new(0, candidate.score.0))
            .collect::<Vec<_>>();
        let mut by_partition = BTreeMap::<usize, Vec<(usize, DocId)>>::new();
        for (rank, Reverse(candidate)) in ranked.iter().enumerate() {
            let partition_ordinal = candidate.document.partition_ordinal();
            let doc_id = candidate.document.doc_id;
            by_partition
                .entry(partition_ordinal)
                .or_default()
                .push((rank, doc_id));
        }
        for (partition_ordinal, entries) in by_partition {
            let documents = self
                .partitions
                .get(partition_ordinal)
                .and_then(|partition| partition.docs.modern())
                .ok_or_else(|| {
                    Error::internal(format!(
                        "resident FTS candidates reference missing modern partition ordinal {partition_ordinal}"
                    ))
                })?;
            let doc_ids = entries
                .iter()
                .map(|(_, doc_id)| *doc_id)
                .collect::<Vec<_>>();
            let Some(resolved) = documents.cached_row_addresses(&doc_ids)? else {
                return Ok(None);
            };
            for ((rank, _), address) in entries.into_iter().zip(resolved) {
                resolved_documents[rank].row_id = address;
            }
        }
        Ok(Some(resolved_documents))
    }

    async fn resolve_deferred_modern_candidates(
        &self,
        ranked: Vec<Reverse<ScoredPartitionDoc>>,
    ) -> Result<Vec<ScoredDoc>> {
        let mut resolved_documents = ranked
            .iter()
            .map(|Reverse(candidate)| ScoredDoc::new(0, candidate.score.0))
            .collect::<Vec<_>>();
        let mut by_partition = BTreeMap::<usize, Vec<(usize, DocId)>>::new();
        for (rank, Reverse(candidate)) in ranked.iter().enumerate() {
            let partition_ordinal = candidate.document.partition_ordinal();
            let doc_id = candidate.document.doc_id;
            by_partition
                .entry(partition_ordinal)
                .or_default()
                .push((rank, doc_id));
        }
        let mut address_reads = Vec::with_capacity(by_partition.len());
        let mut largest_read_bytes = 0;
        for (partition_ordinal, entries) in by_partition {
            let documents = self
                .partitions
                .get(partition_ordinal)
                .and_then(|partition| partition.docs.modern())
                .cloned()
                .ok_or_else(|| {
                    Error::internal(format!(
                        "deferred FTS candidates reference missing modern partition ordinal {partition_ordinal}"
                    ))
                })?;
            let doc_ids = entries
                .iter()
                .map(|(_, doc_id)| *doc_id)
                .collect::<Vec<_>>();
            largest_read_bytes =
                largest_read_bytes.max(documents.estimated_address_read_bytes(&doc_ids));
            address_reads.push(async move {
                let resolved = documents.resolve_document_keys(&doc_ids).await?;
                Result::Ok((entries, resolved))
            });
        }
        let concurrency = address_read_concurrency(self.store.io_parallelism(), largest_read_bytes);
        let mut address_reads = stream::iter(address_reads).buffer_unordered(concurrency);
        while let Some((entries, resolved)) = address_reads.try_next().await? {
            for ((rank, _), (row_id, doc_index)) in entries.into_iter().zip(resolved) {
                resolved_documents[rank].row_id = row_id;
                resolved_documents[rank].doc_index = doc_index;
            }
        }
        Ok(resolved_documents)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn scorer() -> Arc<MemBM25Scorer> {
        Arc::new(MemBM25Scorer::new(10, 2, HashMap::new()))
    }

    #[test]
    fn test_reuse_prepared_scorer_kill_switch_parser() {
        for value in [None, Some(""), Some("1"), Some("on"), Some("false")] {
            assert!(
                reuse_prepared_scorer_enabled_from_value(value),
                "unexpected disable: {value:?}"
            );
        }
        for value in [Some("0"), Some("off"), Some("OFF"), Some(" off ")] {
            assert!(
                !reuse_prepared_scorer_enabled_from_value(value),
                "unexpected enable: {value:?}"
            );
        }
    }

    #[test]
    fn test_select_impact_scorer_reuses_enabled_prepared_arc() {
        let prepared = scorer();
        let impact = select_impact_scorer(prepared.as_ref(), Some(&prepared), || false, || true);

        assert!(Arc::ptr_eq(&impact, &prepared));
        assert_eq!(Arc::strong_count(&prepared), 2);
    }

    #[test]
    fn test_select_impact_scorer_clones_when_disabled() {
        let prepared = scorer();
        let unrelated = scorer();
        let impact = select_impact_scorer(unrelated.as_ref(), Some(&prepared), || false, || false);

        assert!(!Arc::ptr_eq(&impact, &unrelated));
        assert_eq!(Arc::strong_count(&prepared), 1);
        assert_eq!(Arc::strong_count(&unrelated), 1);
        assert_eq!(Arc::strong_count(&impact), 1);
    }

    #[test]
    fn test_select_impact_scorer_clones_legacy_without_reading_switch() {
        let prepared = scorer();
        let unrelated = scorer();
        let impact = select_impact_scorer(
            unrelated.as_ref(),
            Some(&prepared),
            || true,
            || panic!("legacy search must not read the reuse switch"),
        );

        assert!(!Arc::ptr_eq(&impact, &unrelated));
        assert_eq!(Arc::strong_count(&prepared), 1);
        assert_eq!(Arc::strong_count(&unrelated), 1);
        assert_eq!(Arc::strong_count(&impact), 1);
    }

    #[test]
    fn test_select_impact_scorer_clones_nonprepared_without_reading_switch() {
        let scorer = scorer();
        let impact = select_impact_scorer(
            scorer.as_ref(),
            None,
            || panic!("nonprepared search must not inspect the index format"),
            || panic!("nonprepared search must not read the reuse switch"),
        );

        assert!(!Arc::ptr_eq(&impact, &scorer));
        assert_eq!(Arc::strong_count(&scorer), 1);
        assert_eq!(Arc::strong_count(&impact), 1);
    }
}