From 5aa20a989944312ad3efd7bc020847e2c3ed9a87 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 12:51:05 +0800 Subject: [PATCH 01/22] perf(fts): index residual compound rows once --- rust/lance-index/src/scalar/inverted.rs | 1 + .../src/scalar/inverted/compound.rs | 57 +++ rust/lance/src/dataset/mem_wal/index/fts.rs | 344 +++++++++++++++++- rust/lance/src/dataset/scanner.rs | 173 ++++++++- rust/lance/src/dataset/tests/dataset_index.rs | 144 ++++++++ rust/lance/src/io/exec/fts.rs | 272 +++++++++++++- 6 files changed, 961 insertions(+), 30 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted.rs b/rust/lance-index/src/scalar/inverted.rs index 6a95876d174..ec7c22021cc 100644 --- a/rust/lance-index/src/scalar/inverted.rs +++ b/rust/lance-index/src/scalar/inverted.rs @@ -27,6 +27,7 @@ pub use compound::{ compound_search, compound_search_prepared_match, compound_search_prepared_match_with_score_floor, compound_search_with_base_scorer, compound_search_with_base_scorer_and_score_floor, exclusive_scaled_score_floor, + materialized_compound_top_k, }; #[doc(hidden)] pub use cross_column::cross_column_compound_search; diff --git a/rust/lance-index/src/scalar/inverted/compound.rs b/rust/lance-index/src/scalar/inverted/compound.rs index ed774c6337f..2b6653aa40e 100644 --- a/rust/lance-index/src/scalar/inverted/compound.rs +++ b/rust/lance-index/src/scalar/inverted/compound.rs @@ -2162,6 +2162,41 @@ impl TopKCollector { } } +/// Evaluate a compound query over exact, materialized leaf result sets. +/// +/// This is the bridge used by query-local residual postings: it keeps Boolean, +/// Boost, and MultiMatch semantics in the same scorer tree as the on-disk +/// compound path while allowing a different posting source. +#[doc(hidden)] +pub fn materialized_compound_top_k( + query: &FtsQuery, + leaves: Vec>, + limit: usize, + metrics: &dyn MetricsCollector, +) -> Result<(Vec, Vec)> { + let mut leaf_count = 0; + let plan = CompoundScorerPlan::from_query(query, &mut leaf_count)?; + if leaf_count != leaves.len() { + return Err(Error::internal(format!( + "compound FTS planned {leaf_count} leaves but received {} materialized leaves", + leaves.len() + ))); + } + let mut scorers = leaves + .into_iter() + .map(|rows| { + let rows = rows + .into_iter() + .map(|(row_id, score)| ScoredRow { row_id, score }) + .collect(); + MaterializedScorer::try_new(rows).map(|scorer| Some(Box::new(scorer) as BoxScorer<'_>)) + }) + .collect::>>()?; + let mut scorer = plan.build(&mut scorers, metrics)?; + let rows = TopKCollector::new(limit).collect(scorer.as_mut())?; + Ok(rows.into_iter().map(|row| (row.row_id, row.score)).unzip()) +} + #[derive(Debug, Clone, Copy)] pub(super) enum DisjunctionScore { Sum, @@ -4414,6 +4449,7 @@ mod tests { use super::super::scorer::Scorer; use super::*; use crate::metrics::NoOpMetricsCollector; + use crate::scalar::inverted::query::MultiMatchQuery; fn rows(values: &[(u64, f32)]) -> Vec { values @@ -4426,6 +4462,27 @@ mod tests { Box::new(MaterializedScorer::try_new(rows(values)).unwrap()) } + #[test] + fn materialized_compound_top_k_preserves_multimatch_and_tie_order() { + let query = FtsQuery::MultiMatch(MultiMatchQuery { + match_queries: vec![ + MatchQuery::new("alpha".to_string()).with_column(Some("text".to_string())), + MatchQuery::new("alpha".to_string()).with_column(Some("text".to_string())), + ], + }); + let metrics = NoOpMetricsCollector; + let (row_ids, scores) = materialized_compound_top_k( + &query, + vec![vec![(7, 1.0), (3, 2.0)], vec![(7, 3.0), (5, 3.0)]], + 2, + &metrics, + ) + .unwrap(); + + assert_eq!(row_ids, vec![5, 7]); + assert_eq!(scores, vec![3.0, 3.0]); + } + fn zero_weight_wand<'a>( documents: &'a DocSet, scorer: Arc, diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index fb34b9a2f79..db09f68e821 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -51,14 +51,14 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; use arc_swap::ArcSwap; -use arrow_array::RecordBatch; +use arrow_array::{Array, RecordBatch, UInt64Array}; use crossbeam_skiplist::SkipMap; use fst::{Map, Streamer}; use lance_bitpacking::{BitPacker, BitPacker4x}; use lance_core::datatypes::Schema as LanceSchema; use lance_core::{Error, Result}; use lance_index::scalar::InvertedIndexParams; -use lance_index::scalar::inverted::query::{Operator, Tokens}; +use lance_index::scalar::inverted::query::{FtsQuery, Operator, Tokens}; use lance_index::scalar::inverted::tokenizer::document_tokenizer::{DocType, LanceTokenizer}; use lance_index::scalar::inverted::{DocSet, MemBM25Scorer, Scorer, TokenSet}; use lance_tokenizer::TokenStream; @@ -1256,7 +1256,45 @@ impl FtsMemIndex { self.insert_batch(batch, row_offset) } + /// Insert explicit, potentially non-contiguous rows while retaining + /// postings only for query terms. + /// The tokenizer still visits the complete document so BM25 document + /// length and corpus totals remain identical to a full index. + pub(crate) fn insert_with_row_ids_for_terms( + &self, + batch: &RecordBatch, + row_ids: &UInt64Array, + terms: &HashSet, + ) -> Result<()> { + if row_ids.len() != batch.num_rows() || row_ids.null_count() != 0 { + return Err(Error::invalid_input(format!( + "MemWAL FTS explicit row ids require {} non-null values, got len={} nulls={}", + batch.num_rows(), + row_ids.len(), + row_ids.null_count() + ))); + } + self.insert_batch_with_keys(batch, |row_index| Ok(row_ids.value(row_index)), Some(terms)) + } + fn insert_batch(&self, batch: &RecordBatch, row_offset: u64) -> Result<()> { + self.insert_batch_with_keys( + batch, + |row_index| { + row_offset + .checked_add(row_index as u64) + .ok_or_else(|| Error::invalid_input("MemWAL FTS row position overflow")) + }, + None, + ) + } + + fn insert_batch_with_keys( + &self, + batch: &RecordBatch, + row_position: impl Fn(usize) -> Result, + allowed_terms: Option<&HashSet>, + ) -> Result<()> { let st = self.state.load_full(); let document_position_start = st.tail.doc_count(); if self.resolved_field.get().is_none() { @@ -1290,7 +1328,13 @@ impl FtsMemIndex { self.params.get_document_granularity().is_list_element(); let mut index_document = |key: DocumentKey, text: &str| -> Result<()> { let document_position = document_position_start + documents.len() as u64; - let num_tokens = index_text(text, document_position, tokenizer, &mut term_builders)?; + let num_tokens = index_text_filtered( + text, + document_position, + tokenizer, + &mut term_builders, + allowed_terms, + )?; if preserve_zero_token_documents || num_tokens > 0 { documents.push(DocumentMetadata { key, num_tokens }); total_tokens += num_tokens as u64; @@ -1301,7 +1345,7 @@ impl FtsMemIndex { for document in extracted_documents { index_document( DocumentKey { - row_position: row_offset + document.row_index as u64, + row_position: row_position(document.row_index)?, doc_index: document.doc_index, }, &document.text, @@ -1332,6 +1376,127 @@ impl FtsMemIndex { Ok(()) } + /// Analyze every exact leaf and return the deduplicated query terms in + /// canonical leaf traversal order. + pub(crate) fn exact_query_terms(&self, query: &FtsQuery) -> Result> { + fn visit(index: &FtsMemIndex, query: &FtsQuery, terms: &mut Vec) -> Result<()> { + match query { + FtsQuery::Match(query) => { + if query.fuzziness != Some(0) { + return Err(Error::invalid_input( + "residual compound FTS only supports exact Match leaves", + )); + } + terms.extend(index.analyze_for_search(&query.terms)); + } + FtsQuery::Phrase(query) => { + terms.extend(index.analyze_for_search(&query.terms)); + } + FtsQuery::Boost(query) => { + visit(index, &query.positive, terms)?; + visit(index, &query.negative, terms)?; + } + FtsQuery::MultiMatch(query) => { + for query in &query.match_queries { + visit(index, &FtsQuery::Match(query.clone()), terms)?; + } + } + FtsQuery::Boolean(query) => { + for query in query + .should + .iter() + .chain(&query.must) + .chain(&query.must_not) + { + visit(index, query, terms)?; + } + } + } + Ok(()) + } + + let mut terms = Vec::new(); + visit(self, query, &mut terms)?; + let mut seen = HashSet::with_capacity(terms.len()); + terms.retain(|term| seen.insert(term.clone())); + Ok(terms) + } + + /// Build exact residual corpus statistics for the supplied query terms. + pub(crate) fn bm25_stats_for_terms(&self, terms: &[String]) -> MemBM25Scorer { + let st = self.state.load_full(); + let tail = st.tail.snapshot(); + build_scorer(&st, &tail, terms, true) + } + + /// Materialize each exact leaf with a caller-supplied logical-corpus + /// scorer. Compound semantics are deliberately evaluated by the canonical + /// lance-index scorer instead of being duplicated here. + pub(crate) fn exact_leaf_results( + &self, + query: &FtsQuery, + scorer: &MemBM25Scorer, + ) -> Result>> { + fn visit( + index: &FtsMemIndex, + query: &FtsQuery, + scorer: &MemBM25Scorer, + leaves: &mut Vec>, + ) -> Result<()> { + match query { + FtsQuery::Match(query) => { + if query.fuzziness != Some(0) { + return Err(Error::invalid_input( + "residual compound FTS only supports exact Match leaves", + )); + } + let st = index.state.load_full(); + let tokens = index.analyze_for_search(&query.terms); + let rows = index + .search_match_with_scorer(&st, &tokens, query.operator, scorer) + .into_iter() + .map(|entry| (entry.row_position, entry.score)) + .collect(); + leaves.push(rows); + } + FtsQuery::Phrase(query) => { + let st = index.state.load_full(); + let tokens = index.analyze_for_search(&query.terms); + let rows = index + .search_phrase_with_scorer(&st, &tokens, query.slop, scorer) + .into_iter() + .map(|entry| (entry.row_position, entry.score)) + .collect(); + leaves.push(rows); + } + FtsQuery::Boost(query) => { + visit(index, &query.positive, scorer, leaves)?; + visit(index, &query.negative, scorer, leaves)?; + } + FtsQuery::MultiMatch(query) => { + for query in &query.match_queries { + visit(index, &FtsQuery::Match(query.clone()), scorer, leaves)?; + } + } + FtsQuery::Boolean(query) => { + for query in query + .should + .iter() + .chain(&query.must) + .chain(&query.must_not) + { + visit(index, query, scorer, leaves)?; + } + } + } + Ok(()) + } + + let mut leaves = Vec::new(); + visit(self, query, scorer, &mut leaves)?; + Ok(leaves) + } + /// Freeze the current tail into a new immutable partition and publish a /// fresh empty tail. Only the writer calls this; readers snapshotting the /// old `IndexState` keep a consistent view across the freeze. @@ -1569,6 +1734,75 @@ impl FtsMemIndex { } } + fn search_match_with_scorer( + &self, + st: &IndexState, + query_tokens: &Tokens, + operator: Operator, + scorer: &MemBM25Scorer, + ) -> Vec { + if operator == Operator::And && has_grouped_positions(query_tokens) { + let mut result_map: Option> = None; + for group in query_position_groups(query_tokens) { + let group_results = + self.search_match_strings_with_scorer(st, &group, Operator::Or, scorer); + let group_map = group_results + .into_iter() + .map(|entry| (entry.key(), entry.score)) + .collect::>(); + let Some(current) = result_map.as_mut() else { + result_map = Some(group_map); + continue; + }; + current.retain(|key, score| { + if let Some(group_score) = group_map.get(key) { + *score += group_score; + true + } else { + false + } + }); + } + return result_map + .unwrap_or_default() + .into_iter() + .map(|(key, score)| FtsEntry { + row_position: key.row_position, + doc_index: public_doc_index(&key.doc_index), + score, + }) + .collect(); + } + let tokens = query_tokens_to_vec(query_tokens); + self.search_match_strings_with_scorer(st, &tokens, operator, scorer) + } + + fn search_match_strings_with_scorer( + &self, + st: &IndexState, + tokens: &[String], + operator: Operator, + scorer: &MemBM25Scorer, + ) -> Vec { + if tokens.is_empty() || scorer.num_docs() == 0 { + return Vec::new(); + } + let tail = st.tail.snapshot(); + let mut results = Vec::new(); + for partition in st.partitions.iter() { + results.extend(partition.search_match(tokens, operator, scorer)); + } + results.extend(score_terms( + &tail, + &st.tail.terms, + tokens, + operator, + scorer, + f32::NEG_INFINITY, + )); + results + } + fn search_grouped_and( &self, st: &IndexState, @@ -1691,6 +1925,57 @@ impl FtsMemIndex { results } + fn search_phrase_with_scorer( + &self, + st: &IndexState, + query_tokens: &Tokens, + slop: u32, + scorer: &MemBM25Scorer, + ) -> Vec { + if query_tokens.is_empty() || scorer.num_docs() == 0 { + return Vec::new(); + } + let groups = query_position_groups(query_tokens); + if groups.is_empty() { + return Vec::new(); + } + if groups.len() == 1 { + return self.search_match_strings_with_scorer(st, &groups[0], Operator::Or, scorer); + } + if !self.params.has_positions() { + return Vec::new(); + } + let has_grouped_terms = groups.iter().any(|group| group.len() > 1); + let tokens = position_groups_to_tokens(&groups); + let tail = st.tail.snapshot(); + let mut results = Vec::new(); + for partition in st.partitions.iter() { + if has_grouped_terms { + results.extend(partition.search_phrase_groups(&groups, slop, scorer)); + } else { + results.extend(partition.search_phrase(&tokens, slop, scorer)); + } + } + if has_grouped_terms { + results.extend(phrase_search_tail_groups( + &tail, + &st.tail.terms, + &groups, + slop, + scorer, + )); + } else { + results.extend(phrase_search_tail( + &tail, + &st.tail.terms, + &tokens, + slop, + scorer, + )); + } + results + } + fn search_fuzzy_tokens( &self, st: &IndexState, @@ -2258,11 +2543,12 @@ impl BatchTermBuilder { } } -fn index_text( +fn index_text_filtered( text: &str, document_position: u64, tokenizer: &mut dyn LanceTokenizer, term_builders: &mut FxHashMap, BatchTermBuilder>, + allowed_terms: Option<&HashSet>, ) -> Result { let mut stream = tokenizer.token_stream_for_doc(text); let mut num_tokens = 0u32; @@ -2274,13 +2560,15 @@ fn index_text( )) })?; let term = token.text.as_str(); - if let Some(builder) = term_builders.get_mut(term) { - builder.observe(document_position, position); - } else { - term_builders.insert( - Arc::::from(term), - BatchTermBuilder::with_first(document_position, position), - ); + if allowed_terms.is_none_or(|allowed| allowed.contains(term)) { + if let Some(builder) = term_builders.get_mut(term) { + builder.observe(document_position, position); + } else { + term_builders.insert( + Arc::::from(term), + BatchTermBuilder::with_first(document_position, position), + ); + } } num_tokens = num_tokens.checked_add(1).ok_or_else(|| { Error::invalid_input(format!( @@ -4142,6 +4430,38 @@ mod tests { .unwrap() } + #[test] + fn explicit_row_ids_and_query_term_allowlist_preserve_bm25_stats() { + let schema = create_test_schema(); + let batch = create_test_batch(schema.as_ref()); + let row_ids = UInt64Array::from(vec![900, 42, 777]); + let terms = HashSet::from(["hello".to_string()]); + let index = FtsMemIndex::new(1, "description".to_string()); + + index + .insert_with_row_ids_for_terms(&batch, &row_ids, &terms) + .unwrap(); + + assert_eq!(index.doc_count(), 3); + assert_eq!(index.entry_count(), 2); + let scorer = index.bm25_stats_for_terms(&["hello".to_string()]); + assert_eq!(scorer.num_docs, 3); + assert_eq!(scorer.total_tokens, 6); + assert_eq!(scorer.num_docs_containing_token("hello"), 2); + + let query = FtsQuery::Match( + lance_index::scalar::inverted::query::MatchQuery::new("hello".to_string()) + .with_column(Some("description".to_string())), + ); + let leaves = index.exact_leaf_results(&query, &scorer).unwrap(); + let mut actual = leaves[0] + .iter() + .map(|(row_id, _)| *row_id) + .collect::>(); + actual.sort_unstable(); + assert_eq!(actual, vec![777, 900]); + } + fn create_element_test_batch() -> RecordBatch { let mut tags = ListBuilder::new(StringBuilder::new()); tags.values().append_value("alpha beta"); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 8a4fb75a399..95fbeac1561 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -102,8 +102,8 @@ use crate::dataset::utils::SchemaAdapter; use crate::index::DatasetIndexInternalExt; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ - fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, - resolve_fts_field, resolve_query_document_granularity, + fts_index_fragment_bitmap, load_segment_details, load_segment_params, load_segments, + normalize_inverted_details, resolve_fts_field, resolve_query_document_granularity, }; use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_fragment_bitmap}; use crate::index::vector::utils::{ @@ -114,7 +114,8 @@ use crate::io::exec::filtered_read::{ }; use crate::io::exec::fts::{ BoostQueryExec, CompoundQueryExec, CrossColumnCompoundQueryExec, FlatMatchFilterExec, - FlatMatchQueryExec, FtsDocumentExec, MatchQueryExec, PhraseQueryExec, SharedFtsScorer, + FlatMatchQueryExec, FtsDocumentExec, HybridCompoundQueryExec, MatchQueryExec, PhraseQueryExec, + SharedFtsScorer, }; use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; @@ -283,6 +284,65 @@ fn supports_compound_scorer(query: &FtsQuery) -> bool { !columns.is_empty() && (!matches!(query, FtsQuery::MultiMatch(_)) || columns.len() == 1) } +fn supports_exact_residual_compound(query: &FtsQuery) -> bool { + match query { + FtsQuery::Match(query) => query.fuzziness == Some(0), + // MemWAL phrase matching currently collapses tokenizer position gaps. + // Keep phrase queries on the established fallback until it can retain + // those gaps exactly (notably when stop words are configured). + FtsQuery::Phrase(_) => false, + FtsQuery::Boost(query) => { + supports_exact_residual_compound(&query.positive) + && supports_exact_residual_compound(&query.negative) + } + FtsQuery::MultiMatch(query) => query + .match_queries + .iter() + .all(|query| query.fuzziness == Some(0)), + FtsQuery::Boolean(query) => query + .should + .iter() + .chain(&query.must) + .chain(&query.must_not) + .all(supports_exact_residual_compound), + } +} + +fn has_exact_hybrid_fts_coverage( + segments: &[IndexMetadata], + residual_fragments: &[Fragment], + target_fragments: &[Fragment], +) -> bool { + let Some(target) = target_fragments + .iter() + .map(|fragment| u32::try_from(fragment.id).ok()) + .collect::>() + else { + return false; + }; + let Some(residual) = residual_fragments + .iter() + .map(|fragment| u32::try_from(fragment.id).ok()) + .collect::>() + else { + return false; + }; + let mut indexed = RoaringBitmap::new(); + for segment in segments { + let Some(coverage) = segment.fragment_bitmap.as_ref() else { + return false; + }; + if !indexed.is_disjoint(coverage) { + return false; + } + indexed |= coverage; + } + if !indexed.is_subset(&target) || !indexed.is_disjoint(&residual) { + return false; + } + indexed | residual == target +} + fn validate_fts_query_contract(query: &FtsQuery) -> Result<()> { fn validate_multiplier(name: &str, value: f32) -> Result<()> { if value.is_finite() && value >= 0.0 { @@ -4196,6 +4256,7 @@ impl Scanner { &self, query: &FtsQuery, params: &FtsSearchParams, + filter_plan: &ExprFilterPlan, prefilter_source: &PreFilterSource, document_granularity: DocumentGranularity, ) -> Result>> { @@ -4220,6 +4281,17 @@ impl Scanner { } let mut phrase_columns = HashSet::new(); collect_phrase_columns(query, &mut phrase_columns); + let allow_exact_residual = !cross_column + && !self.fast_search + && self.fragments.is_none() + && filter_plan.is_empty() + && self.external_row_mask.is_none() + && params.limit.is_some() + && document_granularity == DocumentGranularity::Row + && target_fragments + .iter() + .all(|fragment| fragment.deletion_file.is_none()) + && supports_exact_residual_compound(query); let segment_groups = futures::future::try_join_all(columns.into_iter().map(|column| { let phrase_columns = &phrase_columns; @@ -4245,6 +4317,8 @@ impl Scanner { let unindexed_fragments = self.retain_target_fragments(unindexed_fragments); if !unindexed_fragments.is_empty() && (!self.fast_search || unindexed_fragments.len() == target_fragments.len()) + && !(allow_exact_residual + && unindexed_fragments.len() < target_fragments.len()) { // Flat and posting-backed leaves do not share a document // domain, so preserve the exact fallback for partial index @@ -4254,10 +4328,6 @@ impl Scanner { // indexed. return Ok(None); } - let unindexed_fragment_ids = unindexed_fragments - .iter() - .map(|fragment| fragment.id as u32) - .collect::(); let segments = match overlay_plan { FtsOverlayPlan::Unchanged(Some(segments)) => segments, FtsOverlayPlan::Unchanged(None) => { @@ -4271,6 +4341,29 @@ impl Scanner { } FtsOverlayPlan::RowLevel { .. } | FtsOverlayPlan::FullScan => return Ok(None), }; + if allow_exact_residual && !unindexed_fragments.is_empty() { + if !has_exact_hybrid_fts_coverage( + &segments, + &unindexed_fragments, + target_fragments, + ) { + return Ok(None); + } + let first_segment = segments.first().ok_or_else(|| { + Error::internal("hybrid compound FTS requires one indexed segment") + })?; + if load_segment_params(&self.dataset, first_segment) + .await? + .posting_block_size() + != 128 + { + // Larger posting blocks quantize document lengths. The + // query-local residual index currently retains exact + // lengths, so the two arms would not have bit-identical + // scores. + return Ok(None); + } + } if cross_column { let details = futures::future::try_join_all( @@ -4306,7 +4399,7 @@ impl Scanner { } } - Ok(Some((column, segments, unindexed_fragment_ids))) + Ok(Some((column, segments, unindexed_fragments))) } })) .await?; @@ -4315,9 +4408,43 @@ impl Scanner { }; if !cross_column { - let (_, segments, _) = segment_groups.into_iter().next().ok_or_else(|| { - Error::internal("compound scorer requires one column".to_string()) - })?; + let (column, segments, unindexed_fragments) = + segment_groups.into_iter().next().ok_or_else(|| { + Error::internal("compound scorer requires one column".to_string()) + })?; + if allow_exact_residual && !unindexed_fragments.is_empty() { + let resolved = + resolve_fts_field(self.dataset.schema(), &column, document_granularity)?; + let scan_column = if resolved.has_lists() { + resolved.root_column.clone() + } else { + resolved.canonical_path.clone() + }; + let scan_projection = self + .dataset + .empty_projection() + .with_row_id() + .union_columns(&[scan_column], OnMissing::Error)?; + let PlannedFilteredScan { plan, .. } = self + .filtered_read( + &ExprFilterPlan::default(), + scan_projection, + /* make_deletions_null */ false, + Some(Arc::new(unindexed_fragments)), + None, + /* is_prefilter */ true, + None, + ) + .await?; + return Ok(Some(Arc::new(HybridCompoundQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + column, + segments, + plan, + )))); + } return Ok(Some(Arc::new( CompoundQueryExec::new_with_segments( self.dataset.clone(), @@ -4334,9 +4461,17 @@ impl Scanner { let Some((_, _, first_unindexed_fragments)) = coverage_groups.next() else { return Ok(None); }; - if coverage_groups - .any(|(_, _, unindexed_fragments)| unindexed_fragments != first_unindexed_fragments) - { + let first_unindexed_fragment_ids = first_unindexed_fragments + .iter() + .map(|fragment| fragment.id as u32) + .collect::(); + if coverage_groups.any(|(_, _, unindexed_fragments)| { + unindexed_fragments + .iter() + .map(|fragment| fragment.id as u32) + .collect::() + != first_unindexed_fragment_ids + }) { // The cross-column scorer builds one shared prefilter. If column // coverage differs, that prefilter's union can re-admit stale // postings from a fragment invalidated only for another column. @@ -4348,7 +4483,6 @@ impl Scanner { .into_iter() .map(|(column, segments, _)| (column, segments)) .collect(); - let exec = CrossColumnCompoundQueryExec::new_with_segments( self.dataset.clone(), query.clone(), @@ -4371,7 +4505,13 @@ impl Scanner { if !document_granularity.is_list_element() && supports_compound_scorer(query) && let Some(plan) = self - .plan_compound_scorer(query, params, prefilter_source, document_granularity) + .plan_compound_scorer( + query, + params, + filter_plan, + prefilter_source, + document_granularity, + ) .await? { return Ok(plan); @@ -4441,6 +4581,7 @@ impl Scanner { .plan_compound_scorer( &child_query, params, + filter_plan, field_prefilter_source, document_granularity, ) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index bdcd6e5719c..65b2da20e04 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2619,12 +2619,70 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { .full_text_search(FullTextSearchQuery::new_query(query.clone())) .unwrap(); exact_scanner.limit(Some(2), None).unwrap(); + let exact_plan = exact_scanner.explain_plan(false).await.unwrap(); + assert!( + exact_plan.contains("HybridCompoundFtsScorer"), + "exact partial coverage should build one query-local residual index:\n{exact_plan}" + ); + assert!( + !exact_plan.contains("FlatMatchQuery"), + "hybrid compound scoring must not scan the residual once per leaf:\n{exact_plan}" + ); let exact = exact_scanner.try_into_batch().await.unwrap(); assert_eq!( exact["id"].as_primitive::().values(), &[0, 2], "exact search should include the appended hit" ); + let (_, exact_stats) = compound_fts_results_with_stats(&dataset, query.clone(), 2).await; + assert_eq!( + exact_stats + .all_counts + .get(crate::io::exec::fts::HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC), + Some(&1) + ); + assert_eq!( + exact_stats + .all_counts + .get(crate::io::exec::fts::HYBRID_COMPOUND_RESIDUAL_DOCS_INDEXED_METRIC), + Some(&1) + ); + assert_eq!( + exact_stats + .all_counts + .get(crate::io::exec::fts::HYBRID_COMPOUND_MERGED_CANDIDATES_METRIC), + Some(&2) + ); + + let mut filtered_scanner = dataset.scan(); + filtered_scanner + .with_row_id() + .filter("id >= 0") + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query.clone())) + .unwrap(); + filtered_scanner.limit(Some(2), None).unwrap(); + let filtered_plan = filtered_scanner.explain_plan(false).await.unwrap(); + assert!( + !filtered_plan.contains("HybridCompoundFtsScorer"), + "filtered residual scoring must retain the exact fallback:\n{filtered_plan}" + ); + + let phrase_query: FtsQuery = BooleanQuery::new([ + ( + Occur::Must, + PhraseQuery::new("fresh alpha".to_string()) + .with_column(Some("text".to_string())) + .into(), + ), + (Occur::Must, compound_match_query("fresh", "text", 1.0)), + ]) + .into(); + let phrase_plan = compound_fts_plan(&dataset, phrase_query, 2).await; + assert!( + !phrase_plan.contains("HybridCompoundFtsScorer"), + "phrase position gaps are not yet supported by the residual index:\n{phrase_plan}" + ); let mut fast_scanner = dataset.scan(); fast_scanner @@ -2673,6 +2731,92 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { ); } +#[tokio::test] +async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { + let initial = arrow_array::record_batch!( + ("text", Utf8, ["fresh alpha", "blocked fresh alpha"]), + ("id", Int32, [0, 1]) + ) + .unwrap(); + let schema = initial.schema(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema), + "memory://", + None, + ) + .await + .unwrap(); + create_fragmented_fts_index(&mut dataset, "text", true).await; + + let appended = arrow_array::record_batch!( + ("text", Utf8, ["fresh alpha", "fresh beta"]), + ("id", Int32, [2, 3]) + ) + .unwrap(); + let schema = appended.schema(); + dataset + .append( + RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), + None, + ) + .await + .unwrap(); + + let positive: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("fresh", "text", 1.0)), + (Occur::Should, compound_match_query("alpha", "text", 1.0)), + (Occur::MustNot, compound_match_query("blocked", "text", 1.0)), + ]) + .into(); + let boost_query: FtsQuery = BoostQuery::new( + positive, + compound_match_query("alpha", "text", 1.0), + Some(0.25), + ) + .into(); + let partial_boost = compound_fts_results(&dataset, boost_query.clone(), Some(10)).await; + assert_eq!( + partial_boost.len(), + 3, + "MUST_NOT must exclude the blocked row" + ); + assert_eq!(partial_boost[0].1.to_bits(), partial_boost[1].1.to_bits()); + assert!( + partial_boost[0].0 < partial_boost[1].0, + "equal-score rows must use ascending row id as the exact tie break" + ); + + let multimatch_query: FtsQuery = MultiMatchQuery::try_new( + "fresh alpha".to_string(), + vec!["text".to_string(), "text".to_string()], + ) + .unwrap() + .try_with_boosts(vec![1.0, 2.0]) + .unwrap() + .into(); + let partial_multimatch = + compound_fts_results(&dataset, multimatch_query.clone(), Some(3)).await; + + dataset + .create_index( + &["text"], + IndexType::Inverted, + Some("text_idx".to_string()), + &InvertedIndexParams::default().with_position(true), + true, + ) + .await + .unwrap(); + let rebuilt_boost = compound_fts_results(&dataset, boost_query, Some(10)).await; + let rebuilt_multimatch = compound_fts_results(&dataset, multimatch_query, Some(3)).await; + assert_scored_rows_close("partial_hybrid_boost", &partial_boost, &rebuilt_boost); + assert_scored_rows_close( + "partial_hybrid_multimatch", + &partial_multimatch, + &rebuilt_multimatch, + ); +} + #[tokio::test] async fn test_boolean_must_scores_sum_across_execution_paths() { let batch = arrow_array::record_batch!( diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index a9bd8938184..3432a70e49e 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -38,6 +38,7 @@ use lance_table::format::IndexMetadata; use super::PreFilterSource; use super::utils::{IndexMetrics, PreFilterMasks, build_prefilter}; +use crate::dataset::mem_wal::index::FtsMemIndex; use crate::index::scalar::inverted::{ ResolvedFtsField, fts_document_schema, load_segment_details, load_segments, transform_fts_document_stream, @@ -56,8 +57,10 @@ use lance_index::scalar::inverted::{ DOC_INDEX_COL, DocumentGranularity, FTS_SCHEMA, FlatBm25SearchOptions, InvertedIndex, MemBM25Scorer, PreparedBm25Query, SCORE_COL, Scorer, build_global_bm25_scorer, compound_search, compound_search_prepared_match, compound_search_prepared_match_with_score_floor, - compound_search_with_base_scorer, cross_column_compound_search, exclusive_scaled_score_floor, - flat_bm25_search_stream_with_options_and_scorer, fts_schema, prepare_bm25_query, + compound_search_with_base_scorer, compound_search_with_base_scorer_and_score_floor, + cross_column_compound_search, exclusive_scaled_score_floor, + flat_bm25_search_stream_with_options_and_scorer, fts_schema, materialized_compound_top_k, + prepare_bm25_query, }; use lance_index::{prefilter::PreFilter, scalar::inverted::query::BooleanQuery}; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; @@ -790,6 +793,262 @@ impl CompoundQueryExec { } } +/// Exact compound FTS over committed postings plus an append-only residual +/// scan. The residual documents are tokenized once into query-local postings, +/// rather than once for every compound leaf. +#[derive(Debug)] +pub(crate) struct HybridCompoundQueryExec { + dataset: Arc, + query: FtsQuery, + params: FtsSearchParams, + column: String, + segments: Arc<[IndexMetadata]>, + residual_input: Arc, + properties: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl HybridCompoundQueryExec { + pub(crate) fn new( + dataset: Arc, + query: FtsQuery, + params: FtsSearchParams, + column: String, + segments: Vec, + residual_input: Arc, + ) -> Self { + Self { + dataset, + query, + params, + column, + segments: Arc::from(segments), + residual_input, + properties: Arc::new(PlanProperties::new( + EquivalenceProperties::new(FTS_SCHEMA.clone()), + Partitioning::RoundRobinBatch(1), + EmissionType::Final, + Boundedness::Bounded, + )), + metrics: ExecutionPlanMetricsSet::new(), + } + } +} + +impl DisplayAs for HybridCompoundQueryExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "HybridCompoundFtsScorer: column={}, query={}", + self.column, self.query + ) + } +} + +impl ExecutionPlan for HybridCompoundQueryExec { + fn name(&self) -> &str { + "HybridCompoundQueryExec" + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.residual_input] + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::SinglePartition] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return Err(DataFusionError::Internal(format!( + "hybrid compound FTS expected one residual child, got {}", + children.len() + ))); + } + let residual_input = children.pop().ok_or_else(|| { + DataFusionError::Internal("hybrid compound FTS lost its residual child".to_string()) + })?; + Ok(Arc::new(Self::new( + self.dataset.clone(), + self.query.clone(), + self.params.clone(), + self.column.clone(), + self.segments.to_vec(), + residual_input, + ))) + } + + #[instrument(name = "hybrid_compound_fts_exec", level = "debug", skip_all)] + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let dataset = self.dataset.clone(); + let query = self.query.clone(); + let params = self.params.clone(); + let column = self.column.clone(); + let segments = self.segments.clone(); + let mut residual_input = self.residual_input.execute(partition, context.clone())?; + let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); + let residual_rows_scanned = self + .metrics + .new_count(HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC, partition); + let residual_docs_indexed = self + .metrics + .new_count(HYBRID_COMPOUND_RESIDUAL_DOCS_INDEXED_METRIC, partition); + let index_candidates = self + .metrics + .new_count(HYBRID_COMPOUND_INDEX_CANDIDATES_METRIC, partition); + let residual_candidates = self + .metrics + .new_count(HYBRID_COMPOUND_RESIDUAL_CANDIDATES_METRIC, partition); + let merged_candidates = self + .metrics + .new_count(HYBRID_COMPOUND_MERGED_CANDIDATES_METRIC, partition); + let schema = self.schema(); + + let stream = stream::once(async move { + let _timer = metrics.baseline_metrics.elapsed_compute().timer(); + let indices = + open_fts_segments(&dataset, &column, &segments, &metrics.index_metrics).await?; + let first_index = indices.first().ok_or_else(|| { + DataFusionError::Execution(format!( + "FTS index for column {column} has no committed segments" + )) + })?; + let field_id = dataset.schema().field_id(&column)?; + let residual = FtsMemIndex::try_with_params( + field_id, + column.clone(), + first_index.params().clone(), + )?; + let terms = residual.exact_query_terms(&query)?; + let allowed_terms = terms.iter().cloned().collect::>(); + + while let Some(batch) = residual_input.try_next().await? { + residual_rows_scanned.add(batch.num_rows()); + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| { + DataFusionError::Execution( + "hybrid compound FTS residual input is missing _rowid".to_string(), + ) + })? + .as_primitive::(); + residual.insert_with_row_ids_for_terms(&batch, row_ids, &allowed_terms)?; + } + residual_docs_indexed.add(residual.doc_count()); + + let query_tokens = Tokens::new(terms.clone(), first_index.tokenizer().doc_type()); + let exact_params = params + .clone() + .with_fuzziness(Some(0)) + .with_phrase_slop(None); + let mut scorer = build_global_bm25_scorer( + &indices, + &query_tokens, + &exact_params, + Some(metrics.as_ref()), + ) + .await?; + let residual_stats = residual.bm25_stats_for_terms(&terms); + scorer.total_tokens = scorer + .total_tokens + .checked_add(residual_stats.total_tokens) + .ok_or_else(|| { + DataFusionError::Execution( + "hybrid compound FTS total token count overflow".to_string(), + ) + })?; + scorer.num_docs = scorer + .num_docs + .checked_add(residual_stats.num_docs) + .ok_or_else(|| { + DataFusionError::Execution( + "hybrid compound FTS document count overflow".to_string(), + ) + })?; + for term in &terms { + let residual_df = residual_stats.num_docs_containing_token(term); + let df = scorer.token_docs.get_mut(term).ok_or_else(|| { + DataFusionError::Execution(format!( + "hybrid compound FTS scorer is missing query term '{term}'" + )) + })?; + *df = df.checked_add(residual_df).ok_or_else(|| { + DataFusionError::Execution(format!( + "hybrid compound FTS document frequency overflow for term '{term}'" + )) + })?; + } + let scorer = Arc::new(scorer); + let limit = params.limit.ok_or_else(|| { + DataFusionError::Execution( + "hybrid compound FTS requires a bounded result limit".to_string(), + ) + })?; + + let prefilter = build_prefilter( + context, + partition, + &PreFilterSource::None, + dataset, + &segments, + None, + None, + )?; + let (indexed_row_ids, indexed_scores) = compound_search_with_base_scorer( + &indices, + &query, + ¶ms, + prefilter, + metrics.clone(), + scorer.clone(), + ) + .await?; + index_candidates.add(indexed_row_ids.len()); + let residual_leaves = residual.exact_leaf_results(&query, scorer.as_ref())?; + let (residual_row_ids, residual_scores) = + materialized_compound_top_k(&query, residual_leaves, limit, metrics.as_ref())?; + residual_candidates.add(residual_row_ids.len()); + + let mut documents = indexed_row_ids + .into_iter() + .zip(indexed_scores) + .chain(residual_row_ids.into_iter().zip(residual_scores)) + .map(|(row_id, score)| ScoredDoc::new(row_id, score)) + .collect::>(); + merged_candidates.add(documents.len()); + documents.sort_unstable_by(|left, right| { + right + .score + .0 + .total_cmp(&left.score.0) + .then_with(|| left.row_id.cmp(&right.row_id)) + }); + documents.truncate(limit); + metrics.baseline_metrics.record_output(documents.len()); + scored_documents_batch(schema, documents).map_err(DataFusionError::from) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream.stream_in_current_span().boxed(), + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn properties(&self) -> &Arc { + &self.properties + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum WandExactnessCertificate { Exhaustive, @@ -1874,6 +2133,15 @@ impl Drop for SharedFtsScorerProducer { /// Time spent resolving an exact ordered UUID selection to committed FTS segments. pub const FTS_SEGMENT_BIND_DURATION_METRIC: &str = "fts_segment_bind_duration"; +pub(crate) const HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC: &str = + "hybrid_compound_residual_rows_scanned"; +pub(crate) const HYBRID_COMPOUND_RESIDUAL_DOCS_INDEXED_METRIC: &str = + "hybrid_compound_residual_docs_indexed"; +pub(crate) const HYBRID_COMPOUND_INDEX_CANDIDATES_METRIC: &str = "hybrid_compound_index_candidates"; +pub(crate) const HYBRID_COMPOUND_RESIDUAL_CANDIDATES_METRIC: &str = + "hybrid_compound_residual_candidates"; +pub(crate) const HYBRID_COMPOUND_MERGED_CANDIDATES_METRIC: &str = + "hybrid_compound_merged_candidates"; #[derive(Debug, Clone)] enum FtsSegmentSelection { From 09f00f9629a5e80d29f102a2ebfe53653ec66c3b Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 13:15:16 +0800 Subject: [PATCH 02/22] fix(fts): reject unsafe hybrid index segments --- rust/lance/src/dataset/scanner.rs | 97 +++++++-- rust/lance/src/dataset/tests/dataset_index.rs | 195 +++++++++++++++++- 2 files changed, 272 insertions(+), 20 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 95fbeac1561..5031dd0755b 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -77,7 +77,8 @@ use lance_index::scalar::inverted::query::{ }; use lance_index::scalar::inverted::{ DOC_INDEX_COL, DOC_INDEX_FIELD, DocumentGranularity, INVERTED_INDEX_VERSION_V2, - INVERTED_INDEX_VERSION_V3, SCORE_COL, SCORE_FIELD, fts_schema, + INVERTED_INDEX_VERSION_V3, InvertedIndex, InvertedIndexParams, SCORE_COL, SCORE_FIELD, + fts_schema, }; use lance_index::scalar::registry::VALUE_COLUMN_NAME; use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, DIST_COL, Query}; @@ -343,6 +344,19 @@ fn has_exact_hybrid_fts_coverage( indexed | residual == target } +fn has_compatible_hybrid_physical_segments( + params: &[InvertedIndexParams], + has_deleted_fragments: &[bool], +) -> bool { + let Some(first) = params.first() else { + return false; + }; + params.len() == has_deleted_fragments.len() + && first.posting_block_size() == 128 + && params.iter().all(|params| params == first) + && has_deleted_fragments.iter().all(|has_deleted| !has_deleted) +} + fn validate_fts_query_contract(query: &FtsQuery) -> Result<()> { fn validate_multiplier(name: &str, value: f32) -> Result<()> { if value.is_finite() && value >= 0.0 { @@ -4349,18 +4363,53 @@ impl Scanner { ) { return Ok(None); } - let first_segment = segments.first().ok_or_else(|| { - Error::internal("hybrid compound FTS requires one indexed segment") - })?; - if load_segment_params(&self.dataset, first_segment) - .await? - .posting_block_size() - != 128 - { - // Larger posting blocks quantize document lengths. The - // query-local residual index currently retains exact - // lengths, so the two arms would not have bit-identical - // scores. + if segments.is_empty() { + return Err(Error::internal( + "hybrid compound FTS requires one indexed segment", + )); + } + // Preserve the established semantic mismatch error before + // applying the narrower physical fast-path gate. + load_segment_details(&self.dataset, &column, &segments).await?; + let segment_params = futures::future::try_join_all( + segments + .iter() + .map(|segment| load_segment_params(&self.dataset, segment)), + ) + .await?; + let has_deleted_fragments = futures::future::try_join_all( + segments.iter().map(|segment| { + let column = &column; + async move { + let index = self + .dataset + .open_scalar_index( + column, + &segment.uuid, + &NoOpMetricsCollector, + ) + .await?; + let index = index + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::internal(format!( + "hybrid compound FTS segment {} is not an inverted index", + segment.uuid + )) + })?; + Ok::<_, Error>(!index.deleted_fragments().is_empty()) + } + }), + ) + .await?; + if !has_compatible_hybrid_physical_segments( + &segment_params, + &has_deleted_fragments, + ) { + // Larger posting blocks quantize document lengths, and + // retired physical documents remain in corpus stats. + // Either would make the two arms incomparable. return Ok(None); } } @@ -7534,6 +7583,28 @@ mod test { assert!(!supports_compound_scorer(&cross_column)); } + #[test] + fn test_hybrid_compound_requires_compatible_live_physical_segments() { + let params = InvertedIndexParams::default(); + assert!(has_compatible_hybrid_physical_segments( + &[params.clone(), params.clone()], + &[false, false] + )); + assert!(!has_compatible_hybrid_physical_segments( + &[params.clone(), params.clone()], + &[false, true] + )); + assert!(!has_compatible_hybrid_physical_segments( + &[params.clone(), params.clone().with_position(true)], + &[false, false] + )); + assert!(!has_compatible_hybrid_physical_segments( + &[params.clone().block_size(256).unwrap()], + &[false] + )); + assert!(!has_compatible_hybrid_physical_segments(&[params], &[])); + } + #[test] fn test_collect_phrase_columns_traverses_prohibited_subtrees() { let phrase = diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 65b2da20e04..fc56f26496e 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -10,6 +10,8 @@ use std::vec; use crate::dataset::ROW_ID; use crate::dataset::builder::DatasetBuilder; +use crate::dataset::index::LanceIndexStoreExt; +use crate::dataset::optimize::{CompactionOptions, compact_files}; use crate::dataset::tests::dataset_migrations::scan_dataset; use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; use crate::dataset::transaction::{Operation, Transaction}; @@ -20,7 +22,7 @@ use crate::{Dataset, Error, Result}; use lance_arrow::FixedSizeListArrayExt; use crate::dataset::write::{WriteMode, WriteParams}; -use crate::index::DatasetIndexExt; +use crate::index::{DatasetIndexExt, DatasetIndexInternalExt}; use arrow::array::{AsArray, GenericListBuilder, GenericStringBuilder}; use arrow::datatypes::UInt64Type; use arrow_array::RecordBatch; @@ -49,7 +51,10 @@ use lance_index::scalar::inverted::{ query::{BooleanQuery, BoostQuery, MatchQuery, Occur, Operator, PhraseQuery}, tokenizer::InvertedIndexParams, }; -use lance_index::scalar::{FullTextSearchQuery, ScalarIndex}; +use lance_index::scalar::lance_format::LanceIndexStore; +use lance_index::scalar::{ + FullTextSearchQuery, OldIndexDataFilter, ScalarIndex, index_files_to_table, +}; use lance_index::{FtsPrewarmOptions, PrewarmOptions}; use lance_index::{IndexType, scalar::ScalarIndexParams, vector::DIST_COL}; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; @@ -63,7 +68,9 @@ use lance_arrow::json::ARROW_JSON_EXT_NAME; use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; use lance_testing::datagen::generate_random_array; use rand::Rng; +use roaring::RoaringBitmap; use rstest::rstest; +use uuid::Uuid; #[rstest] #[tokio::test] @@ -2809,11 +2816,185 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { .unwrap(); let rebuilt_boost = compound_fts_results(&dataset, boost_query, Some(10)).await; let rebuilt_multimatch = compound_fts_results(&dataset, multimatch_query, Some(3)).await; - assert_scored_rows_close("partial_hybrid_boost", &partial_boost, &rebuilt_boost); - assert_scored_rows_close( - "partial_hybrid_multimatch", - &partial_multimatch, - &rebuilt_multimatch, + let score_bits = |rows: &[(u64, f32)]| { + rows.iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect::>() + }; + assert_eq!( + score_bits(&partial_boost), + score_bits(&rebuilt_boost), + "hybrid Boost scores must be bit-identical to a rebuilt index" + ); + assert_eq!( + score_bits(&partial_multimatch), + score_bits(&rebuilt_multimatch), + "hybrid MultiMatch scores must be bit-identical to a rebuilt index" + ); +} + +#[tokio::test] +async fn test_partial_compound_hybrid_rejects_retired_physical_fragments() { + let initial = arrow_array::record_batch!( + ("text", Utf8, ["retired alpha", "live alpha"]), + ("id", Int32, [0, 1]) + ) + .unwrap(); + let schema = initial.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema), + &test_uri, + None, + ) + .await + .unwrap(); + create_fragmented_fts_index(&mut dataset, "text", true).await; + + let initial_segment = dataset + .load_index_by_name("text_idx") + .await + .unwrap() + .unwrap(); + let retired_fragments = initial_segment.fragment_bitmap.clone().unwrap(); + let initial_index = dataset + .open_scalar_index( + "text", + &initial_segment.uuid, + &lance_index::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + + dataset.delete("id = 0").await.unwrap(); + let metrics = compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 10, + materialize_deletions_threshold: 0.0, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert!(metrics.fragments_removed > 0); + assert!( + dataset + .get_fragments() + .iter() + .all(|fragment| fragment.metadata().deletion_file.is_none()) + ); + let indexed_fragments = dataset.fragment_bitmap.as_ref().clone(); + assert!(indexed_fragments.is_disjoint(&retired_fragments)); + + // Model an incremental FTS replacement that retains the old postings and + // records their now-retired fragment ids for merge-on-read filtering. + let resolved = crate::index::scalar::inverted::resolve_fts_field_by_id( + dataset.schema(), + initial_segment.fields[0], + DocumentGranularity::Row, + ) + .unwrap(); + let current_fragments = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.metadata().clone()) + .collect(); + let update_criteria = initial_index.update_criteria(); + let new_data = crate::index::scalar::load_fts_training_data( + &dataset, + &resolved, + &update_criteria.data_criteria, + Some(current_fragments), + true, + None, + ) + .await + .unwrap(); + let updated_uuid = Uuid::new_v4(); + let updated_store = LanceIndexStore::from_dataset_for_new(&dataset, &updated_uuid).unwrap(); + let created = initial_index + .update( + new_data, + &updated_store, + Some(OldIndexDataFilter::Fragments { + to_keep: RoaringBitmap::new(), + to_remove: retired_fragments.clone(), + }), + ) + .await + .unwrap(); + let updated_segment = lance_table::format::IndexMetadata { + uuid: updated_uuid, + dataset_version: dataset.manifest.version, + fragment_bitmap: Some(indexed_fragments), + index_details: Some(Arc::new(created.index_details)), + index_version: created.index_version as i32, + created_at: Some(chrono::Utc::now()), + base_id: None, + files: Some(index_files_to_table(created.files)), + ..initial_segment + }; + dataset + .commit_existing_index_segments("text_idx", "text", vec![updated_segment]) + .await + .unwrap(); + + let committed = dataset + .load_index_by_name("text_idx") + .await + .unwrap() + .unwrap(); + let committed_index = dataset + .open_scalar_index( + "text", + &committed.uuid, + &lance_index::metrics::NoOpMetricsCollector, + ) + .await + .unwrap(); + let committed_index = committed_index + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(committed_index.deleted_fragments(), &retired_fragments); + + let appended = + arrow_array::record_batch!(("text", Utf8, ["tail alpha"]), ("id", Int32, [2])).unwrap(); + let schema = appended.schema(); + dataset + .append( + RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), + None, + ) + .await + .unwrap(); + + let query: FtsQuery = BooleanQuery::new([ + (Occur::Should, compound_match_query("retired", "text", 1.0)), + (Occur::Should, compound_match_query("tail", "text", 1.0)), + ]) + .into(); + let plan = compound_fts_plan(&dataset, query.clone(), 10).await; + assert!( + !plan.contains("HybridCompoundFtsScorer"), + "retired physical docs make unified hybrid BM25 stats unsafe:\n{plan}" + ); + + let mut scanner = dataset.scan(); + scanner + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query)) + .unwrap() + .limit(Some(10), None) + .unwrap(); + let results = scanner.try_into_batch().await.unwrap(); + assert_eq!( + results["id"].as_primitive::().values(), + &[2], + "fallback must prune stale postings from the retired fragment" ); } From 587bfe17ef41fc69369451a3780a0f1f7bfd32d0 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 13:37:00 +0800 Subject: [PATCH 03/22] fix(fts): bound hybrid residual materialization --- rust/lance/src/dataset/mem_wal/index.rs | 1 + rust/lance/src/dataset/mem_wal/index/fts.rs | 100 +++++++++++++++++- rust/lance/src/dataset/scanner.rs | 11 ++ rust/lance/src/dataset/tests/dataset_index.rs | 23 ++++ rust/lance/src/io/exec/fts.rs | 73 +++++++++---- 5 files changed, 189 insertions(+), 19 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index 3daf3e1274a..7653891bcf2 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -49,6 +49,7 @@ pub type RowPosition = u64; // Re-export public types used externally pub use btree::{BTreeIndexConfig, BTreeMemIndex}; +pub(crate) use fts::QueryLocalFtsIndex; pub use fts::{FtsIndexConfig, FtsMemIndex, FtsQueryExpr, SearchOptions}; pub use hnsw::{HnswIndexConfig, HnswMemIndex}; pub use pk_key::encode_pk_tuple; diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index db09f68e821..f822ebd8b1d 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1003,6 +1003,11 @@ pub struct FtsMemIndex { /// The tail freezes into a partition once it reaches this many docs. freeze_threshold_rows: usize, + /// Query-local materializations disable freezes and tiered merges. Their + /// lifetime is bounded by one query, so background maintenance would only + /// outlive cancellation without providing reuse. + background_maintenance: bool, + /// Background tiered-merge slot. `None` = idle; `Some` with `result: None` /// = a merge is running on a worker thread; `Some` with `result: Some` = /// the merged partition is ready for the writer to install. Only the @@ -1011,6 +1016,64 @@ pub struct FtsMemIndex { merge: Arc>>, } +/// Query-owned exact postings for one residual scan. +/// +/// This deliberately exposes only the immutable feature-materialization API +/// needed by hybrid execution. Unlike [`FtsMemIndex`], it never freezes or +/// starts a detached tiered merge; dropping the query drops all residual +/// postings. +#[derive(Debug)] +pub(crate) struct QueryLocalFtsIndex { + inner: FtsMemIndex, +} + +impl QueryLocalFtsIndex { + pub(crate) fn try_with_params( + field_id: i32, + column_name: String, + params: InvertedIndexParams, + ) -> Result { + Ok(Self { + inner: FtsMemIndex::try_with_params_and_maintenance( + field_id, + column_name, + params, + false, + )?, + }) + } + + pub(crate) fn exact_query_terms(&self, query: &FtsQuery) -> Result> { + self.inner.exact_query_terms(query) + } + + pub(crate) fn insert_with_row_ids_for_terms( + &self, + batch: &RecordBatch, + row_ids: &UInt64Array, + terms: &HashSet, + ) -> Result<()> { + self.inner + .insert_with_row_ids_for_terms(batch, row_ids, terms) + } + + pub(crate) fn doc_count(&self) -> usize { + self.inner.doc_count() + } + + pub(crate) fn bm25_stats_for_terms(&self, terms: &[String]) -> MemBM25Scorer { + self.inner.bm25_stats_for_terms(terms) + } + + pub(crate) fn exact_leaf_results( + &self, + query: &FtsQuery, + scorer: &MemBM25Scorer, + ) -> Result>> { + self.inner.exact_leaf_results(query, scorer) + } +} + /// A tiered merge dispatched to a background worker. struct PendingMerge { /// `Arc::as_ptr` of each source partition, for identity-matching the @@ -1090,6 +1153,15 @@ impl FtsMemIndex { field_id: i32, column_name: String, params: InvertedIndexParams, + ) -> Result { + Self::try_with_params_and_maintenance(field_id, column_name, params, true) + } + + fn try_with_params_and_maintenance( + field_id: i32, + column_name: String, + params: InvertedIndexParams, + background_maintenance: bool, ) -> Result { params.validate_format_version()?; let pool = TokenizerPool::new(¶ms, Self::DEFAULT_TOKENIZER_POOL_CAP)?; @@ -1103,6 +1175,7 @@ impl FtsMemIndex { writer_tokenizer: Mutex::new(writer_tokenizer), state: ArcSwap::from(IndexState::empty()), freeze_threshold_rows: Self::DEFAULT_FREEZE_THRESHOLD_ROWS, + background_maintenance, merge: Arc::new(Mutex::new(None)), }) } @@ -1370,7 +1443,7 @@ impl FtsMemIndex { self.params.has_positions(), ); - if st.tail.doc_count() >= self.freeze_threshold_rows as u64 { + if self.background_maintenance && st.tail.doc_count() >= self.freeze_threshold_rows as u64 { self.freeze(&st)?; } Ok(()) @@ -4462,6 +4535,31 @@ mod tests { assert_eq!(actual, vec![777, 900]); } + #[test] + fn query_local_materialization_never_starts_background_maintenance() { + let schema = create_test_schema(); + let batch = create_test_batch(schema.as_ref()); + let row_ids = UInt64Array::from(vec![900, 42, 777]); + let terms = HashSet::from(["hello".to_string()]); + let mut index = QueryLocalFtsIndex::try_with_params( + 1, + "description".to_string(), + InvertedIndexParams::default(), + ) + .unwrap(); + // Crossing the normal freeze threshold would create a partition and + // may launch a detached tiered merge. Query-local materialization must + // remain entirely in its query-owned tail instead. + index.inner.freeze_threshold_rows = 1; + index + .insert_with_row_ids_for_terms(&batch, &row_ids, &terms) + .unwrap(); + + assert!(index.inner.state.load().partitions.is_empty()); + assert!(index.inner.merge.lock().unwrap().is_none()); + assert_eq!(index.doc_count(), 3); + } + fn create_element_test_batch() -> RecordBatch { let mut tags = ListBuilder::new(StringBuilder::new()); tags.values().append_value("alpha beta"); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 5031dd0755b..c530c2beac7 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -354,6 +354,12 @@ fn has_compatible_hybrid_physical_segments( params.len() == has_deleted_fragments.len() && first.posting_block_size() == 128 && params.iter().all(|params| params == first) + && params.iter().all(|params| { + matches!( + params.resolved_format_version().index_version(), + INVERTED_INDEX_VERSION_V2 | INVERTED_INDEX_VERSION_V3 + ) + }) && has_deleted_fragments.iter().all(|has_deleted| !has_deleted) } @@ -7368,6 +7374,7 @@ mod test { }; use lance_file::version::LanceFileVersion; use lance_index::optimize::OptimizeOptions; + use lance_index::scalar::inverted::InvertedListFormatVersion; use lance_index::scalar::inverted::query::{ BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, PhraseQuery, }; @@ -7602,6 +7609,10 @@ mod test { &[params.clone().block_size(256).unwrap()], &[false] )); + assert!(!has_compatible_hybrid_physical_segments( + &[params.clone().format_version(InvertedListFormatVersion::V1)], + &[false] + )); assert!(!has_compatible_hybrid_physical_segments(&[params], &[])); } diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index fc56f26496e..446fa3ce754 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2661,6 +2661,29 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { Some(&2) ); + let empty_terms_query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("", "text", 1.0)), + (Occur::Should, compound_match_query(" ", "text", 1.0)), + ]) + .into(); + let empty_terms_plan = compound_fts_plan(&dataset, empty_terms_query.clone(), 2).await; + assert!( + empty_terms_plan.contains("HybridCompoundFtsScorer"), + "the empty analyzed-term case must exercise the hybrid short circuit:\n{empty_terms_plan}" + ); + let (empty_results, empty_stats) = + compound_fts_results_with_stats(&dataset, empty_terms_query, 2).await; + assert!(empty_results.is_empty()); + assert_eq!( + empty_stats + .all_counts + .get(crate::io::exec::fts::HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC) + .copied() + .unwrap_or_default(), + 0, + "an empty analyzed query must not poll the residual scan" + ); + let mut filtered_scanner = dataset.scan(); filtered_scanner .with_row_id() diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 3432a70e49e..4a7370c9d22 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -30,7 +30,10 @@ use futures::{FutureExt, StreamExt, TryStreamExt}; use itertools::Itertools; use lance_core::{ Error, ROW_ID, Result, - utils::{tokio::get_num_compute_intensive_cpus, tracing::StreamTracingExt}, + utils::{ + tokio::{get_num_compute_intensive_cpus, spawn_cpu}, + tracing::StreamTracingExt, + }, }; use lance_datafusion::utils::{ExecutionPlanMetricsSetExt, MetricsExt, PARTITIONS_SEARCHED_METRIC}; use lance_select::RowAddrMask; @@ -38,7 +41,7 @@ use lance_table::format::IndexMetadata; use super::PreFilterSource; use super::utils::{IndexMetrics, PreFilterMasks, build_prefilter}; -use crate::dataset::mem_wal::index::FtsMemIndex; +use crate::dataset::mem_wal::index::QueryLocalFtsIndex; use crate::index::scalar::inverted::{ ResolvedFtsField, fts_document_schema, load_segment_details, load_segments, transform_fts_document_stream, @@ -793,6 +796,26 @@ impl CompoundQueryExec { } } +async fn index_query_local_residual_batch( + residual: QueryLocalFtsIndex, + batch: RecordBatch, + allowed_terms: Arc>, +) -> Result { + spawn_cpu(move || { + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| { + Error::invalid_input( + "hybrid compound FTS residual input is missing _rowid".to_string(), + ) + })? + .as_primitive::(); + residual.insert_with_row_ids_for_terms(&batch, row_ids, allowed_terms.as_ref())?; + Ok(residual) + }) + .await +} + /// Exact compound FTS over committed postings plus an append-only residual /// scan. The residual documents are tokenized once into query-local postings, /// rather than once for every compound leaf. @@ -892,7 +915,7 @@ impl ExecutionPlan for HybridCompoundQueryExec { let params = self.params.clone(); let column = self.column.clone(); let segments = self.segments.clone(); - let mut residual_input = self.residual_input.execute(partition, context.clone())?; + let residual_input = self.residual_input.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); let residual_rows_scanned = self .metrics @@ -921,25 +944,23 @@ impl ExecutionPlan for HybridCompoundQueryExec { )) })?; let field_id = dataset.schema().field_id(&column)?; - let residual = FtsMemIndex::try_with_params( + let mut residual = QueryLocalFtsIndex::try_with_params( field_id, column.clone(), first_index.params().clone(), )?; let terms = residual.exact_query_terms(&query)?; - let allowed_terms = terms.iter().cloned().collect::>(); + if terms.is_empty() { + metrics.baseline_metrics.record_output(0); + return scored_documents_batch(schema, Vec::new()).map_err(DataFusionError::from); + } + let allowed_terms = Arc::new(terms.iter().cloned().collect::>()); + let mut residual_input = residual_input.execute(partition, context.clone())?; while let Some(batch) = residual_input.try_next().await? { residual_rows_scanned.add(batch.num_rows()); - let row_ids = batch - .column_by_name(ROW_ID) - .ok_or_else(|| { - DataFusionError::Execution( - "hybrid compound FTS residual input is missing _rowid".to_string(), - ) - })? - .as_primitive::(); - residual.insert_with_row_ids_for_terms(&batch, row_ids, &allowed_terms)?; + residual = index_query_local_residual_batch(residual, batch, allowed_terms.clone()) + .await?; } residual_docs_indexed.add(residual.doc_count()); @@ -955,7 +976,12 @@ impl ExecutionPlan for HybridCompoundQueryExec { Some(metrics.as_ref()), ) .await?; - let residual_stats = residual.bm25_stats_for_terms(&terms); + let stats_terms = terms.clone(); + let (residual, residual_stats) = spawn_cpu(move || { + let stats = residual.bm25_stats_for_terms(&stats_terms); + Ok::<_, Error>((residual, stats)) + }) + .await?; scorer.total_tokens = scorer .total_tokens .checked_add(residual_stats.total_tokens) @@ -1011,9 +1037,20 @@ impl ExecutionPlan for HybridCompoundQueryExec { ) .await?; index_candidates.add(indexed_row_ids.len()); - let residual_leaves = residual.exact_leaf_results(&query, scorer.as_ref())?; - let (residual_row_ids, residual_scores) = - materialized_compound_top_k(&query, residual_leaves, limit, metrics.as_ref())?; + let residual_query = query.clone(); + let residual_scorer = scorer.clone(); + let residual_metrics = metrics.clone(); + let (residual_row_ids, residual_scores) = spawn_cpu(move || { + let residual_leaves = + residual.exact_leaf_results(&residual_query, residual_scorer.as_ref())?; + materialized_compound_top_k( + &residual_query, + residual_leaves, + limit, + residual_metrics.as_ref(), + ) + }) + .await?; residual_candidates.add(residual_row_ids.len()); let mut documents = indexed_row_ids From 8905c134ecb97001f00f0e8e65d1a547b24f7d04 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 13:53:02 +0800 Subject: [PATCH 04/22] test(fts): exercise hybrid prefilter fallback --- rust/lance/src/dataset/tests/dataset_index.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 446fa3ce754..7e69e3c8d11 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2691,11 +2691,12 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { .unwrap() .full_text_search(FullTextSearchQuery::new_query(query.clone())) .unwrap(); + filtered_scanner.prefilter(true); filtered_scanner.limit(Some(2), None).unwrap(); let filtered_plan = filtered_scanner.explain_plan(false).await.unwrap(); assert!( !filtered_plan.contains("HybridCompoundFtsScorer"), - "filtered residual scoring must retain the exact fallback:\n{filtered_plan}" + "prefiltered residual scoring must retain the exact fallback:\n{filtered_plan}" ); let phrase_query: FtsQuery = BooleanQuery::new([ From 7542e47c4aa742e629b18996e4515b0fd6424934 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 14:39:08 +0800 Subject: [PATCH 05/22] fix(fts): reject rewritten hybrid index sources --- rust/lance/src/dataset/scanner.rs | 51 ++++++-- rust/lance/src/dataset/tests/dataset_index.rs | 122 +++++++++++++++++- rust/lance/src/index.rs | 81 +++++++++++- 3 files changed, 243 insertions(+), 11 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index c530c2beac7..ef4ea65de26 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -100,7 +100,6 @@ use crate::dataset::overlay::{collect_overlay_stale_rows_for_segment, overlaid_f use crate::dataset::row_offsets_to_row_addresses; use crate::dataset::rowids::{live_row_addrs_to_row_ids, translate_addr_treemap_to_row_ids}; use crate::dataset::utils::SchemaAdapter; -use crate::index::DatasetIndexInternalExt; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ fts_index_fragment_bitmap, load_segment_details, load_segment_params, load_segments, @@ -110,6 +109,7 @@ use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_frag use crate::index::vector::utils::{ default_distance_type_for, get_vector_dim, get_vector_type, validate_distance_type_for, }; +use crate::index::{DatasetIndexInternalExt, has_append_only_indexed_field_history}; use crate::io::exec::filtered_read::{ FilteredReadExec, FilteredReadOptions, FilteredReadThreadingMode, }; @@ -346,18 +346,20 @@ fn has_exact_hybrid_fts_coverage( fn has_compatible_hybrid_physical_segments( params: &[InvertedIndexParams], + details: &[InvertedIndexDetails], has_deleted_fragments: &[bool], ) -> bool { let Some(first) = params.first() else { return false; }; - params.len() == has_deleted_fragments.len() + params.len() == details.len() + && params.len() == has_deleted_fragments.len() && first.posting_block_size() == 128 && params.iter().all(|params| params == first) - && params.iter().all(|params| { + && details.iter().all(|details| { matches!( - params.resolved_format_version().index_version(), - INVERTED_INDEX_VERSION_V2 | INVERTED_INDEX_VERSION_V3 + details.posting_format_version, + Some(INVERTED_INDEX_VERSION_V2 | INVERTED_INDEX_VERSION_V3) ) }) && has_deleted_fragments.iter().all(|has_deleted| !has_deleted) @@ -4377,6 +4379,18 @@ impl Scanner { // Preserve the established semantic mismatch error before // applying the narrower physical fast-path gate. load_segment_details(&self.dataset, &column, &segments).await?; + if !has_append_only_indexed_field_history(&self.dataset, &segments).await { + // Logical coverage can prune a same-id field rewrite + // while the physical segment still contributes the + // obsolete document to BM25 corpus statistics. + return Ok(None); + } + let physical_details = futures::future::try_join_all( + segments.iter().map(|segment| { + load_physical_fts_details(&self.dataset, &column, segment) + }), + ) + .await?; let segment_params = futures::future::try_join_all( segments .iter() @@ -4411,6 +4425,7 @@ impl Scanner { .await?; if !has_compatible_hybrid_physical_segments( &segment_params, + &physical_details, &has_deleted_fragments, ) { // Larger posting blocks quantize document lengths, and @@ -7374,7 +7389,6 @@ mod test { }; use lance_file::version::LanceFileVersion; use lance_index::optimize::OptimizeOptions; - use lance_index::scalar::inverted::InvertedListFormatVersion; use lance_index::scalar::inverted::query::{ BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, PhraseQuery, }; @@ -7593,27 +7607,48 @@ mod test { #[test] fn test_hybrid_compound_requires_compatible_live_physical_segments() { let params = InvertedIndexParams::default(); + let modern_details = InvertedIndexDetails { + posting_format_version: Some(INVERTED_INDEX_VERSION_V3), + ..Default::default() + }; assert!(has_compatible_hybrid_physical_segments( &[params.clone(), params.clone()], + &[modern_details.clone(), modern_details.clone()], &[false, false] )); assert!(!has_compatible_hybrid_physical_segments( &[params.clone(), params.clone()], + &[modern_details.clone(), modern_details.clone()], &[false, true] )); assert!(!has_compatible_hybrid_physical_segments( &[params.clone(), params.clone().with_position(true)], + &[modern_details.clone(), modern_details.clone()], &[false, false] )); assert!(!has_compatible_hybrid_physical_segments( &[params.clone().block_size(256).unwrap()], + std::slice::from_ref(&modern_details), &[false] )); assert!(!has_compatible_hybrid_physical_segments( - &[params.clone().format_version(InvertedListFormatVersion::V1)], + std::slice::from_ref(¶ms), + &[InvertedIndexDetails { + posting_format_version: Some(1), + ..Default::default() + }], &[false] )); - assert!(!has_compatible_hybrid_physical_segments(&[params], &[])); + assert!(!has_compatible_hybrid_physical_segments( + std::slice::from_ref(¶ms), + &[InvertedIndexDetails::default()], + &[false] + )); + assert!(!has_compatible_hybrid_physical_segments( + &[params], + &[modern_details], + &[] + )); } #[test] diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 7e69e3c8d11..36a9cc57e6c 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -9,12 +9,13 @@ use std::sync::{Arc, Mutex}; use std::vec; use crate::dataset::ROW_ID; +use crate::dataset::WriteDestination; use crate::dataset::builder::DatasetBuilder; use crate::dataset::index::LanceIndexStoreExt; use crate::dataset::optimize::{CompactionOptions, compact_files}; use crate::dataset::tests::dataset_migrations::scan_dataset; use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; -use crate::dataset::transaction::{Operation, Transaction}; +use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; use crate::index::vector::VectorIndexParams; use crate::session::Session; use crate::utils::test::covering; @@ -2857,6 +2858,115 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { ); } +#[tokio::test] +async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { + let initial = arrow_array::record_batch!( + ("text", Utf8, ["stable alpha common", "stale alpha common"]), + ("id", Int32, [0, 1]) + ) + .unwrap(); + let schema = initial.schema(); + let test_uri = TempStrDir::default(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema.clone()), + &test_uri, + Some(WriteParams { + max_rows_per_file: 1, + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .create_index( + &["text"], + IndexType::Inverted, + Some("text_idx".to_string()), + &InvertedIndexParams::default().with_position(true), + true, + ) + .await + .unwrap(); + + let rewritten_fragment = dataset.get_fragment(1).unwrap(); + let mut replacement_file = rewritten_fragment.metadata().files[0].clone(); + replacement_file.path = "replacement.lance".to_string(); + let replacement = arrow_array::record_batch!( + ("text", Utf8, ["replacement beta common"]), + ("id", Int32, [1]) + ) + .unwrap(); + let object_writer = dataset + .object_store + .create(&dataset.data_dir().join(&replacement_file.path)) + .await + .unwrap(); + let mut writer = lance_file::versions::v2_1::create_writer( + object_writer, + schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + writer.write_batch(&replacement).await.unwrap(); + writer.finish().await.unwrap(); + + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(1, replacement_file)], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + let committed = dataset + .load_index_by_name("text_idx") + .await + .unwrap() + .unwrap(); + assert!( + !committed.fragment_bitmap.as_ref().unwrap().contains(1), + "the logical index must prune the same-id rewritten fragment" + ); + + let appended = + arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [2])) + .unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), + None, + ) + .await + .unwrap(); + + let query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("common", "text", 1.0)), + (Occur::Should, compound_match_query("alpha", "text", 1.0)), + (Occur::Should, compound_match_query("beta", "text", 1.0)), + (Occur::MustNot, compound_match_query("stale", "text", 1.0)), + ]) + .into(); + let plan = compound_fts_plan(&dataset, query.clone(), 2).await; + assert!( + !plan.contains("HybridCompoundFtsScorer"), + "same-id indexed-field rewrites make physical BM25 stats unsafe:\n{plan}" + ); + + let actual = compound_fts_results(&dataset, query.clone(), Some(2)).await; + let mut flat_oracle = compound_fts_results(&dataset, query, None).await; + flat_oracle.truncate(2); + assert_eq!( + actual, flat_oracle, + "bounded fallback must preserve the flat path's ordered row ids and scores" + ); +} + #[tokio::test] async fn test_partial_compound_hybrid_rejects_retired_physical_fragments() { let initial = arrow_array::record_batch!( @@ -3830,6 +3940,16 @@ async fn test_fts_v1_remains_queryable_after_append_optimize() { let schema = batch.schema(); let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); dataset.append(batches, None).await.unwrap(); + let compound_query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("alpha", "text", 1.0)), + (Occur::Should, compound_match_query("original", "text", 1.0)), + ]) + .into(); + let plan = compound_fts_plan(&dataset, compound_query, 2).await; + assert!( + !plan.contains("HybridCompoundFtsScorer"), + "a physical FTS v1 segment must not enter the modern hybrid path:\n{plan}" + ); dataset .optimize_indices(&OptimizeOptions::append()) .await diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index d10ec624d77..6dbbe9ec700 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -164,9 +164,9 @@ fn fragment_field_paths<'a>( /// segment's carried columns can go stale independently of its keyed column, so /// checking only the keyed subtree would leave a fragment covered after a carried /// column was rewritten, and the segment would answer with the obsolete value. -fn segment_indexed_field_ids(dataset: &Dataset, segment: &IndexSegment) -> Result> { +fn indexed_field_ids(dataset: &Dataset, fields: &[i32]) -> Result> { let mut indexed_field_ids = HashSet::new(); - for field_id in segment.fields() { + for field_id in fields { let field = dataset.schema().field_by_id(*field_id).ok_or_else(|| { Error::invalid_input(format!( "CreateIndex: field id {field_id} does not exist in the current schema" @@ -177,6 +177,83 @@ fn segment_indexed_field_ids(dataset: &Dataset, segment: &IndexSegment) -> Resul Ok(indexed_field_ids) } +fn segment_indexed_field_ids(dataset: &Dataset, segment: &IndexSegment) -> Result> { + indexed_field_ids(dataset, segment.fields()) +} + +/// Prove that every indexed-field file present at a segment's build version is +/// still the current file for the same fragment id. +/// +/// A committed segment's logical bitmap may already have pruned a rewritten +/// fragment even though its physical postings and corpus statistics still +/// contain that fragment. Since the original per-segment coverage is no longer +/// available after that pruning, this check deliberately considers every +/// fragment present at each build version. That is conservative, but it makes +/// the hybrid scorer available only when the history can be proven to be a +/// pure append with respect to the indexed field subtrees. +pub(crate) async fn has_append_only_indexed_field_history( + dataset: &Dataset, + segments: &[IndexMetadata], +) -> bool { + let current_version = dataset.manifest.version; + let current_fragments = dataset + .fragments() + .iter() + .filter_map(|fragment| u32::try_from(fragment.id).ok().map(|id| (id, fragment))) + .collect::>(); + if current_fragments.len() != dataset.fragments().len() { + return false; + } + + let build_versions = segments + .iter() + .map(|segment| segment.dataset_version) + .collect::>(); + for build_version in build_versions { + if build_version > current_version { + return false; + } + if build_version == current_version { + continue; + } + let historical = match dataset.checkout_version(build_version).await { + Ok(historical) => historical, + Err(_) => return false, + }; + let mut indexed_field_ids_at_version = HashSet::new(); + for segment in segments + .iter() + .filter(|segment| segment.dataset_version == build_version) + { + let Ok(historical_field_ids) = indexed_field_ids(&historical, &segment.fields) else { + return false; + }; + let Ok(current_field_ids) = indexed_field_ids(dataset, &segment.fields) else { + return false; + }; + if historical_field_ids != current_field_ids { + return false; + } + indexed_field_ids_at_version.extend(historical_field_ids); + } + + for historical_fragment in historical.fragments() { + let Ok(fragment_id) = u32::try_from(historical_fragment.id) else { + return false; + }; + let Some(current_fragment) = current_fragments.get(&fragment_id) else { + return false; + }; + if fragment_field_paths(historical_fragment, &indexed_field_ids_at_version) + != fragment_field_paths(current_fragment, &indexed_field_ids_at_version) + { + return false; + } + } + } + true +} + async fn prune_stale_segment_coverage( dataset: &Dataset, segments: &mut [IndexSegment], From f8ac1ada6dc930c5cdb16a458e987b42818a3c42 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 15:01:31 +0800 Subject: [PATCH 06/22] fix(fts): include data base in hybrid rewrite checks --- rust/lance/src/dataset/tests/dataset_index.rs | 200 +++++++++++++++++- rust/lance/src/index.rs | 43 +++- 2 files changed, 226 insertions(+), 17 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 36a9cc57e6c..6a6223c112f 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -67,6 +67,7 @@ use futures::{StreamExt, TryStreamExt}; use itertools::Itertools; use lance_arrow::json::ARROW_JSON_EXT_NAME; use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; +use lance_table::format::BasePath; use lance_testing::datagen::generate_random_array; use rand::Rng; use roaring::RoaringBitmap; @@ -1330,6 +1331,12 @@ async fn compound_fts_results( .collect() } +fn scored_row_bits(rows: &[(u64, f32)]) -> Vec<(u64, u32)> { + rows.iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect() +} + fn compound_fts_result_bits(batch: &RecordBatch) -> Vec<(u64, u32)> { let row_ids = batch[ROW_ID].as_primitive::().values(); let scores = batch[SCORE_COL].as_primitive::().values(); @@ -2841,19 +2848,14 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { .unwrap(); let rebuilt_boost = compound_fts_results(&dataset, boost_query, Some(10)).await; let rebuilt_multimatch = compound_fts_results(&dataset, multimatch_query, Some(3)).await; - let score_bits = |rows: &[(u64, f32)]| { - rows.iter() - .map(|(row_id, score)| (*row_id, score.to_bits())) - .collect::>() - }; assert_eq!( - score_bits(&partial_boost), - score_bits(&rebuilt_boost), + scored_row_bits(&partial_boost), + scored_row_bits(&rebuilt_boost), "hybrid Boost scores must be bit-identical to a rebuilt index" ); assert_eq!( - score_bits(&partial_multimatch), - score_bits(&rebuilt_multimatch), + scored_row_bits(&partial_multimatch), + scored_row_bits(&rebuilt_multimatch), "hybrid MultiMatch scores must be bit-identical to a rebuilt index" ); } @@ -2965,6 +2967,186 @@ async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { actual, flat_oracle, "bounded fallback must preserve the flat path's ordered row ids and scores" ); + + let exact_oracle_query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("common", "text", 0.0)), + (Occur::Should, compound_match_query("alpha", "text", 0.0)), + (Occur::Should, compound_match_query("beta", "text", 0.0)), + (Occur::MustNot, compound_match_query("stale", "text", 0.0)), + ]) + .into(); + let exact_actual = compound_fts_results(&dataset, exact_oracle_query.clone(), Some(3)).await; + let mut rebuilt = dataset.clone(); + rebuilt + .create_index( + &["text"], + IndexType::Inverted, + Some("text_idx".to_string()), + &InvertedIndexParams::default().with_position(true), + true, + ) + .await + .unwrap(); + let exact_expected = compound_fts_results(&rebuilt, exact_oracle_query, Some(3)).await; + assert_eq!( + scored_row_bits(&exact_actual), + scored_row_bits(&exact_expected), + "same-id rewrite fallback must match a rebuilt index exactly" + ); +} + +#[tokio::test] +async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() { + let primary = TempStrDir::default(); + let base_one = TempStrDir::default(); + let base_two = TempStrDir::default(); + let initial = arrow_array::record_batch!( + ("text", Utf8, ["stable alpha common", "stale alpha common"]), + ("id", Int32, [0, 1]) + ) + .unwrap(); + let schema = initial.schema(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema.clone()), + &primary, + Some(WriteParams { + max_rows_per_file: 1, + initial_bases: Some(vec![ + BasePath::new(1, base_one.to_string(), Some("base-one".to_string()), false), + BasePath::new(2, base_two.to_string(), Some("base-two".to_string()), false), + ]), + target_bases: Some(vec![1]), + ..Default::default() + }), + ) + .await + .unwrap(); + assert!(dataset.get_fragments().iter().all(|fragment| { + fragment + .metadata() + .files + .iter() + .all(|file| file.base_id == Some(1)) + })); + + let columns = ["text"]; + let params = InvertedIndexParams::default().with_position(true); + let segment = dataset + .create_index_builder(&columns, IndexType::Inverted, ¶ms) + .name("text_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + + let rewritten_fragment = dataset.get_fragment(1).unwrap(); + let relative_path = rewritten_fragment.metadata().files[0].path.clone(); + let replacement = arrow_array::record_batch!( + ("text", Utf8, ["replacement beta common"]), + ("id", Int32, [1]) + ) + .unwrap(); + let replacement_path = dataset + .data_file_dir_for_base(Some(2)) + .unwrap() + .join(&relative_path); + let object_writer = dataset + .object_store(Some(2)) + .await + .unwrap() + .create(&replacement_path) + .await + .unwrap(); + let mut writer = lance_file::versions::v2_1::create_writer( + object_writer, + schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + writer.write_batch(&replacement).await.unwrap(); + writer.finish().await.unwrap(); + let replacement_file = dataset + .create_data_file(&relative_path, Some(2)) + .await + .unwrap(); + assert_eq!(replacement_file.path, relative_path); + assert_eq!(replacement_file.base_id, Some(2)); + + let read_version = dataset.manifest.version; + let mut dataset = Dataset::commit( + WriteDestination::Dataset(Arc::new(dataset)), + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(1, replacement_file)], + }, + Some(read_version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + let rewritten_file = &dataset.get_fragment(1).unwrap().metadata().files[0]; + assert_eq!(rewritten_file.path, relative_path); + assert_eq!(rewritten_file.base_id, Some(2)); + + dataset + .commit_existing_index_segments("text_idx", "text", vec![segment]) + .await + .unwrap(); + let committed = dataset + .load_index_by_name("text_idx") + .await + .unwrap() + .unwrap(); + let coverage = committed.fragment_bitmap.as_ref().unwrap(); + assert!(coverage.contains(0)); + assert!( + !coverage.contains(1), + "changing only the registered base must prune stale logical coverage" + ); + + let appended = + arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [2])) + .unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), + None, + ) + .await + .unwrap(); + + let query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("common", "text", 0.0)), + (Occur::Should, compound_match_query("alpha", "text", 0.0)), + (Occur::Should, compound_match_query("beta", "text", 0.0)), + (Occur::MustNot, compound_match_query("stale", "text", 0.0)), + ]) + .into(); + let plan = compound_fts_plan(&dataset, query.clone(), 3).await; + assert!( + !plan.contains("HybridCompoundFtsScorer"), + "same-path files from different bases make physical BM25 stats unsafe:\n{plan}" + ); + let actual = compound_fts_results(&dataset, query.clone(), Some(3)).await; + + let mut rebuilt = dataset.clone(); + rebuilt + .create_index( + &["text"], + IndexType::Inverted, + Some("text_idx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + let expected = compound_fts_results(&rebuilt, query, Some(3)).await; + assert_eq!( + scored_row_bits(&actual), + scored_row_bits(&expected), + "different-base rewrite fallback must match a rebuilt index exactly" + ); } #[tokio::test] diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 6dbbe9ec700..3307e98fa92 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -58,7 +58,7 @@ use lance_io::utils::{ CachedFileSize, read_last_block, read_message, read_message_from_buf, read_metadata_offset, read_version, }; -use lance_table::format::{Fragment, SelfDescribingFileReader}; +use lance_table::format::{DataFile, Fragment, SelfDescribingFileReader}; use lance_table::format::{IndexFile, IndexMetadata, list_index_files_with_sizes}; use lance_table::io::manifest::read_manifest_indexes; use roaring::RoaringBitmap; @@ -142,10 +142,37 @@ fn collect_subtree_field_ids(field: &Field, field_ids: &mut HashSet) { } } -fn fragment_field_paths<'a>( +/// Stable identity fields for a physical data file. +/// +/// This mirrors transaction rewrite validation and deliberately excludes +/// `file_size_bytes`, which is a mutable cache rather than file identity. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct PhysicalDataFileIdentity<'a> { + base_id: Option, + path: &'a str, + fields: &'a [i32], + column_indices: &'a [i32], + file_major_version: u32, + file_minor_version: u32, +} + +impl<'a> From<&'a DataFile> for PhysicalDataFileIdentity<'a> { + fn from(file: &'a DataFile) -> Self { + Self { + base_id: file.base_id, + path: &file.path, + fields: file.fields.as_ref(), + column_indices: file.column_indices.as_ref(), + file_major_version: file.file_major_version, + file_minor_version: file.file_minor_version, + } + } +} + +fn fragment_field_files<'a>( fragment: &'a Fragment, indexed_field_ids: &HashSet, -) -> HashMap { +) -> HashMap> { fragment .files .iter() @@ -153,7 +180,7 @@ fn fragment_field_paths<'a>( file.fields .iter() .filter(|field_id| indexed_field_ids.contains(field_id)) - .map(|field_id| (*field_id, file.path.as_str())) + .map(|field_id| (*field_id, file.into())) }) .collect() } @@ -244,8 +271,8 @@ pub(crate) async fn has_append_only_indexed_field_history( let Some(current_fragment) = current_fragments.get(&fragment_id) else { return false; }; - if fragment_field_paths(historical_fragment, &indexed_field_ids_at_version) - != fragment_field_paths(current_fragment, &indexed_field_ids_at_version) + if fragment_field_files(historical_fragment, &indexed_field_ids_at_version) + != fragment_field_files(current_fragment, &indexed_field_ids_at_version) { return false; } @@ -299,8 +326,8 @@ async fn prune_stale_segment_coverage( return true; }; let changed_files = - fragment_field_paths(historical_fragment, &indexed_field_ids) - != fragment_field_paths(current_fragment, &indexed_field_ids); + fragment_field_files(historical_fragment, &indexed_field_ids) + != fragment_field_files(current_fragment, &indexed_field_ids); let changed_overlays = prune_newer_overlays && current_fragment.overlays.iter().any(|overlay| { overlay.committed_version > version From d2609f39ec0573b95fcf011c159c1ca05f82c35c Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 15:12:18 +0800 Subject: [PATCH 07/22] fix(fts): reject rebound hybrid data bases --- rust/lance/src/dataset/tests/dataset_index.rs | 164 ++++++++++++++++++ rust/lance/src/index.rs | 58 +++++-- 2 files changed, 210 insertions(+), 12 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 6a6223c112f..b10a100e1d1 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -58,6 +58,7 @@ use lance_index::scalar::{ }; use lance_index::{FtsPrewarmOptions, PrewarmOptions}; use lance_index::{IndexType, scalar::ScalarIndexParams, vector::DIST_COL}; +use lance_io::object_store::ObjectStore; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; use lance_linalg::distance::MetricType; @@ -3149,6 +3150,169 @@ async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() ); } +#[tokio::test] +async fn test_partial_compound_hybrid_rejects_rebound_registered_base() { + let primary = TempStrDir::default(); + let base_a = TempStrDir::default(); + let base_b = TempStrDir::default(); + let stable = + arrow_array::record_batch!(("text", Utf8, ["stable alpha common"]), ("id", Int32, [0])) + .unwrap(); + let schema = stable.schema(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![stable].into_iter().map(Ok), schema.clone()), + &primary, + Some(WriteParams { + max_rows_per_file: 1, + initial_bases: Some(vec![BasePath::new( + 1, + base_a.to_string(), + Some("base-a".to_string()), + false, + )]), + ..Default::default() + }), + ) + .await + .unwrap(); + let indexed_on_base = + arrow_array::record_batch!(("text", Utf8, ["stale alpha common"]), ("id", Int32, [1])) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![indexed_on_base].into_iter().map(Ok), schema.clone()), + Arc::new(dataset), + Some(WriteParams { + mode: WriteMode::Append, + target_bases: Some(vec![1]), + ..Default::default() + }), + ) + .await + .unwrap(); + assert_eq!( + dataset.get_fragment(0).unwrap().metadata().files[0].base_id, + None + ); + assert_eq!( + dataset.get_fragment(1).unwrap().metadata().files[0].base_id, + Some(1) + ); + + let columns = ["text"]; + let params = InvertedIndexParams::default().with_position(true); + let segment = dataset + .create_index_builder(&columns, IndexType::Inverted, ¶ms) + .name("text_idx".to_string()) + .execute_uncommitted() + .await + .unwrap(); + + let fragment_before = dataset.get_fragment(1).unwrap().metadata().clone(); + let relative_path = fragment_before.files[0].path.clone(); + let replacement = arrow_array::record_batch!( + ("text", Utf8, ["replacement beta common"]), + ("id", Int32, [1]) + ) + .unwrap(); + let (base_b_store, base_b_root) = ObjectStore::from_uri(&base_b).await.unwrap(); + let object_writer = base_b_store + .create(&base_b_root.join(&relative_path)) + .await + .unwrap(); + let mut writer = lance_file::versions::v2_1::create_writer( + object_writer, + schema.as_ref().try_into().unwrap(), + Default::default(), + ) + .unwrap(); + writer.write_batch(&replacement).await.unwrap(); + writer.finish().await.unwrap(); + + let binding_before = dataset.manifest.base_paths.get(&1).unwrap().clone(); + dataset = Arc::new(dataset) + .add_bases( + vec![BasePath::new( + 1, + base_b.to_string(), + Some("base-b".to_string()), + false, + )], + None, + ) + .await + .unwrap(); + let binding_after = dataset.manifest.base_paths.get(&1).unwrap(); + assert_eq!(binding_before.path, base_a.to_string()); + assert_eq!(binding_after.path, base_b.to_string()); + assert_eq!( + binding_before.is_dataset_root, + binding_after.is_dataset_root + ); + assert_eq!( + dataset.get_fragment(1).unwrap().metadata(), + &fragment_before, + "UpdateBases must leave the DataFile identity unchanged" + ); + + dataset + .commit_existing_index_segments("text_idx", "text", vec![segment]) + .await + .unwrap(); + let committed = dataset + .load_index_by_name("text_idx") + .await + .unwrap() + .unwrap(); + let coverage = committed.fragment_bitmap.as_ref().unwrap(); + assert!(coverage.contains(0)); + assert!( + !coverage.contains(1), + "rebinding a referenced base must prune stale logical coverage" + ); + + let appended = + arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [2])) + .unwrap(); + dataset + .append( + RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), + None, + ) + .await + .unwrap(); + let query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("common", "text", 0.0)), + (Occur::Should, compound_match_query("alpha", "text", 0.0)), + (Occur::Should, compound_match_query("beta", "text", 0.0)), + (Occur::MustNot, compound_match_query("stale", "text", 0.0)), + ]) + .into(); + let plan = compound_fts_plan(&dataset, query.clone(), 3).await; + assert!( + !plan.contains("HybridCompoundFtsScorer"), + "a rebound registered base makes physical BM25 stats unsafe:\n{plan}" + ); + let actual = compound_fts_results(&dataset, query.clone(), Some(3)).await; + + let mut rebuilt = dataset.clone(); + rebuilt + .create_index( + &["text"], + IndexType::Inverted, + Some("text_idx".to_string()), + ¶ms, + true, + ) + .await + .unwrap(); + let expected = compound_fts_results(&rebuilt, query, Some(3)).await; + assert_eq!( + scored_row_bits(&actual), + scored_row_bits(&expected), + "rebound-base fallback must match a rebuilt index exactly" + ); +} + #[tokio::test] async fn test_partial_compound_hybrid_rejects_retired_physical_fragments() { let initial = arrow_array::record_batch!( diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 3307e98fa92..3a60337e41c 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -144,11 +144,22 @@ fn collect_subtree_field_ids(field: &Field, field_ids: &mut HashSet) { /// Stable identity fields for a physical data file. /// -/// This mirrors transaction rewrite validation and deliberately excludes +/// This mirrors transaction rewrite validation, additionally resolves a +/// registered base to its physical binding, and deliberately excludes /// `file_size_bytes`, which is a mutable cache rather than file identity. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PhysicalBaseBinding<'a> { + Primary, + Registered { + path: &'a str, + is_dataset_root: bool, + }, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct PhysicalDataFileIdentity<'a> { base_id: Option, + base_binding: PhysicalBaseBinding<'a>, path: &'a str, fields: &'a [i32], column_indices: &'a [i32], @@ -156,23 +167,35 @@ struct PhysicalDataFileIdentity<'a> { file_minor_version: u32, } -impl<'a> From<&'a DataFile> for PhysicalDataFileIdentity<'a> { - fn from(file: &'a DataFile) -> Self { - Self { +impl<'a> PhysicalDataFileIdentity<'a> { + fn try_new(dataset: &'a Dataset, file: &'a DataFile) -> Option { + let base_binding = match file.base_id { + Some(base_id) => { + let base = dataset.manifest.base_paths.get(&base_id)?; + PhysicalBaseBinding::Registered { + path: &base.path, + is_dataset_root: base.is_dataset_root, + } + } + None => PhysicalBaseBinding::Primary, + }; + Some(Self { base_id: file.base_id, + base_binding, path: &file.path, fields: file.fields.as_ref(), column_indices: file.column_indices.as_ref(), file_major_version: file.file_major_version, file_minor_version: file.file_minor_version, - } + }) } } fn fragment_field_files<'a>( + dataset: &'a Dataset, fragment: &'a Fragment, indexed_field_ids: &HashSet, -) -> HashMap> { +) -> Option>> { fragment .files .iter() @@ -180,7 +203,10 @@ fn fragment_field_files<'a>( file.fields .iter() .filter(|field_id| indexed_field_ids.contains(field_id)) - .map(|field_id| (*field_id, file.into())) + .map(|field_id| { + PhysicalDataFileIdentity::try_new(dataset, file) + .map(|identity| (*field_id, identity)) + }) }) .collect() } @@ -271,9 +297,14 @@ pub(crate) async fn has_append_only_indexed_field_history( let Some(current_fragment) = current_fragments.get(&fragment_id) else { return false; }; - if fragment_field_files(historical_fragment, &indexed_field_ids_at_version) - != fragment_field_files(current_fragment, &indexed_field_ids_at_version) - { + let historical_files = fragment_field_files( + &historical, + historical_fragment, + &indexed_field_ids_at_version, + ); + let current_files = + fragment_field_files(dataset, current_fragment, &indexed_field_ids_at_version); + if historical_files.is_none() || historical_files != current_files { return false; } } @@ -325,9 +356,12 @@ async fn prune_stale_segment_coverage( let Some(current_fragment) = current_fragments.get(fragment_id) else { return true; }; + let historical_files = + fragment_field_files(&historical, historical_fragment, &indexed_field_ids); + let current_files = + fragment_field_files(dataset, current_fragment, &indexed_field_ids); let changed_files = - fragment_field_files(historical_fragment, &indexed_field_ids) - != fragment_field_files(current_fragment, &indexed_field_ids); + historical_files.is_none() || historical_files != current_files; let changed_overlays = prune_newer_overlays && current_fragment.overlays.iter().any(|overlay| { overlay.committed_version > version From 85097083613a9c19ec3da882063b86f0dcf0851d Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 20:36:27 +0800 Subject: [PATCH 08/22] perf(fts): remove hybrid benchmark metrics --- rust/lance/src/dataset/tests/dataset_index.rs | 31 +------------------ rust/lance/src/io/exec/fts.rs | 29 ----------------- 2 files changed, 1 insertion(+), 59 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index b10a100e1d1..6a23dc35b49 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2650,25 +2650,6 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { &[0, 2], "exact search should include the appended hit" ); - let (_, exact_stats) = compound_fts_results_with_stats(&dataset, query.clone(), 2).await; - assert_eq!( - exact_stats - .all_counts - .get(crate::io::exec::fts::HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC), - Some(&1) - ); - assert_eq!( - exact_stats - .all_counts - .get(crate::io::exec::fts::HYBRID_COMPOUND_RESIDUAL_DOCS_INDEXED_METRIC), - Some(&1) - ); - assert_eq!( - exact_stats - .all_counts - .get(crate::io::exec::fts::HYBRID_COMPOUND_MERGED_CANDIDATES_METRIC), - Some(&2) - ); let empty_terms_query: FtsQuery = BooleanQuery::new([ (Occur::Must, compound_match_query("", "text", 1.0)), @@ -2680,18 +2661,8 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { empty_terms_plan.contains("HybridCompoundFtsScorer"), "the empty analyzed-term case must exercise the hybrid short circuit:\n{empty_terms_plan}" ); - let (empty_results, empty_stats) = - compound_fts_results_with_stats(&dataset, empty_terms_query, 2).await; + let empty_results = compound_fts_results(&dataset, empty_terms_query, Some(2)).await; assert!(empty_results.is_empty()); - assert_eq!( - empty_stats - .all_counts - .get(crate::io::exec::fts::HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC) - .copied() - .unwrap_or_default(), - 0, - "an empty analyzed query must not poll the residual scan" - ); let mut filtered_scanner = dataset.scan(); filtered_scanner diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 4a7370c9d22..7351086cbc7 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -917,21 +917,6 @@ impl ExecutionPlan for HybridCompoundQueryExec { let segments = self.segments.clone(); let residual_input = self.residual_input.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); - let residual_rows_scanned = self - .metrics - .new_count(HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC, partition); - let residual_docs_indexed = self - .metrics - .new_count(HYBRID_COMPOUND_RESIDUAL_DOCS_INDEXED_METRIC, partition); - let index_candidates = self - .metrics - .new_count(HYBRID_COMPOUND_INDEX_CANDIDATES_METRIC, partition); - let residual_candidates = self - .metrics - .new_count(HYBRID_COMPOUND_RESIDUAL_CANDIDATES_METRIC, partition); - let merged_candidates = self - .metrics - .new_count(HYBRID_COMPOUND_MERGED_CANDIDATES_METRIC, partition); let schema = self.schema(); let stream = stream::once(async move { @@ -958,11 +943,9 @@ impl ExecutionPlan for HybridCompoundQueryExec { let mut residual_input = residual_input.execute(partition, context.clone())?; while let Some(batch) = residual_input.try_next().await? { - residual_rows_scanned.add(batch.num_rows()); residual = index_query_local_residual_batch(residual, batch, allowed_terms.clone()) .await?; } - residual_docs_indexed.add(residual.doc_count()); let query_tokens = Tokens::new(terms.clone(), first_index.tokenizer().doc_type()); let exact_params = params @@ -1036,7 +1019,6 @@ impl ExecutionPlan for HybridCompoundQueryExec { scorer.clone(), ) .await?; - index_candidates.add(indexed_row_ids.len()); let residual_query = query.clone(); let residual_scorer = scorer.clone(); let residual_metrics = metrics.clone(); @@ -1051,7 +1033,6 @@ impl ExecutionPlan for HybridCompoundQueryExec { ) }) .await?; - residual_candidates.add(residual_row_ids.len()); let mut documents = indexed_row_ids .into_iter() @@ -1059,7 +1040,6 @@ impl ExecutionPlan for HybridCompoundQueryExec { .chain(residual_row_ids.into_iter().zip(residual_scores)) .map(|(row_id, score)| ScoredDoc::new(row_id, score)) .collect::>(); - merged_candidates.add(documents.len()); documents.sort_unstable_by(|left, right| { right .score @@ -2170,15 +2150,6 @@ impl Drop for SharedFtsScorerProducer { /// Time spent resolving an exact ordered UUID selection to committed FTS segments. pub const FTS_SEGMENT_BIND_DURATION_METRIC: &str = "fts_segment_bind_duration"; -pub(crate) const HYBRID_COMPOUND_RESIDUAL_ROWS_SCANNED_METRIC: &str = - "hybrid_compound_residual_rows_scanned"; -pub(crate) const HYBRID_COMPOUND_RESIDUAL_DOCS_INDEXED_METRIC: &str = - "hybrid_compound_residual_docs_indexed"; -pub(crate) const HYBRID_COMPOUND_INDEX_CANDIDATES_METRIC: &str = "hybrid_compound_index_candidates"; -pub(crate) const HYBRID_COMPOUND_RESIDUAL_CANDIDATES_METRIC: &str = - "hybrid_compound_residual_candidates"; -pub(crate) const HYBRID_COMPOUND_MERGED_CANDIDATES_METRIC: &str = - "hybrid_compound_merged_candidates"; #[derive(Debug, Clone)] enum FtsSegmentSelection { From 7ee4c59c08b988ce176e6f1a770562e8df866e2d Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 23:17:45 +0800 Subject: [PATCH 09/22] fix(fts): repair hybrid integration after restack --- rust/lance/src/index.rs | 2 +- rust/lance/src/io/exec/fts.rs | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 3a60337e41c..a195dcb3da6 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -290,7 +290,7 @@ pub(crate) async fn has_append_only_indexed_field_history( indexed_field_ids_at_version.extend(historical_field_ids); } - for historical_fragment in historical.fragments() { + for historical_fragment in historical.fragments().iter() { let Ok(fragment_id) = u32::try_from(historical_fragment.id) else { return false; }; diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 7351086cbc7..8c73f66e3a2 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -60,8 +60,7 @@ use lance_index::scalar::inverted::{ DOC_INDEX_COL, DocumentGranularity, FTS_SCHEMA, FlatBm25SearchOptions, InvertedIndex, MemBM25Scorer, PreparedBm25Query, SCORE_COL, Scorer, build_global_bm25_scorer, compound_search, compound_search_prepared_match, compound_search_prepared_match_with_score_floor, - compound_search_with_base_scorer, compound_search_with_base_scorer_and_score_floor, - cross_column_compound_search, exclusive_scaled_score_floor, + compound_search_with_base_scorer, cross_column_compound_search, exclusive_scaled_score_floor, flat_bm25_search_stream_with_options_and_scorer, fts_schema, materialized_compound_top_k, prepare_bm25_query, }; @@ -1007,8 +1006,10 @@ impl ExecutionPlan for HybridCompoundQueryExec { &PreFilterSource::None, dataset, &segments, - None, - None, + PreFilterMasks { + overlay_block: None, + external_mask: None, + }, )?; let (indexed_row_ids, indexed_scores) = compound_search_with_base_scorer( &indices, From 8879063baae916d03aeb3666e2849bd79bdf1ba1 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Thu, 27 Aug 2026 23:25:26 +0800 Subject: [PATCH 10/22] fix(fts): address hybrid CI diagnostics --- rust/lance/src/dataset/mem_wal/index/fts.rs | 5 +++-- rust/lance/src/dataset/tests/dataset_index.rs | 6 +++--- rust/lance/src/index.rs | 5 ++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index f822ebd8b1d..5bffeaed517 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1023,7 +1023,7 @@ pub struct FtsMemIndex { /// starts a detached tiered merge; dropping the query drops all residual /// postings. #[derive(Debug)] -pub(crate) struct QueryLocalFtsIndex { +pub struct QueryLocalFtsIndex { inner: FtsMemIndex, } @@ -1057,7 +1057,8 @@ impl QueryLocalFtsIndex { .insert_with_row_ids_for_terms(batch, row_ids, terms) } - pub(crate) fn doc_count(&self) -> usize { + #[cfg(test)] + fn doc_count(&self) -> usize { self.inner.doc_count() } diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 6a23dc35b49..a0b240d2380 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2872,7 +2872,7 @@ async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { .unwrap(); let object_writer = dataset .object_store - .create(&dataset.data_dir().join(&replacement_file.path)) + .create(&dataset.data_dir().join(replacement_file.path.as_str())) .await .unwrap(); let mut writer = lance_file::versions::v2_1::create_writer( @@ -3020,7 +3020,7 @@ async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() let replacement_path = dataset .data_file_dir_for_base(Some(2)) .unwrap() - .join(&relative_path); + .join(relative_path.as_str()); let object_writer = dataset .object_store(Some(2)) .await @@ -3187,7 +3187,7 @@ async fn test_partial_compound_hybrid_rejects_rebound_registered_base() { .unwrap(); let (base_b_store, base_b_root) = ObjectStore::from_uri(&base_b).await.unwrap(); let object_writer = base_b_store - .create(&base_b_root.join(&relative_path)) + .create(&base_b_root.join(relative_path.as_str())) .await .unwrap(); let mut writer = lance_file::versions::v2_1::create_writer( diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index a195dcb3da6..18bcf509569 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -269,9 +269,8 @@ pub(crate) async fn has_append_only_indexed_field_history( if build_version == current_version { continue; } - let historical = match dataset.checkout_version(build_version).await { - Ok(historical) => historical, - Err(_) => return false, + let Ok(historical) = dataset.checkout_version(build_version).await else { + return false; }; let mut indexed_field_ids_at_version = HashSet::new(); for segment in segments From 73699440d5eb43ac7855aedc72a31c974017c9ed Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 00:51:53 +0800 Subject: [PATCH 11/22] fix(fts): update hybrid scorer after metrics cleanup --- rust/lance-index/src/scalar/inverted/compound.rs | 5 +---- rust/lance/src/io/exec/fts.rs | 8 +------- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/rust/lance-index/src/scalar/inverted/compound.rs b/rust/lance-index/src/scalar/inverted/compound.rs index 2b6653aa40e..7bbba53ee52 100644 --- a/rust/lance-index/src/scalar/inverted/compound.rs +++ b/rust/lance-index/src/scalar/inverted/compound.rs @@ -2172,7 +2172,6 @@ pub fn materialized_compound_top_k( query: &FtsQuery, leaves: Vec>, limit: usize, - metrics: &dyn MetricsCollector, ) -> Result<(Vec, Vec)> { let mut leaf_count = 0; let plan = CompoundScorerPlan::from_query(query, &mut leaf_count)?; @@ -2192,7 +2191,7 @@ pub fn materialized_compound_top_k( MaterializedScorer::try_new(rows).map(|scorer| Some(Box::new(scorer) as BoxScorer<'_>)) }) .collect::>>()?; - let mut scorer = plan.build(&mut scorers, metrics)?; + let mut scorer = plan.build(&mut scorers)?; let rows = TopKCollector::new(limit).collect(scorer.as_mut())?; Ok(rows.into_iter().map(|row| (row.row_id, row.score)).unzip()) } @@ -4470,12 +4469,10 @@ mod tests { MatchQuery::new("alpha".to_string()).with_column(Some("text".to_string())), ], }); - let metrics = NoOpMetricsCollector; let (row_ids, scores) = materialized_compound_top_k( &query, vec![vec![(7, 1.0), (3, 2.0)], vec![(7, 3.0), (5, 3.0)]], 2, - &metrics, ) .unwrap(); diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 8c73f66e3a2..7c6e68ccc00 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -1022,16 +1022,10 @@ impl ExecutionPlan for HybridCompoundQueryExec { .await?; let residual_query = query.clone(); let residual_scorer = scorer.clone(); - let residual_metrics = metrics.clone(); let (residual_row_ids, residual_scores) = spawn_cpu(move || { let residual_leaves = residual.exact_leaf_results(&residual_query, residual_scorer.as_ref())?; - materialized_compound_top_k( - &residual_query, - residual_leaves, - limit, - residual_metrics.as_ref(), - ) + materialized_compound_top_k(&residual_query, residual_leaves, limit) }) .await?; From 8dbeea203f365403c1cc81844fa80167e2233157 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 17:49:49 +0800 Subject: [PATCH 12/22] fix(fts): preserve physical provenance across remaps --- protos/index_old.proto | 7 + .../src/scalar/inverted/tokenizer.rs | 2 + rust/lance/src/dataset/tests/dataset_index.rs | 62 ++++++-- rust/lance/src/index.rs | 53 +++++-- rust/lance/src/index/append.rs | 85 +++++++---- rust/lance/src/index/create.rs | 12 +- rust/lance/src/index/scalar/inverted.rs | 139 +++++++++++++++++- 7 files changed, 306 insertions(+), 54 deletions(-) diff --git a/protos/index_old.proto b/protos/index_old.proto index 236d1f110c0..5722b2525f1 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -101,4 +101,11 @@ message InvertedIndexDetails { // The posting-list payload format. This is separate from index_version, // which identifies the overall inverted-index layout. optional uint32 posting_format_version = 15; + // Dataset versions whose physical documents contribute to this segment. + // Unlike IndexMetadata.dataset_version, this provenance is not advanced by + // row-address remapping because remapping can retain documents outside its + // address map. Writers sort and deduplicate this set. An empty set means + // provenance is unknown; readers must not infer it from + // IndexMetadata.dataset_version. + repeated uint64 physical_source_dataset_versions = 16; } diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index fd9b9294d7b..818ce466511 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -427,6 +427,7 @@ impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { ), document_granularity: PbDocumentGranularity::from(params.document_granularity) as i32, posting_format_version: Some(params.resolved_format_version().index_version()), + physical_source_dataset_versions: Vec::new(), }) } } @@ -1559,6 +1560,7 @@ mod tests { code_config: None, document_granularity: PbDocumentGranularity::Row as i32, posting_format_version: None, + physical_source_dataset_versions: Vec::new(), }; let params = InvertedIndexParams::try_from(&old_details).unwrap(); assert_eq!(params.block_size, 128); diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index a0b240d2380..2ec77518205 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -12,7 +12,7 @@ use crate::dataset::ROW_ID; use crate::dataset::WriteDestination; use crate::dataset::builder::DatasetBuilder; use crate::dataset::index::LanceIndexStoreExt; -use crate::dataset::optimize::{CompactionOptions, compact_files}; +use crate::dataset::optimize::{CompactionOptions, compact_files, remapping}; use crate::dataset::tests::dataset_migrations::scan_dataset; use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; @@ -2833,10 +2833,18 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { } #[tokio::test] -async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { +async fn test_partial_compound_hybrid_rejects_same_id_rewrite_after_deferred_remap() { let initial = arrow_array::record_batch!( - ("text", Utf8, ["stable alpha common", "stale alpha common"]), - ("id", Int32, [0, 1]) + ( + "text", + Utf8, + [ + "stable alpha common", + "stable gamma common", + "stale alpha common" + ] + ), + ("id", Int32, [0, 1, 2]) ) .unwrap(); let schema = initial.schema(); @@ -2862,12 +2870,12 @@ async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { .await .unwrap(); - let rewritten_fragment = dataset.get_fragment(1).unwrap(); + let rewritten_fragment = dataset.get_fragment(2).unwrap(); let mut replacement_file = rewritten_fragment.metadata().files[0].clone(); replacement_file.path = "replacement.lance".to_string(); let replacement = arrow_array::record_batch!( ("text", Utf8, ["replacement beta common"]), - ("id", Int32, [1]) + ("id", Int32, [2]) ) .unwrap(); let object_writer = dataset @@ -2888,7 +2896,7 @@ async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { let mut dataset = Dataset::commit( WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { - replacements: vec![DataReplacementGroup(1, replacement_file)], + replacements: vec![DataReplacementGroup(2, replacement_file)], }, Some(read_version), None, @@ -2904,12 +2912,45 @@ async fn test_partial_compound_hybrid_rejects_same_id_indexed_field_rewrite() { .unwrap() .unwrap(); assert!( - !committed.fragment_bitmap.as_ref().unwrap().contains(1), + !committed.fragment_bitmap.as_ref().unwrap().contains(2), "the logical index must prune the same-id rewritten fragment" ); + let physical_source_version = committed.dataset_version; + + // Compact only the two still-covered fragments and defer the address + // remap. The stale physical document from rewritten fragment 2 is absent + // from that remap and therefore remains in the postings. + let metrics = compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 2, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + assert_eq!(metrics.fragments_removed, 2); + assert_eq!(metrics.fragments_added, 1); + + remapping::remap_column_index(&mut dataset, &["text"], Some("text_idx".to_string())) + .await + .unwrap(); + let remapped = dataset + .load_index_by_name("text_idx") + .await + .unwrap() + .unwrap(); + assert!(remapped.dataset_version > physical_source_version); + assert_eq!( + crate::index::scalar::inverted::physical_source_dataset_versions(&remapped).unwrap(), + Some(vec![physical_source_version]), + "deferred remap must not advance immutable physical provenance" + ); let appended = - arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [2])) + arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [3])) .unwrap(); dataset .append( @@ -3057,7 +3098,8 @@ async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() ) .await .unwrap(); - let rewritten_file = &dataset.get_fragment(1).unwrap().metadata().files[0]; + let rewritten_fragment = dataset.get_fragment(1).unwrap(); + let rewritten_file = &rewritten_fragment.metadata().files[0]; assert_eq!(rewritten_file.path, relative_path); assert_eq!(rewritten_file.base_id, Some(2)); diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 18bcf509569..651f554d5b7 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -234,16 +234,16 @@ fn segment_indexed_field_ids(dataset: &Dataset, segment: &IndexSegment) -> Resul indexed_field_ids(dataset, segment.fields()) } -/// Prove that every indexed-field file present at a segment's build version is -/// still the current file for the same fragment id. +/// Prove that every indexed-field file present at a segment's physical source +/// version is still the current file for the same fragment id. /// /// A committed segment's logical bitmap may already have pruned a rewritten /// fragment even though its physical postings and corpus statistics still /// contain that fragment. Since the original per-segment coverage is no longer /// available after that pruning, this check deliberately considers every -/// fragment present at each build version. That is conservative, but it makes -/// the hybrid scorer available only when the history can be proven to be a -/// pure append with respect to the indexed field subtrees. +/// fragment present at each physical source version. That is conservative, but +/// it makes the hybrid scorer available only when the history can be proven to +/// be a pure append with respect to the indexed field subtrees. pub(crate) async fn has_append_only_indexed_field_history( dataset: &Dataset, segments: &[IndexMetadata], @@ -258,9 +258,27 @@ pub(crate) async fn has_append_only_indexed_field_history( return false; } - let build_versions = segments + let mut physical_sources = Vec::with_capacity(segments.len()); + for segment in segments { + let Ok(Some(source_versions)) = scalar::inverted::physical_source_dataset_versions(segment) + else { + return false; + }; + // `dataset_version` is an independent mutable address-remap watermark. + // A merged segment can legitimately contain physical sources newer + // than its oldest address watermark, but neither may be in the future. + if source_versions + .iter() + .any(|source_version| *source_version > current_version) + || segment.dataset_version > current_version + { + return false; + } + physical_sources.push((segment, source_versions)); + } + let build_versions = physical_sources .iter() - .map(|segment| segment.dataset_version) + .flat_map(|(_, source_versions)| source_versions.iter().copied()) .collect::>(); for build_version in build_versions { if build_version > current_version { @@ -273,9 +291,9 @@ pub(crate) async fn has_append_only_indexed_field_history( return false; }; let mut indexed_field_ids_at_version = HashSet::new(); - for segment in segments + for (segment, _) in physical_sources .iter() - .filter(|segment| segment.dataset_version == build_version) + .filter(|(_, source_versions)| source_versions.contains(&build_version)) { let Ok(historical_field_ids) = indexed_field_ids(&historical, &segment.fields) else { return false; @@ -1320,7 +1338,8 @@ pub(crate) async fn remap_index( .as_any() .downcast_ref::() .ok_or(Error::index("expected inverted index".to_string()))?; - if inverted_index.is_legacy() { + let is_legacy = inverted_index.is_legacy(); + let mut created_index = if is_legacy { log::warn!( "reindex because of legacy format, index_type: {}, index_id: {}, field: {}", scalar_index.index_type(), @@ -1346,7 +1365,21 @@ pub(crate) async fn remap_index( .await? } else { scalar_index.remap(row_id_map, &new_store).await? + }; + let source_versions = if is_legacy { + // Legacy remapping performs a full rebuild from the + // current dataset instead of retaining old postings. + Some(vec![dataset.manifest.version]) + } else { + scalar::inverted::physical_source_dataset_versions(matched)? + }; + if let Some(source_versions) = source_versions { + scalar::inverted::set_physical_source_dataset_versions( + &mut created_index.index_details, + source_versions, + )?; } + created_index } _ => scalar_index.remap(row_id_map, &new_store).await?, } diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 34f52c58718..ac9c7d3756b 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -1203,7 +1203,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( ) .await?; let new_uuid = Uuid::new_v4(); - let created_index = super::scalar::build_scalar_index( + let mut created_index = super::scalar::build_scalar_index( dataset.as_ref(), &resolved.canonical_path, new_uuid, @@ -1214,6 +1214,10 @@ pub async fn merge_indices_with_unindexed_frags<'a>( Arc::new(NoopIndexBuildProgress), ) .await?; + crate::index::scalar::inverted::set_physical_source_dataset_versions( + &mut created_index.index_details, + [dataset.manifest.version], + )?; return Ok(Some(IndexMergeResults { new_uuid, removed_indices: old_indices.to_vec(), @@ -1283,38 +1287,61 @@ pub async fn merge_indices_with_unindexed_frags<'a>( let new_uuid = Uuid::new_v4(); let new_store = LanceIndexStore::from_dataset_for_new(&dataset, &new_uuid)?; - let (created_index, new_dataset_version) = if selected_indices.is_empty() { - ( - super::scalar::build_scalar_index( - dataset.as_ref(), - &resolved.canonical_path, - new_uuid, - &reference_index.derive_index_params()?, - true, - None, - Some(new_data_stream), - Arc::new(NoopIndexBuildProgress), - ) - .await?, - dataset.manifest.version, + let new_dataset_version = if selected_indices.is_empty() { + dataset.manifest.version + } else { + selected_old_indices + .iter() + .map(|index| index.dataset_version) + .min() + .unwrap_or(dataset.manifest.version) + }; + let source_versions = if selected_indices.is_empty() { + Some(vec![dataset.manifest.version]) + } else { + let mut source_sets = selected_old_indices + .iter() + .map(|segment| { + crate::index::scalar::inverted::physical_source_dataset_versions( + segment, + ) + }) + .collect::>>()?; + if !unindexed.is_empty() { + source_sets.push(Some(vec![dataset.manifest.version])); + } + crate::index::scalar::inverted::merge_physical_source_dataset_versions( + source_sets, ) + }; + let mut created_index = if selected_indices.is_empty() { + super::scalar::build_scalar_index( + dataset.as_ref(), + &resolved.canonical_path, + new_uuid, + &reference_index.derive_index_params()?, + true, + None, + Some(new_data_stream), + Arc::new(NoopIndexBuildProgress), + ) + .await? } else { - ( - InvertedIndex::merge_segments( - &selected_indices, - new_data_stream, - &new_store, - old_data_filter, - options.progress.clone(), - ) - .await?, - selected_old_indices - .iter() - .map(|index| index.dataset_version) - .min() - .unwrap_or(dataset.manifest.version), + InvertedIndex::merge_segments( + &selected_indices, + new_data_stream, + &new_store, + old_data_filter, + options.progress.clone(), ) + .await? }; + if let Some(source_versions) = source_versions { + crate::index::scalar::inverted::set_physical_source_dataset_versions( + &mut created_index.index_details, + source_versions, + )?; + } Ok(( new_uuid, diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 8168cd1eb2d..35f523b0851 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -330,7 +330,7 @@ impl<'a> CreateIndexBuilder<'a> { let index_id = self.index_uuid.unwrap_or_else(Uuid::new_v4); let mut output_index_uuid = index_id; - let created_index = match (self.index_type, self.params.index_name()) { + let mut created_index = match (self.index_type, self.params.index_name()) { ( IndexType::Bitmap | IndexType::BTree @@ -596,6 +596,16 @@ impl<'a> CreateIndexBuilder<'a> { ))); } }; + if created_index + .index_details + .type_url + .ends_with("InvertedIndexDetails") + { + crate::index::scalar::inverted::set_physical_source_dataset_versions( + &mut created_index.index_details, + [self.dataset.manifest.version], + )?; + } Ok(IndexMetadata { uuid: output_index_uuid, diff --git a/rust/lance/src/index/scalar/inverted.rs b/rust/lance/src/index/scalar/inverted.rs index c7c521050b4..1370b9d2827 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -3,7 +3,10 @@ #![allow(clippy::redundant_pub_crate)] -use std::{collections::BTreeMap, sync::Arc}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, +}; use arrow_array::cast::AsArray; use arrow_array::{ @@ -807,6 +810,11 @@ pub(crate) async fn merge_segments( let document_granularity = DocumentGranularity::try_from(details.document_granularity)?; let resolved = resolve_fts_field_by_id(dataset.schema(), field_id, document_granularity)?; load_segment_details(dataset, &resolved.canonical_path, &segments).await?; + let source_sets = segments + .iter() + .map(physical_source_dataset_versions) + .collect::>>()?; + let merged_source_versions = merge_physical_source_dataset_versions(source_sets); let mut source_indices = Vec::with_capacity(segments.len()); let mut fragment_bitmap = RoaringBitmap::new(); @@ -846,7 +854,7 @@ pub(crate) async fn merge_segments( let new_uuid = Uuid::new_v4(); let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_uuid)?; - let created_index = InvertedIndex::merge_segments( + let mut created_index = InvertedIndex::merge_segments( &source_indices, empty_inverted_update_stream(dataset, &resolved)?, &new_store, @@ -854,6 +862,9 @@ pub(crate) async fn merge_segments( lance_index::progress::noop_progress(), ) .await?; + if let Some(source_versions) = merged_source_versions { + set_physical_source_dataset_versions(&mut created_index.index_details, source_versions)?; + } Ok(IndexMetadata { uuid: new_uuid, @@ -936,6 +947,74 @@ pub(crate) async fn fts_index_fragment_bitmap( Ok(fragment_bitmap) } +/// Return the explicit dataset versions whose physical documents contribute to +/// an inverted segment. +/// +/// `None` means the segment predates explicit physical provenance. It must not +/// fall back to `IndexMetadata::dataset_version`: older Lance versions may have +/// advanced that mutable remap watermark while retaining stale postings. +pub(crate) fn physical_source_dataset_versions( + segment: &IndexMetadata, +) -> Result>> { + let Some(details_any) = segment.index_details.as_ref() else { + return Ok(None); + }; + let details = InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| { + Error::io(format!( + "failed to decode InvertedIndexDetails physical provenance: {error}" + )) + })?; + if details.physical_source_dataset_versions.is_empty() { + return Ok(None); + } + let mut versions = details.physical_source_dataset_versions; + versions.sort_unstable(); + versions.dedup(); + Ok(Some(versions)) +} + +/// Union known physical provenance sets. If any input is unknown, the merged +/// corpus is also unknown because its retained postings cannot be attributed +/// exactly. +pub(crate) fn merge_physical_source_dataset_versions( + source_sets: impl IntoIterator>>, +) -> Option> { + let mut merged = BTreeSet::new(); + let mut has_input = false; + for source_versions in source_sets { + has_input = true; + let source_versions = source_versions?; + if source_versions.is_empty() { + return None; + } + merged.extend(source_versions); + } + (has_input && !merged.is_empty()).then(|| merged.into_iter().collect()) +} + +/// Stamp immutable physical provenance into an inverted-details payload while +/// preserving its existing type URL. +pub(crate) fn set_physical_source_dataset_versions( + details_any: &mut prost_types::Any, + dataset_versions: impl IntoIterator, +) -> Result<()> { + let mut details = + InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| { + Error::io(format!( + "failed to decode InvertedIndexDetails physical provenance: {error}" + )) + })?; + let dataset_versions = dataset_versions.into_iter().collect::>(); + if dataset_versions.is_empty() { + return Err(Error::invalid_input( + "physical source dataset versions must not be empty".to_string(), + )); + } + details.physical_source_dataset_versions = dataset_versions.into_iter().collect(); + details_any.value = details.encode_to_vec(); + Ok(()) +} + /// Load and validate the shared [`InvertedIndexDetails`] across committed /// segments returned by [`load_segments`]. /// @@ -984,14 +1063,18 @@ pub async fn load_segment_details( fn canonicalize_inverted_index_details( details: InvertedIndexDetails, ) -> Result { + let physical_source_dataset_versions = details.physical_source_dataset_versions.clone(); let params = InvertedIndexParams::try_from(&details)?; - InvertedIndexDetails::try_from(¶ms) + let mut canonical = InvertedIndexDetails::try_from(¶ms)?; + canonical.physical_source_dataset_versions = physical_source_dataset_versions; + Ok(canonical) } /// Compare canonicalized inverted-index details for shared semantic configuration. /// /// `posting_format_version` records how a single segment physically stores -/// postings, so mixed-version FTS segments may disagree on it without being +/// postings, and `physical_source_dataset_versions` records its provenance, so +/// mixed FTS segments may disagree on either without being semantically /// incompatible. Every other field remains part of the equality check. fn inverted_index_details_semantically_equal( left: &InvertedIndexDetails, @@ -1001,6 +1084,8 @@ fn inverted_index_details_semantically_equal( let mut right = right.clone(); left.posting_format_version = None; right.posting_format_version = None; + left.physical_source_dataset_versions.clear(); + right.physical_source_dataset_versions.clear(); left == right } @@ -1183,6 +1268,52 @@ mod tests { ); } + #[test] + fn canonicalize_inverted_details_preserves_physical_provenance() { + let mut details = InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(); + details.physical_source_dataset_versions = vec![17, 23]; + + let canonical = canonicalize_inverted_index_details(details).unwrap(); + + assert_eq!(canonical.physical_source_dataset_versions, vec![17, 23]); + let mut other_source = canonical.clone(); + other_source.physical_source_dataset_versions = vec![18]; + assert!(inverted_index_details_semantically_equal( + &canonical, + &other_source + )); + } + + #[test] + fn canonicalize_inverted_details_preserves_unknown_physical_provenance() { + let unknown = canonicalize_inverted_index_details(InvertedIndexDetails::default()).unwrap(); + assert!(unknown.physical_source_dataset_versions.is_empty()); + + let mut known = unknown.clone(); + known.physical_source_dataset_versions = vec![17]; + assert!(inverted_index_details_semantically_equal(&unknown, &known)); + } + + #[test] + fn merge_physical_provenance_unions_known_sources_and_preserves_unknown() { + assert_eq!( + merge_physical_source_dataset_versions([Some(vec![19, 17, 19]), Some(vec![18, 19])]), + Some(vec![17, 18, 19]) + ); + assert_eq!( + merge_physical_source_dataset_versions([Some(vec![17]), None]), + None + ); + assert_eq!( + merge_physical_source_dataset_versions([Some(Vec::new())]), + None + ); + assert_eq!( + merge_physical_source_dataset_versions(std::iter::empty::>>()), + None + ); + } + #[test] fn inverted_details_equal_when_only_posting_format_version_differs() { let left = canonicalize_inverted_index_details( From f1bda89a77a8b98ba6f886c8ecb2fd869e9f6735 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 18:01:48 +0800 Subject: [PATCH 13/22] fix(fts): store physical provenance in index sidecar --- protos/index_old.proto | 7 - .../src/scalar/inverted/tokenizer.rs | 2 - rust/lance/src/dataset/tests/dataset_index.rs | 4 +- rust/lance/src/index.rs | 16 +- rust/lance/src/index/append.rs | 35 ++-- rust/lance/src/index/create.rs | 14 +- rust/lance/src/index/scalar/inverted.rs | 180 +++++++++++------- 7 files changed, 153 insertions(+), 105 deletions(-) diff --git a/protos/index_old.proto b/protos/index_old.proto index 5722b2525f1..236d1f110c0 100644 --- a/protos/index_old.proto +++ b/protos/index_old.proto @@ -101,11 +101,4 @@ message InvertedIndexDetails { // The posting-list payload format. This is separate from index_version, // which identifies the overall inverted-index layout. optional uint32 posting_format_version = 15; - // Dataset versions whose physical documents contribute to this segment. - // Unlike IndexMetadata.dataset_version, this provenance is not advanced by - // row-address remapping because remapping can retain documents outside its - // address map. Writers sort and deduplicate this set. An empty set means - // provenance is unknown; readers must not infer it from - // IndexMetadata.dataset_version. - repeated uint64 physical_source_dataset_versions = 16; } diff --git a/rust/lance-index/src/scalar/inverted/tokenizer.rs b/rust/lance-index/src/scalar/inverted/tokenizer.rs index 818ce466511..fd9b9294d7b 100644 --- a/rust/lance-index/src/scalar/inverted/tokenizer.rs +++ b/rust/lance-index/src/scalar/inverted/tokenizer.rs @@ -427,7 +427,6 @@ impl TryFrom<&InvertedIndexParams> for pbold::InvertedIndexDetails { ), document_granularity: PbDocumentGranularity::from(params.document_granularity) as i32, posting_format_version: Some(params.resolved_format_version().index_version()), - physical_source_dataset_versions: Vec::new(), }) } } @@ -1560,7 +1559,6 @@ mod tests { code_config: None, document_granularity: PbDocumentGranularity::Row as i32, posting_format_version: None, - physical_source_dataset_versions: Vec::new(), }; let params = InvertedIndexParams::try_from(&old_details).unwrap(); assert_eq!(params.block_size, 128); diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 2ec77518205..4617317ec33 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2944,7 +2944,9 @@ async fn test_partial_compound_hybrid_rejects_same_id_rewrite_after_deferred_rem .unwrap(); assert!(remapped.dataset_version > physical_source_version); assert_eq!( - crate::index::scalar::inverted::physical_source_dataset_versions(&remapped).unwrap(), + crate::index::scalar::inverted::physical_source_dataset_versions(&dataset, &remapped) + .await + .unwrap(), Some(vec![physical_source_version]), "deferred remap must not advance immutable physical provenance" ); diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 651f554d5b7..ab8835d5bdb 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -260,7 +260,8 @@ pub(crate) async fn has_append_only_indexed_field_history( let mut physical_sources = Vec::with_capacity(segments.len()); for segment in segments { - let Ok(Some(source_versions)) = scalar::inverted::physical_source_dataset_versions(segment) + let Ok(Some(source_versions)) = + scalar::inverted::physical_source_dataset_versions(dataset, segment).await else { return false; }; @@ -1371,13 +1372,16 @@ pub(crate) async fn remap_index( // current dataset instead of retaining old postings. Some(vec![dataset.manifest.version]) } else { - scalar::inverted::physical_source_dataset_versions(matched)? + scalar::inverted::physical_source_dataset_versions(dataset, matched).await? }; if let Some(source_versions) = source_versions { - scalar::inverted::set_physical_source_dataset_versions( - &mut created_index.index_details, - source_versions, - )?; + created_index.files.push( + scalar::inverted::write_physical_source_dataset_versions( + &new_store, + source_versions, + ) + .await?, + ); } created_index } diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index ac9c7d3756b..148e7a06dc5 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -1214,10 +1214,14 @@ pub async fn merge_indices_with_unindexed_frags<'a>( Arc::new(NoopIndexBuildProgress), ) .await?; - crate::index::scalar::inverted::set_physical_source_dataset_versions( - &mut created_index.index_details, - [dataset.manifest.version], - )?; + let new_store = LanceIndexStore::from_dataset_for_new(&dataset, &new_uuid)?; + created_index.files.push( + crate::index::scalar::inverted::write_physical_source_dataset_versions( + &new_store, + [dataset.manifest.version], + ) + .await?, + ); return Ok(Some(IndexMergeResults { new_uuid, removed_indices: old_indices.to_vec(), @@ -1299,14 +1303,16 @@ pub async fn merge_indices_with_unindexed_frags<'a>( let source_versions = if selected_indices.is_empty() { Some(vec![dataset.manifest.version]) } else { - let mut source_sets = selected_old_indices - .iter() - .map(|segment| { + let mut source_sets = Vec::with_capacity(selected_old_indices.len() + 1); + for segment in &selected_old_indices { + source_sets.push( crate::index::scalar::inverted::physical_source_dataset_versions( + dataset.as_ref(), segment, ) - }) - .collect::>>()?; + .await?, + ); + } if !unindexed.is_empty() { source_sets.push(Some(vec![dataset.manifest.version])); } @@ -1337,10 +1343,13 @@ pub async fn merge_indices_with_unindexed_frags<'a>( .await? }; if let Some(source_versions) = source_versions { - crate::index::scalar::inverted::set_physical_source_dataset_versions( - &mut created_index.index_details, - source_versions, - )?; + created_index.files.push( + crate::index::scalar::inverted::write_physical_source_dataset_versions( + &new_store, + source_versions, + ) + .await?, + ); } Ok(( diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 35f523b0851..f52da5ee441 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -5,6 +5,7 @@ use crate::{ Error, Result, dataset::{ Dataset, + index::LanceIndexStoreExt, transaction::{Operation, TransactionBuilder}, }, index::{ @@ -601,10 +602,17 @@ impl<'a> CreateIndexBuilder<'a> { .type_url .ends_with("InvertedIndexDetails") { - crate::index::scalar::inverted::set_physical_source_dataset_versions( - &mut created_index.index_details, - [self.dataset.manifest.version], + let store = lance_index::scalar::lance_format::LanceIndexStore::from_dataset_for_new( + self.dataset, + &output_index_uuid, )?; + created_index.files.push( + crate::index::scalar::inverted::write_physical_source_dataset_versions( + &store, + [self.dataset.manifest.version], + ) + .await?, + ); } Ok(IndexMetadata { diff --git a/rust/lance/src/index/scalar/inverted.rs b/rust/lance/src/index/scalar/inverted.rs index 1370b9d2827..1483ca31ac7 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -24,12 +24,12 @@ use lance_core::{ }; use lance_index::metrics::NoOpMetricsCollector; use lance_index::pbold::InvertedIndexDetails; -use lance_index::scalar::index_files_to_table; use lance_index::scalar::inverted::{ DocumentGranularity, InvertedIndex, InvertedIndexParams, doc_index_storage_column, }; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::VALUE_COLUMN_NAME; +use lance_index::scalar::{IndexFile, IndexStore, index_files_to_table}; use lance_table::format::IndexMetadata; use prost::Message; use roaring::RoaringBitmap; @@ -41,6 +41,9 @@ use crate::{ index::{DatasetIndexExt, scalar::fetch_index_details}, }; +const PHYSICAL_SOURCE_VERSIONS_FILE: &str = "physical_source_versions.lance"; +const PHYSICAL_SOURCE_VERSION_COLUMN: &str = "dataset_version"; + #[derive(Debug, Clone)] enum FtsTraversal { Text, @@ -810,10 +813,10 @@ pub(crate) async fn merge_segments( let document_granularity = DocumentGranularity::try_from(details.document_granularity)?; let resolved = resolve_fts_field_by_id(dataset.schema(), field_id, document_granularity)?; load_segment_details(dataset, &resolved.canonical_path, &segments).await?; - let source_sets = segments - .iter() - .map(physical_source_dataset_versions) - .collect::>>()?; + let mut source_sets = Vec::with_capacity(segments.len()); + for segment in &segments { + source_sets.push(physical_source_dataset_versions(dataset, segment).await?); + } let merged_source_versions = merge_physical_source_dataset_versions(source_sets); let mut source_indices = Vec::with_capacity(segments.len()); @@ -863,7 +866,9 @@ pub(crate) async fn merge_segments( ) .await?; if let Some(source_versions) = merged_source_versions { - set_physical_source_dataset_versions(&mut created_index.index_details, source_versions)?; + created_index + .files + .push(write_physical_source_dataset_versions(&new_store, source_versions).await?); } Ok(IndexMetadata { @@ -947,27 +952,83 @@ pub(crate) async fn fts_index_fragment_bitmap( Ok(fragment_bitmap) } -/// Return the explicit dataset versions whose physical documents contribute to -/// an inverted segment. +/// Read the dataset versions whose physical documents contribute to an +/// inverted segment. /// -/// `None` means the segment predates explicit physical provenance. It must not -/// fall back to `IndexMetadata::dataset_version`: older Lance versions may have -/// advanced that mutable remap watermark while retaining stale postings. -pub(crate) fn physical_source_dataset_versions( +/// Missing sidecars represent legacy segments with unknown provenance. A +/// malformed sidecar returns an error so corruption is never mistaken for +/// missing provenance. +pub(crate) async fn physical_source_dataset_versions( + dataset: &Dataset, segment: &IndexMetadata, ) -> Result>> { - let Some(details_any) = segment.index_details.as_ref() else { - return Ok(None); - }; - let details = InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| { - Error::io(format!( - "failed to decode InvertedIndexDetails physical provenance: {error}" - )) - })?; - if details.physical_source_dataset_versions.is_empty() { - return Ok(None); + match &segment.files { + Some(files) => { + if !files + .iter() + .any(|file| file.path == PHYSICAL_SOURCE_VERSIONS_FILE) + { + return Ok(None); + } + } + None => { + let sidecar_path = dataset + .indice_files_dir(segment)? + .join(segment.uuid.to_string()) + .join(PHYSICAL_SOURCE_VERSIONS_FILE); + let object_store = dataset.object_store_for_index(segment).await?; + if !object_store.exists(&sidecar_path).await? { + return Ok(None); + } + } + } + + let store = LanceIndexStore::from_dataset_for_existing(dataset, segment).await?; + let reader = store.open_index_file(PHYSICAL_SOURCE_VERSIONS_FILE).await?; + let num_rows = reader.num_rows(); + if num_rows == 0 { + return Err(Error::io(format!( + "physical provenance sidecar for index {} is empty", + segment.uuid + ))); + } + let batch = reader.read_range(0..num_rows, None).await?; + let schema = batch.schema(); + if schema.fields().len() != 1 + || schema.field(0).name() != PHYSICAL_SOURCE_VERSION_COLUMN + || schema.field(0).data_type() != &DataType::UInt64 + || schema.field(0).is_nullable() + { + return Err(Error::io(format!( + "physical provenance sidecar for index {} has invalid schema {:?}", + segment.uuid, schema + ))); } - let mut versions = details.physical_source_dataset_versions; + if batch.num_rows() != num_rows { + return Err(Error::io(format!( + "physical provenance sidecar for index {} declared {} rows but read {}", + segment.uuid, + num_rows, + batch.num_rows() + ))); + } + let versions = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::io(format!( + "physical provenance sidecar for index {} has a non-UInt64 column", + segment.uuid + )) + })?; + if versions.null_count() != 0 { + return Err(Error::io(format!( + "physical provenance sidecar for index {} contains null versions", + segment.uuid + ))); + } + let mut versions = versions.values().to_vec(); versions.sort_unstable(); versions.dedup(); Ok(Some(versions)) @@ -992,27 +1053,31 @@ pub(crate) fn merge_physical_source_dataset_versions( (has_input && !merged.is_empty()).then(|| merged.into_iter().collect()) } -/// Stamp immutable physical provenance into an inverted-details payload while -/// preserving its existing type URL. -pub(crate) fn set_physical_source_dataset_versions( - details_any: &mut prost_types::Any, +/// Write immutable physical provenance as an index-local Lance sidecar. +pub(crate) async fn write_physical_source_dataset_versions( + store: &dyn IndexStore, dataset_versions: impl IntoIterator, -) -> Result<()> { - let mut details = - InvertedIndexDetails::decode(details_any.value.as_slice()).map_err(|error| { - Error::io(format!( - "failed to decode InvertedIndexDetails physical provenance: {error}" - )) - })?; +) -> Result { let dataset_versions = dataset_versions.into_iter().collect::>(); if dataset_versions.is_empty() { return Err(Error::invalid_input( "physical source dataset versions must not be empty".to_string(), )); } - details.physical_source_dataset_versions = dataset_versions.into_iter().collect(); - details_any.value = details.encode_to_vec(); - Ok(()) + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + PHYSICAL_SOURCE_VERSION_COLUMN, + DataType::UInt64, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(UInt64Array::from_iter_values(dataset_versions))], + )?; + let mut writer = store + .new_index_file(PHYSICAL_SOURCE_VERSIONS_FILE, schema) + .await?; + writer.write_record_batch(batch).await?; + writer.finish().await } /// Load and validate the shared [`InvertedIndexDetails`] across committed @@ -1063,19 +1128,16 @@ pub async fn load_segment_details( fn canonicalize_inverted_index_details( details: InvertedIndexDetails, ) -> Result { - let physical_source_dataset_versions = details.physical_source_dataset_versions.clone(); let params = InvertedIndexParams::try_from(&details)?; - let mut canonical = InvertedIndexDetails::try_from(¶ms)?; - canonical.physical_source_dataset_versions = physical_source_dataset_versions; - Ok(canonical) + InvertedIndexDetails::try_from(¶ms) } /// Compare canonicalized inverted-index details for shared semantic configuration. /// /// `posting_format_version` records how a single segment physically stores -/// postings, and `physical_source_dataset_versions` records its provenance, so -/// mixed FTS segments may disagree on either without being semantically -/// incompatible. Every other field remains part of the equality check. +/// postings, so mixed-version FTS segments may disagree on it without being +/// semantically incompatible. Every other field remains part of the equality +/// check. fn inverted_index_details_semantically_equal( left: &InvertedIndexDetails, right: &InvertedIndexDetails, @@ -1084,8 +1146,6 @@ fn inverted_index_details_semantically_equal( let mut right = right.clone(); left.posting_format_version = None; right.posting_format_version = None; - left.physical_source_dataset_versions.clear(); - right.physical_source_dataset_versions.clear(); left == right } @@ -1269,33 +1329,7 @@ mod tests { } #[test] - fn canonicalize_inverted_details_preserves_physical_provenance() { - let mut details = InvertedIndexDetails::try_from(&InvertedIndexParams::default()).unwrap(); - details.physical_source_dataset_versions = vec![17, 23]; - - let canonical = canonicalize_inverted_index_details(details).unwrap(); - - assert_eq!(canonical.physical_source_dataset_versions, vec![17, 23]); - let mut other_source = canonical.clone(); - other_source.physical_source_dataset_versions = vec![18]; - assert!(inverted_index_details_semantically_equal( - &canonical, - &other_source - )); - } - - #[test] - fn canonicalize_inverted_details_preserves_unknown_physical_provenance() { - let unknown = canonicalize_inverted_index_details(InvertedIndexDetails::default()).unwrap(); - assert!(unknown.physical_source_dataset_versions.is_empty()); - - let mut known = unknown.clone(); - known.physical_source_dataset_versions = vec![17]; - assert!(inverted_index_details_semantically_equal(&unknown, &known)); - } - - #[test] - fn merge_physical_provenance_unions_known_sources_and_preserves_unknown() { + fn merge_physical_provenance_sidecar_sources() { assert_eq!( merge_physical_source_dataset_versions([Some(vec![19, 17, 19]), Some(vec![18, 19])]), Some(vec![17, 18, 19]) From 913e19206737fdca86b82e040d16e140b5ca3857 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 18:12:31 +0800 Subject: [PATCH 14/22] test(cleanup): account for FTS provenance sidecars --- rust/lance/src/dataset/cleanup.rs | 184 ++++++++++++++++-------------- 1 file changed, 96 insertions(+), 88 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index ad35669f37b..b820eca849c 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -3673,6 +3673,14 @@ mod tests { // Note: branch2 is stored as "dev/branch2"; comments may refer to it as branch2 for brevity. // Important: auto_cleanup_hook uses policy derived from manifest config; it does not flip // clean_referenced_branches unless tests call cleanup_old_versions with a custom policy. + // Cleanup retains or removes an index UUID directory as a unit. Each + // inverted-index generation has seven core files plus its provenance sidecar. + const INDEX_FILES_PER_GENERATION: usize = 8; + + fn index_file_count(generations: usize) -> usize { + generations * INDEX_FILES_PER_GENERATION + } + struct LineageSetup { main: BranchDatasetFixture, branch1: BranchDatasetFixture, @@ -4092,13 +4100,13 @@ mod tests { // - 1 manifest file // - 1 data file // - 1 deletion file - // - 4 index files + // - one index generation // The left is the counts for the latest version of appending assert_eq!(setup.branch1.counts.num_manifest_files, 2); assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 2); - assert_eq!(setup.branch1.counts.num_index_files, 14); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(2)); setup.assert_all_unchanged().await; setup.branch1.compact().await.unwrap(); @@ -4107,13 +4115,13 @@ mod tests { // - 1 manifest file // - 1 data file // - 1 deletion file - // - 4 index files - // The left (1, 1, 1, 0, 4) is the counts for the latest version of compaction + // - one index generation + // Counts include one retained index generation. assert_eq!(setup.branch1.counts.num_manifest_files, 2); assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, 14); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(2)); setup.assert_all_unchanged().await; // Now we clean the referenced files of branch1 by branch2 and branch3 @@ -4122,28 +4130,28 @@ mod tests { setup.branch3.run_cleanup().await.unwrap(); setup.branch2.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version of compaction + // Counts include one retained index generation. assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version of compaction + // Counts include one retained index generation. assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup.branch1.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version of compaction + // Counts include one retained index generation. assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); setup.assert_all_unchanged().await; } @@ -4160,7 +4168,7 @@ mod tests { assert_eq!(setup.branch3.counts.num_data_files, 2); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 2); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup .assert_unchanged(&["branch1", "branch2", "branch4", "main"]) .await; @@ -4176,17 +4184,17 @@ mod tests { assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); setup.branch3.compact().await.unwrap(); setup.branch3.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version + // Counts include one retained index generation. assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup .assert_unchanged(&["branch1", "branch2", "branch4", "main"]) .await; @@ -4194,12 +4202,12 @@ mod tests { setup.branch2.compact().await.unwrap(); setup.branch2.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version + // Counts include one retained index generation. assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); } #[tokio::test] @@ -4216,7 +4224,7 @@ mod tests { assert_eq!(setup.branch4.counts.num_data_files, 2); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 2); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); setup.assert_all_unchanged().await; setup.main.compact().await.unwrap(); @@ -4225,28 +4233,28 @@ mod tests { // - 1 manifest file // - 2 data files // - 1 deletion file - // - 4 index files + // - one index generation // Branch4 holds references from main: // - 1 manifest file // - 3 data files // - 1 deletion file - // - 4 index files + // - one index generation // The left(1, 1, 1, 0, 0) is the counts for the latest version of compaction assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 14); + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); setup.branch4.compact().await.unwrap(); setup.branch4.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts of one version + // Counts include one retained index generation. assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); setup.assert_all_unchanged().await; setup.main.run_cleanup().await.unwrap(); @@ -4254,13 +4262,13 @@ mod tests { // - 1 manifest file // - 2 data files // - 1 deletion file - // - 4 index files - // The left(1, 1, 1, 0, 4) is the counts for the latest version of compaction + // - one index generation + // Counts include one retained index generation. assert_eq!(setup.main.counts.num_manifest_files, 2); assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, 14); + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); } #[tokio::test] @@ -4274,18 +4282,18 @@ mod tests { // - 1 manifest file // - 2 data files // - 1 deletion file - // - 4 index files(only for branch1) + // - one index generation(only for branch1) // Branch4 holds references from main: // - 1 manifest file // - 3 data files // - 1 deletion file - // - 4 index files - // The left(1, 1, 1, 1, 4) is the counts for the latest version of compaction + // - one index generation + // Counts include one retained index generation. assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 3); - assert_eq!(setup.main.counts.num_index_files, 21); + assert_eq!(setup.main.counts.num_index_files, index_file_count(3)); setup.assert_all_unchanged().await; setup.main.compact().await.unwrap(); @@ -4296,7 +4304,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 21); + assert_eq!(setup.main.counts.num_index_files, index_file_count(3)); setup.assert_all_unchanged().await; setup.branch1.write_data().await.unwrap(); @@ -4312,19 +4320,19 @@ mod tests { // - 1 manifest file // - 1 data files // - 2 deletion files - // - 4 index files + // - one index generation assert_eq!(setup.branch2.counts.num_manifest_files, 2); assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, 14); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(2)); setup.branch1.run_cleanup().await.unwrap(); - // Cleanup 4 index files referenced from branch2 + // Cleanup one index generation referenced from branch2 assert_eq!(setup.branch1.counts.num_manifest_files, 2); assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); setup.main.run_cleanup().await.unwrap(); // Branch3 holds references from main: @@ -4335,12 +4343,12 @@ mod tests { // - 1 manifest file // - 3 data files // - 2 deletion files - // - 4 index files + // - one index generation assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 14); + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); setup.branch3.write_data().await.unwrap(); setup.branch3.compact().await.unwrap(); @@ -4350,7 +4358,7 @@ mod tests { assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup.main.run_cleanup().await.unwrap(); // Cleanup doesn't take effects if we don't clean branch2 and branch1 first @@ -4358,7 +4366,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 14); + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); // Cleanup doesn't take effect if we don't clean branch2 first setup.branch1.run_cleanup().await.unwrap(); @@ -4366,57 +4374,57 @@ mod tests { assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); setup.branch2.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version + // Counts include one retained index generation. assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); setup.branch1.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version + // Counts include one retained index generation. assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); setup.main.run_cleanup().await.unwrap(); // Branch4 holds references from main: // - 1 manifest file // - 3 data files // - 2 deletion files - // - 4 index files + // - one index generation assert_eq!(setup.main.counts.num_manifest_files, 2); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 14); + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); setup.branch4.write_data().await.unwrap(); setup.branch4.compact().await.unwrap(); setup.branch4.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version + // Counts include one retained index generation. assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); setup.main.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts for the latest version + // Counts include one retained index generation. assert_eq!(setup.main.counts.num_manifest_files, 1); assert_eq!(setup.main.counts.num_data_files, 1); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 0); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); } #[tokio::test] @@ -4440,7 +4448,7 @@ mod tests { assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); // After auto-clean: branch3 // 2 appends produced 2 data files // 2 deletes produced 2 deletion files @@ -4448,7 +4456,7 @@ mod tests { assert_eq!(setup.branch3.counts.num_data_files, 2); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 2); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup .assert_unchanged(&["branch1", "branch4", "main"]) .await; @@ -4465,19 +4473,19 @@ mod tests { .unwrap(); setup.branch3.refresh().await.unwrap(); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts of one version + // Counts include one retained index generation. assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); // Only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts of one version + // Counts include one retained index generation. assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup .assert_unchanged(&["branch1", "branch4", "main"]) .await; @@ -4502,12 +4510,12 @@ mod tests { // - 1 manifest file // - 3 data files // - 1 deletion file - // - 4 index files + // - one index generation assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 3); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); setup.main.compact().await.unwrap(); setup @@ -4527,7 +4535,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); setup.branch4.compact().await.unwrap(); setup @@ -4544,13 +4552,13 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, 7); - // (1, 1, 1, 0, 4) is the counts of one version + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + // Counts include one retained index generation. assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); setup.branch1.write_data().await.unwrap(); setup.branch1.compact().await.unwrap(); @@ -4568,7 +4576,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); // Branch3 and branch2 still hold references from branch1: // - 1 manifest file // - 1 data files @@ -4577,7 +4585,7 @@ mod tests { assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); setup.branch2.write_data().await.unwrap(); setup.branch2.compact().await.unwrap(); @@ -4595,7 +4603,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); // Branch3 still holds references from branch1: // - 1 manifest file // - 1 data files @@ -4604,7 +4612,7 @@ mod tests { assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); // Branch3 still holds references from branch2: // - 1 manifest file // - 1 data files @@ -4613,7 +4621,7 @@ mod tests { assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); setup.branch3.write_data().await.unwrap(); setup.branch3.compact().await.unwrap(); @@ -4626,27 +4634,27 @@ mod tests { setup.branch2.refresh().await.unwrap(); setup.branch3.refresh().await.unwrap(); // For all branches, only the latest manifest is retained. - // (1, 1, 1, 0, 4) is the counts of one version + // Counts include one retained index generation. assert_eq!(setup.main.counts.num_manifest_files, 1); assert_eq!(setup.main.counts.num_data_files, 1); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 0); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); setup.assert_unchanged(&["branch4"]).await; } @@ -4753,29 +4761,29 @@ mod tests { setup.branch3.refresh().await.unwrap(); setup.branch4.refresh().await.unwrap(); // Two tags hold two manifest references - // Main tag holds 1 tx file, 3 data files, 2 deletion files and 4 index files + // Main tag holds 1 tx file, 3 data files, 2 deletion files and one index generation assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 2); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 14); - // Branch3 tag holds branch1 with 1 tx file, 1 data files, 1 deletion files and 4 index files + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + // Branch3 tag holds branch1 with 1 tx file, 1 data files, 1 deletion files and one index generation assert_eq!(setup.branch2.counts.num_manifest_files, 2); assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, 7); - // Branch3 tag holds branch2 with 1 tx file, 1 data files, 1 deletion files and 4 index files + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + // Branch3 tag holds branch2 with 1 tx file, 1 data files, 1 deletion files and one index generation assert_eq!(setup.branch2.counts.num_manifest_files, 2); assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); setup .branch3 @@ -4798,27 +4806,27 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 2); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, 14); + assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, 7); + assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); setup.main.dataset.tags().delete("main-tag").await.unwrap(); setup @@ -4834,22 +4842,22 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 1); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 0); - assert_eq!(setup.main.counts.num_index_files, 7); + assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, 7); + assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, 7); + assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, 7); + assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); } #[test] From e49d579a1e2ad9bf4da5994ed769c67ab1fd7169 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 18:27:38 +0800 Subject: [PATCH 15/22] test(fts): use exact BM25 scoring oracles --- rust/lance/src/dataset/mem_wal/index/fts.rs | 12 ++++-- rust/lance/src/dataset/tests/dataset_index.rs | 37 +++++++++++++------ 2 files changed, 34 insertions(+), 15 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 5bffeaed517..793a14b1007 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -4511,17 +4511,23 @@ mod tests { let row_ids = UInt64Array::from(vec![900, 42, 777]); let terms = HashSet::from(["hello".to_string()]); let index = FtsMemIndex::new(1, "description".to_string()); + let full_index = FtsMemIndex::new(1, "description".to_string()); index .insert_with_row_ids_for_terms(&batch, &row_ids, &terms) .unwrap(); + full_index.insert(&batch, 0).unwrap(); assert_eq!(index.doc_count(), 3); assert_eq!(index.entry_count(), 2); let scorer = index.bm25_stats_for_terms(&["hello".to_string()]); - assert_eq!(scorer.num_docs, 3); - assert_eq!(scorer.total_tokens, 6); - assert_eq!(scorer.num_docs_containing_token("hello"), 2); + let full_scorer = full_index.bm25_stats_for_terms(&["hello".to_string()]); + assert_eq!(scorer.num_docs, full_scorer.num_docs); + assert_eq!(scorer.total_tokens, full_scorer.total_tokens); + assert_eq!( + scorer.num_docs_containing_token("hello"), + full_scorer.num_docs_containing_token("hello") + ); let query = FtsQuery::Match( lance_index::scalar::inverted::query::MatchQuery::new("hello".to_string()) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 4617317ec33..9774a66bc89 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -1833,7 +1833,7 @@ async fn test_top_level_cross_column_multimatch_uses_field_local_compound_scorer .await .unwrap(); // Index only the title after the append so it can retain a bounded plan - // while the partially covered body uses the exhaustive leaf fallback. + // while the partially covered body uses a query-local hybrid scorer. create_fragmented_fts_index(&mut partial_dataset, "title", true).await; partial_dataset .create_index( @@ -1845,13 +1845,8 @@ async fn test_top_level_cross_column_multimatch_uses_field_local_compound_scorer ) .await .unwrap(); - assert_compound_matches_independent_oracle( - &partial_dataset, - "partial_top_level_cross_column_multimatch", - &explicit_query, - LIMIT, - ) - .await; + let partial_results = + compound_fts_results(&partial_dataset, explicit_query.clone(), Some(LIMIT as i64)).await; let partial_plan = compound_fts_plan(&partial_dataset, explicit_query.clone(), LIMIT).await; assert!( !partial_plan.contains(CROSS_COLUMN_COMPOUND_FTS_SCORER), @@ -1859,18 +1854,23 @@ async fn test_top_level_cross_column_multimatch_uses_field_local_compound_scorer ); assert_eq!( partial_plan.matches("CompoundFtsScorer").count(), + 2, + "both fields should retain field-local bounded compound scorers:\n{partial_plan}" + ); + assert_eq!( + partial_plan.matches("HybridCompoundFtsScorer").count(), 1, - "the fully indexed title should retain its bounded compound scorer:\n{partial_plan}" + "only the partially covered body should use a query-local hybrid scorer:\n{partial_plan}" ); assert!( - partial_plan.contains("FlatMatchQuery"), - "the partially covered body should use the exact indexed-plus-flat fallback:\n{partial_plan}" + !partial_plan.contains("FlatMatchQuery"), + "the hybrid body scorer should replace the indexed-plus-flat fallback:\n{partial_plan}" ); let mut fast_scanner = partial_dataset.scan(); fast_scanner .with_row_id() - .full_text_search(FullTextSearchQuery::new_query(explicit_query)) + .full_text_search(FullTextSearchQuery::new_query(explicit_query.clone())) .unwrap() .fast_search(); fast_scanner.limit(Some(LIMIT as i64), None).unwrap(); @@ -1888,6 +1888,19 @@ async fn test_top_level_cross_column_multimatch_uses_field_local_compound_scorer !fast_plan.contains("FlatMatchQuery"), "fast search must skip the partially covered body's flat path:\n{fast_plan}" ); + + create_fragmented_fts_index(&mut partial_dataset, "body", true).await; + let mut rebuilt_oracle = compound_fts_results(&partial_dataset, explicit_query, None).await; + assert!( + rebuilt_oracle.len() > LIMIT, + "the rebuilt exact oracle must contain candidates beyond k" + ); + rebuilt_oracle.truncate(LIMIT); + assert_scored_rows_close( + "partial_top_level_cross_column_multimatch", + &partial_results, + &rebuilt_oracle, + ); } #[rstest] From 1800f31407f3b17cbc8b5ea367aeb78724cd75d4 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 18:34:53 +0800 Subject: [PATCH 16/22] test(fts): replace index for rebuilt oracle --- rust/lance/src/dataset/tests/dataset_index.rs | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 9774a66bc89..142caf289d0 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -1889,7 +1889,16 @@ async fn test_top_level_cross_column_multimatch_uses_field_local_compound_scorer "fast search must skip the partially covered body's flat path:\n{fast_plan}" ); - create_fragmented_fts_index(&mut partial_dataset, "body", true).await; + partial_dataset + .create_index( + &["body"], + IndexType::Inverted, + Some("body_idx".to_owned()), + &InvertedIndexParams::default().with_position(true), + true, + ) + .await + .unwrap(); let mut rebuilt_oracle = compound_fts_results(&partial_dataset, explicit_query, None).await; assert!( rebuilt_oracle.len() > LIMIT, From 1284b21e255ce80564b4e8dea1becfc945d4f296 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Fri, 28 Aug 2026 23:19:02 +0800 Subject: [PATCH 17/22] perf(fts): parallelize residual compound indexing --- rust/lance/src/dataset/tests/dataset_index.rs | 17 +- rust/lance/src/io/exec/fts.rs | 242 ++++++++++++++---- 2 files changed, 203 insertions(+), 56 deletions(-) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 142caf289d0..fb7db678d56 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2782,15 +2782,24 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { create_fragmented_fts_index(&mut dataset, "text", true).await; let appended = arrow_array::record_batch!( - ("text", Utf8, ["fresh alpha", "fresh beta"]), - ("id", Int32, [2, 3]) + ( + "text", + Utf8, + ["fresh alpha", "fresh beta", "fresh alpha", "fresh gamma"] + ), + ("id", Int32, [2, 3, 4, 5]) ) .unwrap(); let schema = appended.schema(); dataset .append( RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), - None, + Some(WriteParams { + // Keep the residual rows in separate scan batches so the + // parallel query-local shards must merge scores and ties. + max_rows_per_file: 1, + ..Default::default() + }), ) .await .unwrap(); @@ -2810,7 +2819,7 @@ async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { let partial_boost = compound_fts_results(&dataset, boost_query.clone(), Some(10)).await; assert_eq!( partial_boost.len(), - 3, + 5, "MUST_NOT must exclude the blocked row" ); assert_eq!(partial_boost[0].1.to_bits(), partial_boost[1].1.to_bits()); diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 7c6e68ccc00..0fbb29e2fa4 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -25,7 +25,7 @@ use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion_physical_plan::metrics::{BaselineMetrics, Count, Time}; use futures::future::try_join_all; -use futures::stream::{self}; +use futures::stream::{self, FuturesUnordered}; use futures::{FutureExt, StreamExt, TryStreamExt}; use itertools::Itertools; use lance_core::{ @@ -48,6 +48,7 @@ use crate::index::scalar::inverted::{ }; use crate::{Dataset, index::DatasetIndexInternalExt}; use lance_index::metrics::MetricsCollector; +use lance_index::scalar::InvertedIndexParams; use lance_index::scalar::inverted::builder::ScoredDoc; use lance_index::scalar::inverted::builder::document_input; use lance_index::scalar::inverted::document_tokenizer::{DocType, JsonTokenizer, LanceTokenizer}; @@ -796,11 +797,18 @@ impl CompoundQueryExec { } async fn index_query_local_residual_batch( - residual: QueryLocalFtsIndex, + residual: Option, + field_id: i32, + column: String, + index_params: InvertedIndexParams, batch: RecordBatch, allowed_terms: Arc>, ) -> Result { spawn_cpu(move || { + let residual = match residual { + Some(residual) => residual, + None => QueryLocalFtsIndex::try_with_params(field_id, column, index_params)?, + }; let row_ids = batch .column_by_name(ROW_ID) .ok_or_else(|| { @@ -815,6 +823,120 @@ async fn index_query_local_residual_batch( .await } +/// Build a bounded set of independent residual posting shards. +/// +/// A single [`QueryLocalFtsIndex`] intentionally has one writer. Reusing one +/// index per CPU worker preserves that contract while allowing different scan +/// batches to tokenize in parallel. Completed workers immediately take the +/// next batch, so the stream is never collected in memory and the number of +/// live tokenizers/posting maps is bounded by the CPU pool size. +#[allow(clippy::too_many_arguments)] +async fn index_query_local_residual( + mut residual_input: SendableRecordBatchStream, + seed: QueryLocalFtsIndex, + field_id: i32, + column: String, + index_params: InvertedIndexParams, + allowed_terms: Arc>, +) -> DataFusionResult> { + let parallelism = get_num_compute_intensive_cpus().max(1); + let mut seed = Some(seed); + let mut in_flight = FuturesUnordered::new(); + let mut is_input_exhausted = false; + + while in_flight.len() < parallelism { + let Some(batch) = residual_input.try_next().await? else { + is_input_exhausted = true; + break; + }; + in_flight.push(index_query_local_residual_batch( + seed.take(), + field_id, + column.clone(), + index_params.clone(), + batch, + allowed_terms.clone(), + )); + } + + if in_flight.is_empty() { + return Ok(vec![seed.ok_or_else(|| { + DataFusionError::Internal( + "hybrid compound FTS lost its empty residual seed".to_string(), + ) + })?]); + } + + let mut shards = Vec::with_capacity(parallelism.min(in_flight.len())); + while let Some(shard) = in_flight.try_next().await? { + if is_input_exhausted { + shards.push(shard); + continue; + } + match residual_input.try_next().await? { + Some(batch) => in_flight.push(index_query_local_residual_batch( + Some(shard), + field_id, + column.clone(), + index_params.clone(), + batch, + allowed_terms.clone(), + )), + None => { + is_input_exhausted = true; + shards.push(shard); + } + } + } + Ok(shards) +} + +async fn query_local_residual_stats( + shards: Vec, + terms: Arc<[String]>, +) -> Result> { + stream::iter(shards.into_iter().map(|shard| { + let terms = terms.clone(); + spawn_cpu(move || { + let stats = shard.bm25_stats_for_terms(terms.as_ref()); + Ok::<_, Error>((shard, stats)) + }) + })) + .buffered(get_num_compute_intensive_cpus().max(1)) + .try_collect() + .await +} + +async fn query_local_residual_leaves( + shards: Vec, + query: FtsQuery, + scorer: Arc, +) -> Result>> { + let shard_leaves = stream::iter(shards.into_iter().map(|shard| { + let query = query.clone(); + let scorer = scorer.clone(); + spawn_cpu(move || shard.exact_leaf_results(&query, scorer.as_ref())) + })) + .buffered(get_num_compute_intensive_cpus().max(1)) + .try_collect::>() + .await?; + + let leaf_count = shard_leaves.first().map_or(0, Vec::len); + let mut merged = vec![Vec::new(); leaf_count]; + for leaves in shard_leaves { + if leaves.len() != leaf_count { + return Err(Error::internal(format!( + "hybrid compound FTS residual shards produced inconsistent leaf counts: expected {leaf_count}, got {}", + leaves.len() + ))); + } + for (merged, rows) in merged.iter_mut().zip(leaves) { + merged.extend(rows); + } + } + Ok(merged) +} + /// Exact compound FTS over committed postings plus an append-only residual /// scan. The residual documents are tokenized once into query-local postings, /// rather than once for every compound leaf. @@ -928,23 +1050,27 @@ impl ExecutionPlan for HybridCompoundQueryExec { )) })?; let field_id = dataset.schema().field_id(&column)?; - let mut residual = QueryLocalFtsIndex::try_with_params( + let residual_seed = QueryLocalFtsIndex::try_with_params( field_id, column.clone(), first_index.params().clone(), )?; - let terms = residual.exact_query_terms(&query)?; + let terms = residual_seed.exact_query_terms(&query)?; if terms.is_empty() { metrics.baseline_metrics.record_output(0); return scored_documents_batch(schema, Vec::new()).map_err(DataFusionError::from); } let allowed_terms = Arc::new(terms.iter().cloned().collect::>()); - let mut residual_input = residual_input.execute(partition, context.clone())?; - - while let Some(batch) = residual_input.try_next().await? { - residual = index_query_local_residual_batch(residual, batch, allowed_terms.clone()) - .await?; - } + let residual_input = residual_input.execute(partition, context.clone())?; + let residual_shards = index_query_local_residual( + residual_input, + residual_seed, + field_id, + column.clone(), + first_index.params().clone(), + allowed_terms, + ) + .await?; let query_tokens = Tokens::new(terms.clone(), first_index.tokenizer().doc_type()); let exact_params = params @@ -958,41 +1084,46 @@ impl ExecutionPlan for HybridCompoundQueryExec { Some(metrics.as_ref()), ) .await?; - let stats_terms = terms.clone(); - let (residual, residual_stats) = spawn_cpu(move || { - let stats = residual.bm25_stats_for_terms(&stats_terms); - Ok::<_, Error>((residual, stats)) - }) + let residual_shards = query_local_residual_stats( + residual_shards, + Arc::from(terms.clone().into_boxed_slice()), + ) .await?; - scorer.total_tokens = scorer - .total_tokens - .checked_add(residual_stats.total_tokens) - .ok_or_else(|| { - DataFusionError::Execution( - "hybrid compound FTS total token count overflow".to_string(), - ) - })?; - scorer.num_docs = scorer - .num_docs - .checked_add(residual_stats.num_docs) - .ok_or_else(|| { - DataFusionError::Execution( - "hybrid compound FTS document count overflow".to_string(), - ) - })?; - for term in &terms { - let residual_df = residual_stats.num_docs_containing_token(term); - let df = scorer.token_docs.get_mut(term).ok_or_else(|| { - DataFusionError::Execution(format!( - "hybrid compound FTS scorer is missing query term '{term}'" - )) - })?; - *df = df.checked_add(residual_df).ok_or_else(|| { - DataFusionError::Execution(format!( - "hybrid compound FTS document frequency overflow for term '{term}'" - )) - })?; + for (_, residual_stats) in &residual_shards { + scorer.total_tokens = scorer + .total_tokens + .checked_add(residual_stats.total_tokens) + .ok_or_else(|| { + DataFusionError::Execution( + "hybrid compound FTS total token count overflow".to_string(), + ) + })?; + scorer.num_docs = scorer + .num_docs + .checked_add(residual_stats.num_docs) + .ok_or_else(|| { + DataFusionError::Execution( + "hybrid compound FTS document count overflow".to_string(), + ) + })?; + for term in &terms { + let residual_df = residual_stats.num_docs_containing_token(term); + let df = scorer.token_docs.get_mut(term).ok_or_else(|| { + DataFusionError::Execution(format!( + "hybrid compound FTS scorer is missing query term '{term}'" + )) + })?; + *df = df.checked_add(residual_df).ok_or_else(|| { + DataFusionError::Execution(format!( + "hybrid compound FTS document frequency overflow for term '{term}'" + )) + })?; + } } + let residual_shards = residual_shards + .into_iter() + .map(|(shard, _)| shard) + .collect::>(); let scorer = Arc::new(scorer); let limit = params.limit.ok_or_else(|| { DataFusionError::Execution( @@ -1011,23 +1142,30 @@ impl ExecutionPlan for HybridCompoundQueryExec { external_mask: None, }, )?; - let (indexed_row_ids, indexed_scores) = compound_search_with_base_scorer( + let indexed_search = compound_search_with_base_scorer( &indices, &query, ¶ms, prefilter, metrics.clone(), scorer.clone(), - ) - .await?; + ); let residual_query = query.clone(); let residual_scorer = scorer.clone(); - let (residual_row_ids, residual_scores) = spawn_cpu(move || { - let residual_leaves = - residual.exact_leaf_results(&residual_query, residual_scorer.as_ref())?; - materialized_compound_top_k(&residual_query, residual_leaves, limit) - }) - .await?; + let residual_search = async move { + let residual_leaves = query_local_residual_leaves( + residual_shards, + residual_query.clone(), + residual_scorer, + ) + .await?; + spawn_cpu(move || { + materialized_compound_top_k(&residual_query, residual_leaves, limit) + }) + .await + }; + let ((indexed_row_ids, indexed_scores), (residual_row_ids, residual_scores)) = + futures::future::try_join(indexed_search, residual_search).await?; let mut documents = indexed_row_ids .into_iter() From 933a8b0ecfe77101d60507a25e2923c47915f1ff Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sat, 29 Aug 2026 01:20:05 +0800 Subject: [PATCH 18/22] perf(fts): reuse tokenizer assets across residual shards --- rust/lance/src/dataset/mem_wal/index/fts.rs | 42 +++++++++++++ rust/lance/src/io/exec/fts.rs | 69 ++++++++------------- 2 files changed, 68 insertions(+), 43 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 793a14b1007..5a9e950b3e9 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1043,6 +1043,34 @@ impl QueryLocalFtsIndex { }) } + /// Create an empty query-local shard without rebuilding tokenizer assets. + /// + /// The tokenizer pool and its loaded template are shared with the seed; + /// each shard only clones a writer tokenizer from that in-memory template. + pub(crate) fn empty_sibling(&self) -> Self { + let resolved_field = OnceLock::new(); + if let Some(resolved) = self.inner.resolved_field.get() { + resolved_field + .set(resolved.clone()) + .expect("new query-local shard traversal is empty"); + } + + Self { + inner: FtsMemIndex { + field_id: self.inner.field_id, + source_column_name: self.inner.source_column_name.clone(), + params: self.inner.params.clone(), + resolved_field, + tokenizer_pool: self.inner.tokenizer_pool.clone(), + writer_tokenizer: Mutex::new(self.inner.tokenizer_pool.acquire()), + state: ArcSwap::from(IndexState::empty()), + freeze_threshold_rows: self.inner.freeze_threshold_rows, + background_maintenance: false, + merge: Arc::new(Mutex::new(None)), + }, + } + } + pub(crate) fn exact_query_terms(&self, query: &FtsQuery) -> Result> { self.inner.exact_query_terms(query) } @@ -4565,6 +4593,20 @@ mod tests { assert!(index.inner.state.load().partitions.is_empty()); assert!(index.inner.merge.lock().unwrap().is_none()); assert_eq!(index.doc_count(), 3); + + let sibling = index.empty_sibling(); + assert!(Arc::ptr_eq( + &index.inner.tokenizer_pool, + &sibling.inner.tokenizer_pool + )); + assert_eq!(sibling.doc_count(), 0); + sibling + .insert_with_row_ids_for_terms(&batch, &UInt64Array::from(vec![901, 43, 778]), &terms) + .unwrap(); + assert_eq!(index.doc_count(), 3); + assert_eq!(sibling.doc_count(), 3); + assert!(sibling.inner.state.load().partitions.is_empty()); + assert!(sibling.inner.merge.lock().unwrap().is_none()); } fn create_element_test_batch() -> RecordBatch { diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 0fbb29e2fa4..84f5a575a9a 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -48,7 +48,6 @@ use crate::index::scalar::inverted::{ }; use crate::{Dataset, index::DatasetIndexInternalExt}; use lance_index::metrics::MetricsCollector; -use lance_index::scalar::InvertedIndexParams; use lance_index::scalar::inverted::builder::ScoredDoc; use lance_index::scalar::inverted::builder::document_input; use lance_index::scalar::inverted::document_tokenizer::{DocType, JsonTokenizer, LanceTokenizer}; @@ -797,18 +796,11 @@ impl CompoundQueryExec { } async fn index_query_local_residual_batch( - residual: Option, - field_id: i32, - column: String, - index_params: InvertedIndexParams, + residual: QueryLocalFtsIndex, batch: RecordBatch, allowed_terms: Arc>, ) -> Result { spawn_cpu(move || { - let residual = match residual { - Some(residual) => residual, - None => QueryLocalFtsIndex::try_with_params(field_id, column, index_params)?, - }; let row_ids = batch .column_by_name(ROW_ID) .ok_or_else(|| { @@ -828,45 +820,46 @@ async fn index_query_local_residual_batch( /// A single [`QueryLocalFtsIndex`] intentionally has one writer. Reusing one /// index per CPU worker preserves that contract while allowing different scan /// batches to tokenize in parallel. Completed workers immediately take the -/// next batch, so the stream is never collected in memory and the number of -/// live tokenizers/posting maps is bounded by the CPU pool size. -#[allow(clippy::too_many_arguments)] +/// next batch, so the entire stream is never collected in memory and the +/// number of live tokenizers/posting maps is bounded by the CPU pool size. async fn index_query_local_residual( mut residual_input: SendableRecordBatchStream, seed: QueryLocalFtsIndex, - field_id: i32, - column: String, - index_params: InvertedIndexParams, allowed_terms: Arc>, ) -> DataFusionResult> { let parallelism = get_num_compute_intensive_cpus().max(1); - let mut seed = Some(seed); - let mut in_flight = FuturesUnordered::new(); + let mut initial_batches = Vec::with_capacity(parallelism); let mut is_input_exhausted = false; - while in_flight.len() < parallelism { + while initial_batches.len() < parallelism { let Some(batch) = residual_input.try_next().await? else { is_input_exhausted = true; break; }; + initial_batches.push(batch); + } + + if initial_batches.is_empty() { + return Ok(vec![seed]); + } + + // Construct every shard from the already-loaded seed before dispatching + // CPU work. This keeps tokenizer model I/O out of `spawn_cpu` closures. + let mut initial_shards = Vec::with_capacity(initial_batches.len()); + for _ in 1..initial_batches.len() { + initial_shards.push(seed.empty_sibling()); + } + initial_shards.push(seed); + + let mut in_flight = FuturesUnordered::new(); + for (shard, batch) in initial_shards.into_iter().zip(initial_batches) { in_flight.push(index_query_local_residual_batch( - seed.take(), - field_id, - column.clone(), - index_params.clone(), + shard, batch, allowed_terms.clone(), )); } - if in_flight.is_empty() { - return Ok(vec![seed.ok_or_else(|| { - DataFusionError::Internal( - "hybrid compound FTS lost its empty residual seed".to_string(), - ) - })?]); - } - let mut shards = Vec::with_capacity(parallelism.min(in_flight.len())); while let Some(shard) = in_flight.try_next().await? { if is_input_exhausted { @@ -875,10 +868,7 @@ async fn index_query_local_residual( } match residual_input.try_next().await? { Some(batch) => in_flight.push(index_query_local_residual_batch( - Some(shard), - field_id, - column.clone(), - index_params.clone(), + shard, batch, allowed_terms.clone(), )), @@ -1062,15 +1052,8 @@ impl ExecutionPlan for HybridCompoundQueryExec { } let allowed_terms = Arc::new(terms.iter().cloned().collect::>()); let residual_input = residual_input.execute(partition, context.clone())?; - let residual_shards = index_query_local_residual( - residual_input, - residual_seed, - field_id, - column.clone(), - first_index.params().clone(), - allowed_terms, - ) - .await?; + let residual_shards = + index_query_local_residual(residual_input, residual_seed, allowed_terms).await?; let query_tokens = Tokens::new(terms.clone(), first_index.tokenizer().doc_type()); let exact_params = params From 3889ecaafe33d3bd09fbf95ae6c2e2e235f81474 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sun, 30 Aug 2026 16:57:31 +0800 Subject: [PATCH 19/22] perf(fts): use indexed stats for residual compound rows --- rust/lance/src/dataset/cleanup.rs | 184 ++-- rust/lance/src/dataset/mem_wal/index/fts.rs | 257 ++++-- rust/lance/src/dataset/scanner.rs | 194 ++-- rust/lance/src/dataset/tests/dataset_index.rs | 869 ++++-------------- rust/lance/src/index.rs | 124 +-- rust/lance/src/index/append.rs | 92 +- rust/lance/src/index/create.rs | 20 +- rust/lance/src/index/scalar/inverted.rs | 173 +--- rust/lance/src/io/exec/fts.rs | 125 +-- 9 files changed, 610 insertions(+), 1428 deletions(-) diff --git a/rust/lance/src/dataset/cleanup.rs b/rust/lance/src/dataset/cleanup.rs index b820eca849c..ad35669f37b 100644 --- a/rust/lance/src/dataset/cleanup.rs +++ b/rust/lance/src/dataset/cleanup.rs @@ -3673,14 +3673,6 @@ mod tests { // Note: branch2 is stored as "dev/branch2"; comments may refer to it as branch2 for brevity. // Important: auto_cleanup_hook uses policy derived from manifest config; it does not flip // clean_referenced_branches unless tests call cleanup_old_versions with a custom policy. - // Cleanup retains or removes an index UUID directory as a unit. Each - // inverted-index generation has seven core files plus its provenance sidecar. - const INDEX_FILES_PER_GENERATION: usize = 8; - - fn index_file_count(generations: usize) -> usize { - generations * INDEX_FILES_PER_GENERATION - } - struct LineageSetup { main: BranchDatasetFixture, branch1: BranchDatasetFixture, @@ -4100,13 +4092,13 @@ mod tests { // - 1 manifest file // - 1 data file // - 1 deletion file - // - one index generation + // - 4 index files // The left is the counts for the latest version of appending assert_eq!(setup.branch1.counts.num_manifest_files, 2); assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 2); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.branch1.counts.num_index_files, 14); setup.assert_all_unchanged().await; setup.branch1.compact().await.unwrap(); @@ -4115,13 +4107,13 @@ mod tests { // - 1 manifest file // - 1 data file // - 1 deletion file - // - one index generation - // Counts include one retained index generation. + // - 4 index files + // The left (1, 1, 1, 0, 4) is the counts for the latest version of compaction assert_eq!(setup.branch1.counts.num_manifest_files, 2); assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.branch1.counts.num_index_files, 14); setup.assert_all_unchanged().await; // Now we clean the referenced files of branch1 by branch2 and branch3 @@ -4130,28 +4122,28 @@ mod tests { setup.branch3.run_cleanup().await.unwrap(); setup.branch2.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version of compaction assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version of compaction assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup.branch1.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version of compaction assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); setup.assert_all_unchanged().await; } @@ -4168,7 +4160,7 @@ mod tests { assert_eq!(setup.branch3.counts.num_data_files, 2); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 2); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup .assert_unchanged(&["branch1", "branch2", "branch4", "main"]) .await; @@ -4184,17 +4176,17 @@ mod tests { assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); setup.branch3.compact().await.unwrap(); setup.branch3.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup .assert_unchanged(&["branch1", "branch2", "branch4", "main"]) .await; @@ -4202,12 +4194,12 @@ mod tests { setup.branch2.compact().await.unwrap(); setup.branch2.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); } #[tokio::test] @@ -4224,7 +4216,7 @@ mod tests { assert_eq!(setup.branch4.counts.num_data_files, 2); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 2); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); setup.assert_all_unchanged().await; setup.main.compact().await.unwrap(); @@ -4233,28 +4225,28 @@ mod tests { // - 1 manifest file // - 2 data files // - 1 deletion file - // - one index generation + // - 4 index files // Branch4 holds references from main: // - 1 manifest file // - 3 data files // - 1 deletion file - // - one index generation + // - 4 index files // The left(1, 1, 1, 0, 0) is the counts for the latest version of compaction assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.main.counts.num_index_files, 14); setup.branch4.compact().await.unwrap(); setup.branch4.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts of one version assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); setup.assert_all_unchanged().await; setup.main.run_cleanup().await.unwrap(); @@ -4262,13 +4254,13 @@ mod tests { // - 1 manifest file // - 2 data files // - 1 deletion file - // - one index generation - // Counts include one retained index generation. + // - 4 index files + // The left(1, 1, 1, 0, 4) is the counts for the latest version of compaction assert_eq!(setup.main.counts.num_manifest_files, 2); assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.main.counts.num_index_files, 14); } #[tokio::test] @@ -4282,18 +4274,18 @@ mod tests { // - 1 manifest file // - 2 data files // - 1 deletion file - // - one index generation(only for branch1) + // - 4 index files(only for branch1) // Branch4 holds references from main: // - 1 manifest file // - 3 data files // - 1 deletion file - // - one index generation - // Counts include one retained index generation. + // - 4 index files + // The left(1, 1, 1, 1, 4) is the counts for the latest version of compaction assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 3); - assert_eq!(setup.main.counts.num_index_files, index_file_count(3)); + assert_eq!(setup.main.counts.num_index_files, 21); setup.assert_all_unchanged().await; setup.main.compact().await.unwrap(); @@ -4304,7 +4296,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(3)); + assert_eq!(setup.main.counts.num_index_files, 21); setup.assert_all_unchanged().await; setup.branch1.write_data().await.unwrap(); @@ -4320,19 +4312,19 @@ mod tests { // - 1 manifest file // - 1 data files // - 2 deletion files - // - one index generation + // - 4 index files assert_eq!(setup.branch2.counts.num_manifest_files, 2); assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.branch2.counts.num_index_files, 14); setup.branch1.run_cleanup().await.unwrap(); - // Cleanup one index generation referenced from branch2 + // Cleanup 4 index files referenced from branch2 assert_eq!(setup.branch1.counts.num_manifest_files, 2); assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); setup.main.run_cleanup().await.unwrap(); // Branch3 holds references from main: @@ -4343,12 +4335,12 @@ mod tests { // - 1 manifest file // - 3 data files // - 2 deletion files - // - one index generation + // - 4 index files assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.main.counts.num_index_files, 14); setup.branch3.write_data().await.unwrap(); setup.branch3.compact().await.unwrap(); @@ -4358,7 +4350,7 @@ mod tests { assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup.main.run_cleanup().await.unwrap(); // Cleanup doesn't take effects if we don't clean branch2 and branch1 first @@ -4366,7 +4358,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.main.counts.num_index_files, 14); // Cleanup doesn't take effect if we don't clean branch2 first setup.branch1.run_cleanup().await.unwrap(); @@ -4374,57 +4366,57 @@ mod tests { assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); setup.branch2.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); setup.branch1.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); setup.main.run_cleanup().await.unwrap(); // Branch4 holds references from main: // - 1 manifest file // - 3 data files // - 2 deletion files - // - one index generation + // - 4 index files assert_eq!(setup.main.counts.num_manifest_files, 2); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.main.counts.num_index_files, 14); setup.branch4.write_data().await.unwrap(); setup.branch4.compact().await.unwrap(); setup.branch4.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); setup.main.run_cleanup().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts for the latest version assert_eq!(setup.main.counts.num_manifest_files, 1); assert_eq!(setup.main.counts.num_data_files, 1); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 0); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); } #[tokio::test] @@ -4448,7 +4440,7 @@ mod tests { assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); // After auto-clean: branch3 // 2 appends produced 2 data files // 2 deletes produced 2 deletion files @@ -4456,7 +4448,7 @@ mod tests { assert_eq!(setup.branch3.counts.num_data_files, 2); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 2); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup .assert_unchanged(&["branch1", "branch4", "main"]) .await; @@ -4473,19 +4465,19 @@ mod tests { .unwrap(); setup.branch3.refresh().await.unwrap(); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts of one version assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); // Only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts of one version assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup .assert_unchanged(&["branch1", "branch4", "main"]) .await; @@ -4510,12 +4502,12 @@ mod tests { // - 1 manifest file // - 3 data files // - 1 deletion file - // - one index generation + // - 4 index files assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 3); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); setup.main.compact().await.unwrap(); setup @@ -4535,7 +4527,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); setup.branch4.compact().await.unwrap(); setup @@ -4552,13 +4544,13 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); - // Counts include one retained index generation. + assert_eq!(setup.main.counts.num_index_files, 7); + // (1, 1, 1, 0, 4) is the counts of one version assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); setup.branch1.write_data().await.unwrap(); setup.branch1.compact().await.unwrap(); @@ -4576,7 +4568,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); // Branch3 and branch2 still hold references from branch1: // - 1 manifest file // - 1 data files @@ -4585,7 +4577,7 @@ mod tests { assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); setup.branch2.write_data().await.unwrap(); setup.branch2.compact().await.unwrap(); @@ -4603,7 +4595,7 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 3); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 1); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); // Branch3 still holds references from branch1: // - 1 manifest file // - 1 data files @@ -4612,7 +4604,7 @@ mod tests { assert_eq!(setup.branch1.counts.num_data_files, 2); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 1); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); // Branch3 still holds references from branch2: // - 1 manifest file // - 1 data files @@ -4621,7 +4613,7 @@ mod tests { assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); setup.branch3.write_data().await.unwrap(); setup.branch3.compact().await.unwrap(); @@ -4634,27 +4626,27 @@ mod tests { setup.branch2.refresh().await.unwrap(); setup.branch3.refresh().await.unwrap(); // For all branches, only the latest manifest is retained. - // Counts include one retained index generation. + // (1, 1, 1, 0, 4) is the counts of one version assert_eq!(setup.main.counts.num_manifest_files, 1); assert_eq!(setup.main.counts.num_data_files, 1); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 0); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); setup.assert_unchanged(&["branch4"]).await; } @@ -4761,29 +4753,29 @@ mod tests { setup.branch3.refresh().await.unwrap(); setup.branch4.refresh().await.unwrap(); // Two tags hold two manifest references - // Main tag holds 1 tx file, 3 data files, 2 deletion files and one index generation + // Main tag holds 1 tx file, 3 data files, 2 deletion files and 4 index files assert_eq!(setup.main.counts.num_manifest_files, 3); assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 2); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); - // Branch3 tag holds branch1 with 1 tx file, 1 data files, 1 deletion files and one index generation + assert_eq!(setup.main.counts.num_index_files, 14); + // Branch3 tag holds branch1 with 1 tx file, 1 data files, 1 deletion files and 4 index files assert_eq!(setup.branch2.counts.num_manifest_files, 2); assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); - // Branch3 tag holds branch2 with 1 tx file, 1 data files, 1 deletion files and one index generation + assert_eq!(setup.branch2.counts.num_index_files, 7); + // Branch3 tag holds branch2 with 1 tx file, 1 data files, 1 deletion files and 4 index files assert_eq!(setup.branch2.counts.num_manifest_files, 2); assert_eq!(setup.branch2.counts.num_data_files, 2); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 1); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); setup .branch3 @@ -4806,27 +4798,27 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 4); assert_eq!(setup.main.counts.num_tx_files, 2); assert_eq!(setup.main.counts.num_delete_files, 2); - assert_eq!(setup.main.counts.num_index_files, index_file_count(2)); + assert_eq!(setup.main.counts.num_index_files, 14); assert_eq!(setup.branch1.counts.num_manifest_files, 1); assert_eq!(setup.branch1.counts.num_data_files, 1); assert_eq!(setup.branch1.counts.num_tx_files, 1); assert_eq!(setup.branch1.counts.num_delete_files, 0); - assert_eq!(setup.branch1.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch1.counts.num_index_files, 7); assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); setup.main.dataset.tags().delete("main-tag").await.unwrap(); setup @@ -4842,22 +4834,22 @@ mod tests { assert_eq!(setup.main.counts.num_data_files, 1); assert_eq!(setup.main.counts.num_tx_files, 1); assert_eq!(setup.main.counts.num_delete_files, 0); - assert_eq!(setup.main.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.main.counts.num_index_files, 7); assert_eq!(setup.branch2.counts.num_manifest_files, 1); assert_eq!(setup.branch2.counts.num_data_files, 1); assert_eq!(setup.branch2.counts.num_tx_files, 1); assert_eq!(setup.branch2.counts.num_delete_files, 0); - assert_eq!(setup.branch2.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch2.counts.num_index_files, 7); assert_eq!(setup.branch3.counts.num_manifest_files, 1); assert_eq!(setup.branch3.counts.num_data_files, 1); assert_eq!(setup.branch3.counts.num_tx_files, 1); assert_eq!(setup.branch3.counts.num_delete_files, 0); - assert_eq!(setup.branch3.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch3.counts.num_index_files, 7); assert_eq!(setup.branch4.counts.num_manifest_files, 1); assert_eq!(setup.branch4.counts.num_data_files, 1); assert_eq!(setup.branch4.counts.num_tx_files, 1); assert_eq!(setup.branch4.counts.num_delete_files, 0); - assert_eq!(setup.branch4.counts.num_index_files, index_file_count(1)); + assert_eq!(setup.branch4.counts.num_index_files, 7); } #[test] diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 5a9e950b3e9..5f67cf78e11 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -63,7 +63,7 @@ use lance_index::scalar::inverted::tokenizer::document_tokenizer::{DocType, Lanc use lance_index::scalar::inverted::{DocSet, MemBM25Scorer, Scorer, TokenSet}; use lance_tokenizer::TokenStream; use rayon::prelude::*; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use super::RowPosition; use crate::index::scalar::inverted::{ResolvedFtsField, resolve_fts_field}; @@ -770,12 +770,15 @@ impl std::fmt::Debug for TokenizerPool { impl TokenizerPool { fn new(params: &InvertedIndexParams, cap: usize) -> Result { - let template = params.build()?; - Ok(Self { + Ok(Self::from_template(params.build()?, cap)) + } + + fn from_template(template: Box, cap: usize) -> Self { + Self { template, free: Mutex::new(Vec::new()), cap: cap.max(1), - }) + } } /// Acquire a tokenizer. Pops from the free list, otherwise clones the @@ -1016,7 +1019,7 @@ pub struct FtsMemIndex { merge: Arc>>, } -/// Query-owned exact postings for one residual scan. +/// Query-owned term-only postings for one residual scan. /// /// This deliberately exposes only the immutable feature-materialization API /// needed by hybrid execution. Unlike [`FtsMemIndex`], it never freezes or @@ -1028,6 +1031,7 @@ pub struct QueryLocalFtsIndex { } impl QueryLocalFtsIndex { + #[cfg(test)] pub(crate) fn try_with_params( field_id: i32, column_name: String, @@ -1043,6 +1047,25 @@ impl QueryLocalFtsIndex { }) } + pub(crate) fn try_with_loaded_tokenizer( + field_id: i32, + column_name: String, + params: InvertedIndexParams, + tokenizer: Box, + ) -> Result { + params.validate_format_version()?; + let pool = TokenizerPool::from_template(tokenizer, FtsMemIndex::DEFAULT_TOKENIZER_POOL_CAP); + Ok(Self { + inner: FtsMemIndex::with_tokenizer_pool_and_maintenance( + field_id, + column_name, + params, + pool, + false, + ), + }) + } + /// Create an empty query-local shard without rebuilding tokenizer assets. /// /// The tokenizer pool and its loaded template are shared with the seed; @@ -1079,7 +1102,7 @@ impl QueryLocalFtsIndex { &self, batch: &RecordBatch, row_ids: &UInt64Array, - terms: &HashSet, + terms: &FxHashSet, ) -> Result<()> { self.inner .insert_with_row_ids_for_terms(batch, row_ids, terms) @@ -1090,10 +1113,6 @@ impl QueryLocalFtsIndex { self.inner.doc_count() } - pub(crate) fn bm25_stats_for_terms(&self, terms: &[String]) -> MemBM25Scorer { - self.inner.bm25_stats_for_terms(terms) - } - pub(crate) fn exact_leaf_results( &self, query: &FtsQuery, @@ -1194,8 +1213,24 @@ impl FtsMemIndex { ) -> Result { params.validate_format_version()?; let pool = TokenizerPool::new(¶ms, Self::DEFAULT_TOKENIZER_POOL_CAP)?; + Ok(Self::with_tokenizer_pool_and_maintenance( + field_id, + column_name, + params, + pool, + background_maintenance, + )) + } + + fn with_tokenizer_pool_and_maintenance( + field_id: i32, + column_name: String, + params: InvertedIndexParams, + pool: TokenizerPool, + background_maintenance: bool, + ) -> Self { let writer_tokenizer = pool.template.box_clone(); - Ok(Self { + Self { field_id, source_column_name: column_name, params, @@ -1206,7 +1241,7 @@ impl FtsMemIndex { freeze_threshold_rows: Self::DEFAULT_FREEZE_THRESHOLD_ROWS, background_maintenance, merge: Arc::new(Mutex::new(None)), - }) + } } pub(crate) fn try_with_resolved_field( @@ -1361,12 +1396,12 @@ impl FtsMemIndex { /// Insert explicit, potentially non-contiguous rows while retaining /// postings only for query terms. /// The tokenizer still visits the complete document so BM25 document - /// length and corpus totals remain identical to a full index. + /// lengths remain accurate when scoring with committed-index statistics. pub(crate) fn insert_with_row_ids_for_terms( &self, batch: &RecordBatch, row_ids: &UInt64Array, - terms: &HashSet, + terms: &FxHashSet, ) -> Result<()> { if row_ids.len() != batch.num_rows() || row_ids.null_count() != 0 { return Err(Error::invalid_input(format!( @@ -1395,7 +1430,7 @@ impl FtsMemIndex { &self, batch: &RecordBatch, row_position: impl Fn(usize) -> Result, - allowed_terms: Option<&HashSet>, + allowed_terms: Option<&FxHashSet>, ) -> Result<()> { let st = self.state.load_full(); let document_position_start = st.tail.doc_count(); @@ -1424,20 +1459,35 @@ impl FtsMemIndex { // per-document map and per-`(term, doc)` `Vec` allocation that // dominated insert cost. `FxHashMap` skips SipHash on the hot lookup. let mut term_builders: FxHashMap, BatchTermBuilder> = FxHashMap::default(); - let mut documents: Vec = Vec::with_capacity(batch.num_rows()); + let mut documents: Vec = if allowed_terms.is_some() { + Vec::new() + } else { + Vec::with_capacity(batch.num_rows()) + }; let mut total_tokens: u64 = 0; let preserve_zero_token_documents = self.params.get_document_granularity().is_list_element(); let mut index_document = |key: DocumentKey, text: &str| -> Result<()> { let document_position = document_position_start + documents.len() as u64; - let num_tokens = index_text_filtered( - text, - document_position, - tokenizer, - &mut term_builders, - allowed_terms, - )?; - if preserve_zero_token_documents || num_tokens > 0 { + let (num_tokens, retained_term) = match allowed_terms { + Some(allowed_terms) => index_text_filtered( + text, + document_position, + tokenizer, + &mut term_builders, + allowed_terms, + )?, + None => ( + index_text(text, document_position, tokenizer, &mut term_builders)?, + false, + ), + }; + let retain_document = if allowed_terms.is_some() { + retained_term + } else { + preserve_zero_token_documents || num_tokens > 0 + }; + if retain_document { documents.push(DocumentMetadata { key, num_tokens }); total_tokens += num_tokens as u64; } @@ -1524,16 +1574,9 @@ impl FtsMemIndex { Ok(terms) } - /// Build exact residual corpus statistics for the supplied query terms. - pub(crate) fn bm25_stats_for_terms(&self, terms: &[String]) -> MemBM25Scorer { - let st = self.state.load_full(); - let tail = st.tail.snapshot(); - build_scorer(&st, &tail, terms, true) - } - - /// Materialize each exact leaf with a caller-supplied logical-corpus - /// scorer. Compound semantics are deliberately evaluated by the canonical - /// lance-index scorer instead of being duplicated here. + /// Materialize each exact leaf with a caller-supplied scorer. Compound + /// semantics are deliberately evaluated by the canonical lance-index + /// scorer instead of being duplicated here. pub(crate) fn exact_leaf_results( &self, query: &FtsQuery, @@ -1810,6 +1853,7 @@ impl FtsMemIndex { Operator::Or, &scorer, theta, + false, ) { topk.offer(e.score, e.key()); } @@ -1829,6 +1873,7 @@ impl FtsMemIndex { operator, &scorer, f32::NEG_INFINITY, + false, )); } results @@ -1886,7 +1931,7 @@ impl FtsMemIndex { operator: Operator, scorer: &MemBM25Scorer, ) -> Vec { - if tokens.is_empty() || scorer.num_docs() == 0 { + if tokens.is_empty() { return Vec::new(); } let tail = st.tail.snapshot(); @@ -1901,6 +1946,7 @@ impl FtsMemIndex { operator, scorer, f32::NEG_INFINITY, + true, )); results } @@ -2645,15 +2691,39 @@ impl BatchTermBuilder { } } -fn index_text_filtered( +fn index_text( text: &str, document_position: u64, tokenizer: &mut dyn LanceTokenizer, term_builders: &mut FxHashMap, BatchTermBuilder>, - allowed_terms: Option<&HashSet>, ) -> Result { + index_text_with_predicate(text, document_position, tokenizer, term_builders, |_| true) + .map(|(num_tokens, _)| num_tokens) +} + +fn index_text_filtered( + text: &str, + document_position: u64, + tokenizer: &mut dyn LanceTokenizer, + term_builders: &mut FxHashMap, BatchTermBuilder>, + allowed_terms: &FxHashSet, +) -> Result<(u32, bool)> { + index_text_with_predicate(text, document_position, tokenizer, term_builders, |term| { + allowed_terms.contains(term) + }) +} + +#[inline] +fn index_text_with_predicate( + text: &str, + document_position: u64, + tokenizer: &mut dyn LanceTokenizer, + term_builders: &mut FxHashMap, BatchTermBuilder>, + mut retain_term: impl FnMut(&str) -> bool, +) -> Result<(u32, bool)> { let mut stream = tokenizer.token_stream_for_doc(text); let mut num_tokens = 0u32; + let mut retained_term = false; while let Some(token) = stream.next() { let position = u32::try_from(token.position).map_err(|_| { Error::invalid_input(format!( @@ -2662,7 +2732,8 @@ fn index_text_filtered( )) })?; let term = token.text.as_str(); - if allowed_terms.is_none_or(|allowed| allowed.contains(term)) { + if retain_term(term) { + retained_term = true; if let Some(builder) = term_builders.get_mut(term) { builder.observe(document_position, position); } else { @@ -2678,7 +2749,7 @@ fn index_text_filtered( )) })?; } - Ok(num_tokens) + Ok((num_tokens, retained_term)) } fn has_visible_chunk(slice: &TermSlice, visible_count: usize) -> bool { @@ -2769,6 +2840,11 @@ fn tail_token_df( /// Score `tokens` against the visible tail, summing each token's BM25 /// contribution per document. Uses the shared corpus-wide `scorer`. +/// +/// `retain_zero_weight_matches` is reserved for query-local residual postings +/// scored with committed-index statistics. A term absent from the committed +/// corpus has zero BM25 weight, but its fresh matching rows must remain visible +/// to compound membership and MUST_NOT evaluation. fn score_terms( snap: &Snapshot, terms: &SkipMap, Arc>>, @@ -2776,6 +2852,7 @@ fn score_terms( operator: Operator, scorer: &MemBM25Scorer, theta: f32, + retain_zero_weight_matches: bool, ) -> Vec { // Per-token tail data + its score upper bound (max freq over visible chunks, // scored at the most generous doc length of 1). If even the sum of those @@ -2791,7 +2868,7 @@ fn score_terms( continue; }; let qw = scorer.query_weight(token); - if qw == 0.0 { + if qw == 0.0 && !retain_zero_weight_matches { continue; } let slice = entry.value().load_full(); @@ -2801,7 +2878,9 @@ fn score_terms( .map(|c| c.max_freq) .max() .unwrap_or(0); - tail_ub += qw * scorer.doc_weight(max_freq, 1); + if qw != 0.0 { + tail_ub += qw * scorer.doc_weight(max_freq, 1); + } tail_terms.push((qw, slice)); } if tail_ub <= theta { @@ -2818,8 +2897,12 @@ fn score_terms( continue; }; for (i, &document_position) in chunk.row_positions.iter().enumerate() { - let dl = meta.dl(document_position).unwrap_or(1); - let score = qw * scorer.doc_weight(chunk.frequencies[i], dl); + let score = if qw == 0.0 { + 0.0 + } else { + let dl = meta.dl(document_position).unwrap_or(1); + qw * scorer.doc_weight(chunk.frequencies[i], dl) + }; *doc_scores.entry(document_position).or_default() += score; if let Some(doc_hits) = &mut doc_hits { *doc_hits.entry(document_position).or_default() += 1; @@ -4533,12 +4616,17 @@ mod tests { } #[test] - fn explicit_row_ids_and_query_term_allowlist_preserve_bm25_stats() { + fn query_term_allowlist_preserves_document_lengths_with_external_scorer() { let schema = create_test_schema(); let batch = create_test_batch(schema.as_ref()); let row_ids = UInt64Array::from(vec![900, 42, 777]); - let terms = HashSet::from(["hello".to_string()]); - let index = FtsMemIndex::new(1, "description".to_string()); + let terms = FxHashSet::from_iter(["hello".to_string()]); + let index = QueryLocalFtsIndex::try_with_params( + 1, + "description".to_string(), + InvertedIndexParams::default(), + ) + .unwrap(); let full_index = FtsMemIndex::new(1, "description".to_string()); index @@ -4546,42 +4634,85 @@ mod tests { .unwrap(); full_index.insert(&batch, 0).unwrap(); - assert_eq!(index.doc_count(), 3); - assert_eq!(index.entry_count(), 2); - let scorer = index.bm25_stats_for_terms(&["hello".to_string()]); - let full_scorer = full_index.bm25_stats_for_terms(&["hello".to_string()]); - assert_eq!(scorer.num_docs, full_scorer.num_docs); - assert_eq!(scorer.total_tokens, full_scorer.total_tokens); - assert_eq!( - scorer.num_docs_containing_token("hello"), - full_scorer.num_docs_containing_token("hello") - ); + // The unmatched row (row id 42) contributes neither postings nor + // metadata; retained documents still keep their full token counts. + assert_eq!(index.doc_count(), 2); + assert_eq!(index.inner.entry_count(), 2); + let committed_scorer = MemBM25Scorer::new(6, 3, HashMap::from([("hello".to_string(), 2)])); let query = FtsQuery::Match( lance_index::scalar::inverted::query::MatchQuery::new("hello".to_string()) .with_column(Some("description".to_string())), ); - let leaves = index.exact_leaf_results(&query, &scorer).unwrap(); + let leaves = index.exact_leaf_results(&query, &committed_scorer).unwrap(); + let full_leaves = full_index + .exact_leaf_results(&query, &committed_scorer) + .unwrap(); let mut actual = leaves[0] .iter() .map(|(row_id, _)| *row_id) .collect::>(); actual.sort_unstable(); assert_eq!(actual, vec![777, 900]); + let mut actual_scores = leaves[0] + .iter() + .map(|(_, score)| score.to_bits()) + .collect::>(); + let mut full_scores = full_leaves[0] + .iter() + .map(|(_, score)| score.to_bits()) + .collect::>(); + actual_scores.sort_unstable(); + full_scores.sort_unstable(); + assert_eq!(actual_scores, full_scores); } #[test] - fn query_local_materialization_never_starts_background_maintenance() { + fn query_local_external_empty_scorer_retains_zero_score_membership() { let schema = create_test_schema(); let batch = create_test_batch(schema.as_ref()); let row_ids = UInt64Array::from(vec![900, 42, 777]); - let terms = HashSet::from(["hello".to_string()]); - let mut index = QueryLocalFtsIndex::try_with_params( + let terms = FxHashSet::from_iter(["hello".to_string()]); + let index = QueryLocalFtsIndex::try_with_params( 1, "description".to_string(), InvertedIndexParams::default(), ) .unwrap(); + index + .insert_with_row_ids_for_terms(&batch, &row_ids, &terms) + .unwrap(); + + let committed_scorer = MemBM25Scorer::new(0, 0, HashMap::from([("hello".to_string(), 0)])); + let query = FtsQuery::Match( + lance_index::scalar::inverted::query::MatchQuery::new("hello".to_string()) + .with_column(Some("description".to_string())), + ); + let leaves = index.exact_leaf_results(&query, &committed_scorer).unwrap(); + + let mut actual = leaves[0].clone(); + actual.sort_unstable_by_key(|(row_id, _)| *row_id); + assert_eq!(actual.len(), 2); + assert_eq!(actual[0].0, 777); + assert_eq!(actual[1].0, 900); + assert!(actual.iter().all(|(_, score)| score.to_bits() == 0)); + } + + #[test] + fn query_local_materialization_never_starts_background_maintenance() { + let schema = create_test_schema(); + let batch = create_test_batch(schema.as_ref()); + let row_ids = UInt64Array::from(vec![900, 42, 777]); + let terms = FxHashSet::from_iter(["hello".to_string()]); + let params = InvertedIndexParams::default(); + let tokenizer = params.build().unwrap(); + let mut index = QueryLocalFtsIndex::try_with_loaded_tokenizer( + 1, + "description".to_string(), + params, + tokenizer, + ) + .unwrap(); // Crossing the normal freeze threshold would create a partition and // may launch a detached tiered merge. Query-local materialization must // remain entirely in its query-owned tail instead. @@ -4592,7 +4723,7 @@ mod tests { assert!(index.inner.state.load().partitions.is_empty()); assert!(index.inner.merge.lock().unwrap().is_none()); - assert_eq!(index.doc_count(), 3); + assert_eq!(index.doc_count(), 2); let sibling = index.empty_sibling(); assert!(Arc::ptr_eq( @@ -4603,8 +4734,8 @@ mod tests { sibling .insert_with_row_ids_for_terms(&batch, &UInt64Array::from(vec![901, 43, 778]), &terms) .unwrap(); - assert_eq!(index.doc_count(), 3); - assert_eq!(sibling.doc_count(), 3); + assert_eq!(index.doc_count(), 2); + assert_eq!(sibling.doc_count(), 2); assert!(sibling.inner.state.load().partitions.is_empty()); assert!(sibling.inner.merge.lock().unwrap().is_none()); } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index ef4ea65de26..fddd98b3339 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -77,8 +77,7 @@ use lance_index::scalar::inverted::query::{ }; use lance_index::scalar::inverted::{ DOC_INDEX_COL, DOC_INDEX_FIELD, DocumentGranularity, INVERTED_INDEX_VERSION_V2, - INVERTED_INDEX_VERSION_V3, InvertedIndex, InvertedIndexParams, SCORE_COL, SCORE_FIELD, - fts_schema, + INVERTED_INDEX_VERSION_V3, SCORE_COL, SCORE_FIELD, fts_schema, }; use lance_index::scalar::registry::VALUE_COLUMN_NAME; use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, DIST_COL, Query}; @@ -100,16 +99,16 @@ use crate::dataset::overlay::{collect_overlay_stale_rows_for_segment, overlaid_f use crate::dataset::row_offsets_to_row_addresses; use crate::dataset::rowids::{live_row_addrs_to_row_ids, translate_addr_treemap_to_row_ids}; use crate::dataset::utils::SchemaAdapter; +use crate::index::DatasetIndexInternalExt; use crate::index::scalar::fetch_index_details; use crate::index::scalar::inverted::{ - fts_index_fragment_bitmap, load_segment_details, load_segment_params, load_segments, - normalize_inverted_details, resolve_fts_field, resolve_query_document_granularity, + fts_index_fragment_bitmap, load_segment_details, load_segments, normalize_inverted_details, + resolve_fts_field, resolve_query_document_granularity, }; use crate::index::scalar_logical::{load_named_scalar_segments, scalar_index_fragment_bitmap}; use crate::index::vector::utils::{ default_distance_type_for, get_vector_dim, get_vector_type, validate_distance_type_for, }; -use crate::index::{DatasetIndexInternalExt, has_append_only_indexed_field_history}; use crate::io::exec::filtered_read::{ FilteredReadExec, FilteredReadOptions, FilteredReadThreadingMode, }; @@ -285,7 +284,7 @@ fn supports_compound_scorer(query: &FtsQuery) -> bool { !columns.is_empty() && (!matches!(query, FtsQuery::MultiMatch(_)) || columns.len() == 1) } -fn supports_exact_residual_compound(query: &FtsQuery) -> bool { +fn supports_indexed_stats_residual_compound(query: &FtsQuery) -> bool { match query { FtsQuery::Match(query) => query.fuzziness == Some(0), // MemWAL phrase matching currently collapses tokenizer position gaps. @@ -293,8 +292,8 @@ fn supports_exact_residual_compound(query: &FtsQuery) -> bool { // those gaps exactly (notably when stop words are configured). FtsQuery::Phrase(_) => false, FtsQuery::Boost(query) => { - supports_exact_residual_compound(&query.positive) - && supports_exact_residual_compound(&query.negative) + supports_indexed_stats_residual_compound(&query.positive) + && supports_indexed_stats_residual_compound(&query.negative) } FtsQuery::MultiMatch(query) => query .match_queries @@ -305,11 +304,22 @@ fn supports_exact_residual_compound(query: &FtsQuery) -> bool { .iter() .chain(&query.must) .chain(&query.must_not) - .all(supports_exact_residual_compound), + .all(supports_indexed_stats_residual_compound), } } -fn has_exact_hybrid_fts_coverage( +const MAX_QUERY_LOCAL_RESIDUAL_ROWS: usize = 100_000; + +fn has_bounded_query_local_residual_rows(fragments: &[Fragment]) -> bool { + fragments + .iter() + .try_fold(0usize, |total, fragment| { + total.checked_add(fragment.physical_rows?) + }) + .is_some_and(|total| total <= MAX_QUERY_LOCAL_RESIDUAL_ROWS) +} + +fn has_complete_hybrid_fts_coverage( segments: &[IndexMetadata], residual_fragments: &[Fragment], target_fragments: &[Fragment], @@ -344,27 +354,6 @@ fn has_exact_hybrid_fts_coverage( indexed | residual == target } -fn has_compatible_hybrid_physical_segments( - params: &[InvertedIndexParams], - details: &[InvertedIndexDetails], - has_deleted_fragments: &[bool], -) -> bool { - let Some(first) = params.first() else { - return false; - }; - params.len() == details.len() - && params.len() == has_deleted_fragments.len() - && first.posting_block_size() == 128 - && params.iter().all(|params| params == first) - && details.iter().all(|details| { - matches!( - details.posting_format_version, - Some(INVERTED_INDEX_VERSION_V2 | INVERTED_INDEX_VERSION_V3) - ) - }) - && has_deleted_fragments.iter().all(|has_deleted| !has_deleted) -} - fn validate_fts_query_contract(query: &FtsQuery) -> Result<()> { fn validate_multiplier(name: &str, value: f32) -> Result<()> { if value.is_finite() && value >= 0.0 { @@ -4303,7 +4292,11 @@ impl Scanner { } let mut phrase_columns = HashSet::new(); collect_phrase_columns(query, &mut phrase_columns); - let allow_exact_residual = !cross_column + // Query-local residual scoring intentionally reuses committed-index + // BM25 statistics. Matching remains exact for the supported leaf + // shapes, but ranking is approximate until the appended rows are + // incorporated into a persistent index. + let allow_indexed_stats_residual = !cross_column && !self.fast_search && self.fragments.is_none() && filter_plan.is_empty() @@ -4313,7 +4306,7 @@ impl Scanner { && target_fragments .iter() .all(|fragment| fragment.deletion_file.is_none()) - && supports_exact_residual_compound(query); + && supports_indexed_stats_residual_compound(query); let segment_groups = futures::future::try_join_all(columns.into_iter().map(|column| { let phrase_columns = &phrase_columns; @@ -4337,9 +4330,11 @@ impl Scanner { ) .await?; let unindexed_fragments = self.retain_target_fragments(unindexed_fragments); + let has_bounded_residual = allow_indexed_stats_residual + && has_bounded_query_local_residual_rows(&unindexed_fragments); if !unindexed_fragments.is_empty() && (!self.fast_search || unindexed_fragments.len() == target_fragments.len()) - && !(allow_exact_residual + && !(has_bounded_residual && unindexed_fragments.len() < target_fragments.len()) { // Flat and posting-backed leaves do not share a document @@ -4363,8 +4358,8 @@ impl Scanner { } FtsOverlayPlan::RowLevel { .. } | FtsOverlayPlan::FullScan => return Ok(None), }; - if allow_exact_residual && !unindexed_fragments.is_empty() { - if !has_exact_hybrid_fts_coverage( + if has_bounded_residual && !unindexed_fragments.is_empty() { + if !has_complete_hybrid_fts_coverage( &segments, &unindexed_fragments, target_fragments, @@ -4377,62 +4372,8 @@ impl Scanner { )); } // Preserve the established semantic mismatch error before - // applying the narrower physical fast-path gate. + // constructing query-local postings with the same tokenizer. load_segment_details(&self.dataset, &column, &segments).await?; - if !has_append_only_indexed_field_history(&self.dataset, &segments).await { - // Logical coverage can prune a same-id field rewrite - // while the physical segment still contributes the - // obsolete document to BM25 corpus statistics. - return Ok(None); - } - let physical_details = futures::future::try_join_all( - segments.iter().map(|segment| { - load_physical_fts_details(&self.dataset, &column, segment) - }), - ) - .await?; - let segment_params = futures::future::try_join_all( - segments - .iter() - .map(|segment| load_segment_params(&self.dataset, segment)), - ) - .await?; - let has_deleted_fragments = futures::future::try_join_all( - segments.iter().map(|segment| { - let column = &column; - async move { - let index = self - .dataset - .open_scalar_index( - column, - &segment.uuid, - &NoOpMetricsCollector, - ) - .await?; - let index = index - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::internal(format!( - "hybrid compound FTS segment {} is not an inverted index", - segment.uuid - )) - })?; - Ok::<_, Error>(!index.deleted_fragments().is_empty()) - } - }), - ) - .await?; - if !has_compatible_hybrid_physical_segments( - &segment_params, - &physical_details, - &has_deleted_fragments, - ) { - // Larger posting blocks quantize document lengths, and - // retired physical documents remain in corpus stats. - // Either would make the two arms incomparable. - return Ok(None); - } } if cross_column { @@ -4482,7 +4423,7 @@ impl Scanner { segment_groups.into_iter().next().ok_or_else(|| { Error::internal("compound scorer requires one column".to_string()) })?; - if allow_exact_residual && !unindexed_fragments.is_empty() { + if allow_indexed_stats_residual && !unindexed_fragments.is_empty() { let resolved = resolve_fts_field(self.dataset.schema(), &column, document_granularity)?; let scan_column = if resolved.has_lists() { @@ -7437,6 +7378,28 @@ mod test { assert!(error.to_string().contains("BoostQuery negative_boost")); } + #[test] + fn test_query_local_residual_row_bound() { + let fragment_with_rows = |id, physical_rows| { + let mut fragment = Fragment::new(id); + fragment.physical_rows = physical_rows; + fragment + }; + + assert!(has_bounded_query_local_residual_rows(&[ + fragment_with_rows(0, Some(40_000)), + fragment_with_rows(1, Some(60_000)), + ])); + assert!(!has_bounded_query_local_residual_rows(&[ + fragment_with_rows(0, Some(40_000)), + fragment_with_rows(1, Some(60_001)), + ])); + assert!(!has_bounded_query_local_residual_rows(&[ + fragment_with_rows(0, Some(1)), + fragment_with_rows(1, None), + ])); + } + #[test] fn test_normalize_fts_zero_boosts_recurses_and_preserves_nonzero_values() { fn boost_bits(query: &FtsQuery) -> Vec { @@ -7604,53 +7567,6 @@ mod test { assert!(!supports_compound_scorer(&cross_column)); } - #[test] - fn test_hybrid_compound_requires_compatible_live_physical_segments() { - let params = InvertedIndexParams::default(); - let modern_details = InvertedIndexDetails { - posting_format_version: Some(INVERTED_INDEX_VERSION_V3), - ..Default::default() - }; - assert!(has_compatible_hybrid_physical_segments( - &[params.clone(), params.clone()], - &[modern_details.clone(), modern_details.clone()], - &[false, false] - )); - assert!(!has_compatible_hybrid_physical_segments( - &[params.clone(), params.clone()], - &[modern_details.clone(), modern_details.clone()], - &[false, true] - )); - assert!(!has_compatible_hybrid_physical_segments( - &[params.clone(), params.clone().with_position(true)], - &[modern_details.clone(), modern_details.clone()], - &[false, false] - )); - assert!(!has_compatible_hybrid_physical_segments( - &[params.clone().block_size(256).unwrap()], - std::slice::from_ref(&modern_details), - &[false] - )); - assert!(!has_compatible_hybrid_physical_segments( - std::slice::from_ref(¶ms), - &[InvertedIndexDetails { - posting_format_version: Some(1), - ..Default::default() - }], - &[false] - )); - assert!(!has_compatible_hybrid_physical_segments( - std::slice::from_ref(¶ms), - &[InvertedIndexDetails::default()], - &[false] - )); - assert!(!has_compatible_hybrid_physical_segments( - &[params], - &[modern_details], - &[] - )); - } - #[test] fn test_collect_phrase_columns_traverses_prohibited_subtrees() { let phrase = diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index fb7db678d56..f378c57ae45 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -11,8 +11,6 @@ use std::vec; use crate::dataset::ROW_ID; use crate::dataset::WriteDestination; use crate::dataset::builder::DatasetBuilder; -use crate::dataset::index::LanceIndexStoreExt; -use crate::dataset::optimize::{CompactionOptions, compact_files, remapping}; use crate::dataset::tests::dataset_migrations::scan_dataset; use crate::dataset::tests::dataset_transactions::{assert_results, execute_sql}; use crate::dataset::transaction::{DataReplacementGroup, Operation, Transaction}; @@ -23,7 +21,7 @@ use crate::{Dataset, Error, Result}; use lance_arrow::FixedSizeListArrayExt; use crate::dataset::write::{WriteMode, WriteParams}; -use crate::index::{DatasetIndexExt, DatasetIndexInternalExt}; +use crate::index::DatasetIndexExt; use arrow::array::{AsArray, GenericListBuilder, GenericStringBuilder}; use arrow::datatypes::UInt64Type; use arrow_array::RecordBatch; @@ -52,13 +50,9 @@ use lance_index::scalar::inverted::{ query::{BooleanQuery, BoostQuery, MatchQuery, Occur, Operator, PhraseQuery}, tokenizer::InvertedIndexParams, }; -use lance_index::scalar::lance_format::LanceIndexStore; -use lance_index::scalar::{ - FullTextSearchQuery, OldIndexDataFilter, ScalarIndex, index_files_to_table, -}; +use lance_index::scalar::{FullTextSearchQuery, ScalarIndex}; use lance_index::{FtsPrewarmOptions, PrewarmOptions}; use lance_index::{IndexType, scalar::ScalarIndexParams, vector::DIST_COL}; -use lance_io::object_store::ObjectStore; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_io::utils::CachedFileSize; use lance_linalg::distance::MetricType; @@ -71,9 +65,7 @@ use lance_index::scalar::inverted::query::{FtsQuery, MultiMatchQuery}; use lance_table::format::BasePath; use lance_testing::datagen::generate_random_array; use rand::Rng; -use roaring::RoaringBitmap; use rstest::rstest; -use uuid::Uuid; #[rstest] #[tokio::test] @@ -1889,26 +1881,14 @@ async fn test_top_level_cross_column_multimatch_uses_field_local_compound_scorer "fast search must skip the partially covered body's flat path:\n{fast_plan}" ); - partial_dataset - .create_index( - &["body"], - IndexType::Inverted, - Some("body_idx".to_owned()), - &InvertedIndexParams::default().with_position(true), - true, - ) - .await - .unwrap(); - let mut rebuilt_oracle = compound_fts_results(&partial_dataset, explicit_query, None).await; - assert!( - rebuilt_oracle.len() > LIMIT, - "the rebuilt exact oracle must contain candidates beyond k" + assert_eq!( + partial_results.len(), + LIMIT, + "the approximate residual path must still return a bounded top-k" ); - rebuilt_oracle.truncate(LIMIT); - assert_scored_rows_close( - "partial_top_level_cross_column_multimatch", - &partial_results, - &rebuilt_oracle, + assert!( + partial_results.iter().all(|(_, score)| score.is_finite()), + "committed-index statistics must produce finite residual scores" ); } @@ -2650,27 +2630,27 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { ]) .into(); - let mut exact_scanner = dataset.scan(); - exact_scanner + let mut hybrid_scanner = dataset.scan(); + hybrid_scanner .project(&["id"]) .unwrap() .full_text_search(FullTextSearchQuery::new_query(query.clone())) .unwrap(); - exact_scanner.limit(Some(2), None).unwrap(); - let exact_plan = exact_scanner.explain_plan(false).await.unwrap(); + hybrid_scanner.limit(Some(2), None).unwrap(); + let hybrid_plan = hybrid_scanner.explain_plan(false).await.unwrap(); assert!( - exact_plan.contains("HybridCompoundFtsScorer"), - "exact partial coverage should build one query-local residual index:\n{exact_plan}" + hybrid_plan.contains("HybridCompoundFtsScorer"), + "partial coverage should build one indexed-statistics query-local residual index:\n{hybrid_plan}" ); assert!( - !exact_plan.contains("FlatMatchQuery"), - "hybrid compound scoring must not scan the residual once per leaf:\n{exact_plan}" + !hybrid_plan.contains("FlatMatchQuery"), + "hybrid compound scoring must not scan the residual once per leaf:\n{hybrid_plan}" ); - let exact = exact_scanner.try_into_batch().await.unwrap(); + let hybrid = hybrid_scanner.try_into_batch().await.unwrap(); assert_eq!( - exact["id"].as_primitive::().values(), + hybrid["id"].as_primitive::().values(), &[0, 2], - "exact search should include the appended hit" + "approximate residual search should include the appended hit" ); let empty_terms_query: FtsQuery = BooleanQuery::new([ @@ -2765,301 +2745,24 @@ async fn test_same_column_compound_fast_search_excludes_unindexed_rows() { } #[tokio::test] -async fn test_partial_compound_hybrid_matches_rebuilt_index_scores_and_ties() { - let initial = arrow_array::record_batch!( - ("text", Utf8, ["fresh alpha", "blocked fresh alpha"]), - ("id", Int32, [0, 1]) - ) - .unwrap(); - let schema = initial.schema(); - let mut dataset = Dataset::write( - RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema), - "memory://", - None, - ) - .await - .unwrap(); - create_fragmented_fts_index(&mut dataset, "text", true).await; - - let appended = arrow_array::record_batch!( - ( - "text", - Utf8, - ["fresh alpha", "fresh beta", "fresh alpha", "fresh gamma"] - ), - ("id", Int32, [2, 3, 4, 5]) - ) - .unwrap(); - let schema = appended.schema(); - dataset - .append( - RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), - Some(WriteParams { - // Keep the residual rows in separate scan batches so the - // parallel query-local shards must merge scores and ties. - max_rows_per_file: 1, - ..Default::default() - }), - ) - .await - .unwrap(); - - let positive: FtsQuery = BooleanQuery::new([ - (Occur::Must, compound_match_query("fresh", "text", 1.0)), - (Occur::Should, compound_match_query("alpha", "text", 1.0)), - (Occur::MustNot, compound_match_query("blocked", "text", 1.0)), - ]) - .into(); - let boost_query: FtsQuery = BoostQuery::new( - positive, - compound_match_query("alpha", "text", 1.0), - Some(0.25), - ) - .into(); - let partial_boost = compound_fts_results(&dataset, boost_query.clone(), Some(10)).await; - assert_eq!( - partial_boost.len(), - 5, - "MUST_NOT must exclude the blocked row" - ); - assert_eq!(partial_boost[0].1.to_bits(), partial_boost[1].1.to_bits()); - assert!( - partial_boost[0].0 < partial_boost[1].0, - "equal-score rows must use ascending row id as the exact tie break" - ); - - let multimatch_query: FtsQuery = MultiMatchQuery::try_new( - "fresh alpha".to_string(), - vec!["text".to_string(), "text".to_string()], - ) - .unwrap() - .try_with_boosts(vec![1.0, 2.0]) - .unwrap() - .into(); - let partial_multimatch = - compound_fts_results(&dataset, multimatch_query.clone(), Some(3)).await; - - dataset - .create_index( - &["text"], - IndexType::Inverted, - Some("text_idx".to_string()), - &InvertedIndexParams::default().with_position(true), - true, - ) - .await - .unwrap(); - let rebuilt_boost = compound_fts_results(&dataset, boost_query, Some(10)).await; - let rebuilt_multimatch = compound_fts_results(&dataset, multimatch_query, Some(3)).await; - assert_eq!( - scored_row_bits(&partial_boost), - scored_row_bits(&rebuilt_boost), - "hybrid Boost scores must be bit-identical to a rebuilt index" - ); - assert_eq!( - scored_row_bits(&partial_multimatch), - scored_row_bits(&rebuilt_multimatch), - "hybrid MultiMatch scores must be bit-identical to a rebuilt index" - ); -} - -#[tokio::test] -async fn test_partial_compound_hybrid_rejects_same_id_rewrite_after_deferred_remap() { - let initial = arrow_array::record_batch!( - ( - "text", - Utf8, - [ - "stable alpha common", - "stable gamma common", - "stale alpha common" - ] - ), - ("id", Int32, [0, 1, 2]) - ) - .unwrap(); - let schema = initial.schema(); - let test_uri = TempStrDir::default(); - let mut dataset = Dataset::write( - RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema.clone()), - &test_uri, - Some(WriteParams { - max_rows_per_file: 1, - ..Default::default() - }), - ) - .await - .unwrap(); - dataset - .create_index( - &["text"], - IndexType::Inverted, - Some("text_idx".to_string()), - &InvertedIndexParams::default().with_position(true), - true, - ) - .await - .unwrap(); - - let rewritten_fragment = dataset.get_fragment(2).unwrap(); - let mut replacement_file = rewritten_fragment.metadata().files[0].clone(); - replacement_file.path = "replacement.lance".to_string(); - let replacement = arrow_array::record_batch!( - ("text", Utf8, ["replacement beta common"]), - ("id", Int32, [2]) - ) - .unwrap(); - let object_writer = dataset - .object_store - .create(&dataset.data_dir().join(replacement_file.path.as_str())) - .await - .unwrap(); - let mut writer = lance_file::versions::v2_1::create_writer( - object_writer, - schema.as_ref().try_into().unwrap(), - Default::default(), - ) - .unwrap(); - writer.write_batch(&replacement).await.unwrap(); - writer.finish().await.unwrap(); - - let read_version = dataset.manifest.version; - let mut dataset = Dataset::commit( - WriteDestination::Dataset(Arc::new(dataset)), - Operation::DataReplacement { - replacements: vec![DataReplacementGroup(2, replacement_file)], - }, - Some(read_version), - None, - None, - Arc::new(Default::default()), - false, - ) - .await - .unwrap(); - let committed = dataset - .load_index_by_name("text_idx") - .await - .unwrap() - .unwrap(); - assert!( - !committed.fragment_bitmap.as_ref().unwrap().contains(2), - "the logical index must prune the same-id rewritten fragment" - ); - let physical_source_version = committed.dataset_version; - - // Compact only the two still-covered fragments and defer the address - // remap. The stale physical document from rewritten fragment 2 is absent - // from that remap and therefore remains in the postings. - let metrics = compact_files( - &mut dataset, - CompactionOptions { - target_rows_per_fragment: 2, - defer_index_remap: true, - ..Default::default() - }, - None, - ) - .await - .unwrap(); - assert_eq!(metrics.fragments_removed, 2); - assert_eq!(metrics.fragments_added, 1); - - remapping::remap_column_index(&mut dataset, &["text"], Some("text_idx".to_string())) - .await - .unwrap(); - let remapped = dataset - .load_index_by_name("text_idx") - .await - .unwrap() - .unwrap(); - assert!(remapped.dataset_version > physical_source_version); - assert_eq!( - crate::index::scalar::inverted::physical_source_dataset_versions(&dataset, &remapped) - .await - .unwrap(), - Some(vec![physical_source_version]), - "deferred remap must not advance immutable physical provenance" - ); - - let appended = - arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [3])) - .unwrap(); - dataset - .append( - RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), - None, - ) - .await - .unwrap(); - - let query: FtsQuery = BooleanQuery::new([ - (Occur::Must, compound_match_query("common", "text", 1.0)), - (Occur::Should, compound_match_query("alpha", "text", 1.0)), - (Occur::Should, compound_match_query("beta", "text", 1.0)), - (Occur::MustNot, compound_match_query("stale", "text", 1.0)), - ]) - .into(); - let plan = compound_fts_plan(&dataset, query.clone(), 2).await; - assert!( - !plan.contains("HybridCompoundFtsScorer"), - "same-id indexed-field rewrites make physical BM25 stats unsafe:\n{plan}" - ); - - let actual = compound_fts_results(&dataset, query.clone(), Some(2)).await; - let mut flat_oracle = compound_fts_results(&dataset, query, None).await; - flat_oracle.truncate(2); - assert_eq!( - actual, flat_oracle, - "bounded fallback must preserve the flat path's ordered row ids and scores" - ); - - let exact_oracle_query: FtsQuery = BooleanQuery::new([ - (Occur::Must, compound_match_query("common", "text", 0.0)), - (Occur::Should, compound_match_query("alpha", "text", 0.0)), - (Occur::Should, compound_match_query("beta", "text", 0.0)), - (Occur::MustNot, compound_match_query("stale", "text", 0.0)), - ]) - .into(); - let exact_actual = compound_fts_results(&dataset, exact_oracle_query.clone(), Some(3)).await; - let mut rebuilt = dataset.clone(); - rebuilt - .create_index( - &["text"], - IndexType::Inverted, - Some("text_idx".to_string()), - &InvertedIndexParams::default().with_position(true), - true, - ) - .await - .unwrap(); - let exact_expected = compound_fts_results(&rebuilt, exact_oracle_query, Some(3)).await; - assert_eq!( - scored_row_bits(&exact_actual), - scored_row_bits(&exact_expected), - "same-id rewrite fallback must match a rebuilt index exactly" - ); -} - -#[tokio::test] -async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() { +async fn test_partial_compound_hybrid_prunes_same_path_different_base_rewrite() { let primary = TempStrDir::default(); let base_one = TempStrDir::default(); let base_two = TempStrDir::default(); let initial = arrow_array::record_batch!( - ("text", Utf8, ["stable alpha common", "stale alpha common"]), + ("text", Utf8, ["stable alpha", "stale alpha"]), ("id", Int32, [0, 1]) ) .unwrap(); let schema = initial.schema(); - let mut dataset = Dataset::write( + let dataset = Dataset::write( RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema.clone()), &primary, Some(WriteParams { max_rows_per_file: 1, initial_bases: Some(vec![ - BasePath::new(1, base_one.to_string(), Some("base-one".to_string()), false), - BasePath::new(2, base_two.to_string(), Some("base-two".to_string()), false), + BasePath::new(1, base_one.to_string(), None, false), + BasePath::new(2, base_two.to_string(), None, false), ]), target_bases: Some(vec![1]), ..Default::default() @@ -3067,30 +2770,29 @@ async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() ) .await .unwrap(); - assert!(dataset.get_fragments().iter().all(|fragment| { - fragment - .metadata() - .files + assert_eq!(dataset.get_fragments().len(), 2); + assert!( + dataset + .get_fragments() .iter() - .all(|file| file.base_id == Some(1)) - })); - - let columns = ["text"]; - let params = InvertedIndexParams::default().with_position(true); + .all(|fragment| { fragment.metadata().files[0].base_id == Some(1) }) + ); let segment = dataset - .create_index_builder(&columns, IndexType::Inverted, ¶ms) + .create_index_builder( + &["text"], + IndexType::Inverted, + &InvertedIndexParams::default().with_position(true), + ) .name("text_idx".to_string()) .execute_uncommitted() .await .unwrap(); - let rewritten_fragment = dataset.get_fragment(1).unwrap(); - let relative_path = rewritten_fragment.metadata().files[0].path.clone(); - let replacement = arrow_array::record_batch!( - ("text", Utf8, ["replacement beta common"]), - ("id", Int32, [1]) - ) - .unwrap(); + let relative_path = dataset.get_fragment(1).unwrap().metadata().files[0] + .path + .clone(); + let replacement = + arrow_array::record_batch!(("text", Utf8, ["current beta"]), ("id", Int32, [1])).unwrap(); let replacement_path = dataset .data_file_dir_for_base(Some(2)) .unwrap() @@ -3131,11 +2833,6 @@ async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() ) .await .unwrap(); - let rewritten_fragment = dataset.get_fragment(1).unwrap(); - let rewritten_file = &rewritten_fragment.metadata().files[0]; - assert_eq!(rewritten_file.path, relative_path); - assert_eq!(rewritten_file.base_id, Some(2)); - dataset .commit_existing_index_segments("text_idx", "text", vec![segment]) .await @@ -3146,381 +2843,203 @@ async fn test_partial_compound_hybrid_rejects_same_path_different_base_rewrite() .unwrap() .unwrap(); let coverage = committed.fragment_bitmap.as_ref().unwrap(); - assert!(coverage.contains(0)); - assert!( - !coverage.contains(1), - "changing only the registered base must prune stale logical coverage" - ); - - let appended = - arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [2])) - .unwrap(); - dataset - .append( - RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), - None, - ) - .await - .unwrap(); - - let query: FtsQuery = BooleanQuery::new([ - (Occur::Must, compound_match_query("common", "text", 0.0)), - (Occur::Should, compound_match_query("alpha", "text", 0.0)), - (Occur::Should, compound_match_query("beta", "text", 0.0)), - (Occur::MustNot, compound_match_query("stale", "text", 0.0)), - ]) - .into(); - let plan = compound_fts_plan(&dataset, query.clone(), 3).await; assert!( - !plan.contains("HybridCompoundFtsScorer"), - "same-path files from different bases make physical BM25 stats unsafe:\n{plan}" - ); - let actual = compound_fts_results(&dataset, query.clone(), Some(3)).await; - - let mut rebuilt = dataset.clone(); - rebuilt - .create_index( - &["text"], - IndexType::Inverted, - Some("text_idx".to_string()), - ¶ms, - true, - ) - .await - .unwrap(); - let expected = compound_fts_results(&rebuilt, query, Some(3)).await; - assert_eq!( - scored_row_bits(&actual), - scored_row_bits(&expected), - "different-base rewrite fallback must match a rebuilt index exactly" - ); -} - -#[tokio::test] -async fn test_partial_compound_hybrid_rejects_rebound_registered_base() { - let primary = TempStrDir::default(); - let base_a = TempStrDir::default(); - let base_b = TempStrDir::default(); - let stable = - arrow_array::record_batch!(("text", Utf8, ["stable alpha common"]), ("id", Int32, [0])) - .unwrap(); - let schema = stable.schema(); - let dataset = Dataset::write( - RecordBatchIterator::new(vec![stable].into_iter().map(Ok), schema.clone()), - &primary, - Some(WriteParams { - max_rows_per_file: 1, - initial_bases: Some(vec![BasePath::new( - 1, - base_a.to_string(), - Some("base-a".to_string()), - false, - )]), - ..Default::default() - }), - ) - .await - .unwrap(); - let indexed_on_base = - arrow_array::record_batch!(("text", Utf8, ["stale alpha common"]), ("id", Int32, [1])) - .unwrap(); - let mut dataset = Dataset::write( - RecordBatchIterator::new(vec![indexed_on_base].into_iter().map(Ok), schema.clone()), - Arc::new(dataset), - Some(WriteParams { - mode: WriteMode::Append, - target_bases: Some(vec![1]), - ..Default::default() - }), - ) - .await - .unwrap(); - assert_eq!( - dataset.get_fragment(0).unwrap().metadata().files[0].base_id, - None + coverage.contains(0), + "the unchanged physical file must remain covered" ); - assert_eq!( - dataset.get_fragment(1).unwrap().metadata().files[0].base_id, - Some(1) - ); - - let columns = ["text"]; - let params = InvertedIndexParams::default().with_position(true); - let segment = dataset - .create_index_builder(&columns, IndexType::Inverted, ¶ms) - .name("text_idx".to_string()) - .execute_uncommitted() - .await - .unwrap(); - - let fragment_before = dataset.get_fragment(1).unwrap().metadata().clone(); - let relative_path = fragment_before.files[0].path.clone(); - let replacement = arrow_array::record_batch!( - ("text", Utf8, ["replacement beta common"]), - ("id", Int32, [1]) - ) - .unwrap(); - let (base_b_store, base_b_root) = ObjectStore::from_uri(&base_b).await.unwrap(); - let object_writer = base_b_store - .create(&base_b_root.join(relative_path.as_str())) - .await - .unwrap(); - let mut writer = lance_file::versions::v2_1::create_writer( - object_writer, - schema.as_ref().try_into().unwrap(), - Default::default(), - ) - .unwrap(); - writer.write_batch(&replacement).await.unwrap(); - writer.finish().await.unwrap(); - - let binding_before = dataset.manifest.base_paths.get(&1).unwrap().clone(); - dataset = Arc::new(dataset) - .add_bases( - vec![BasePath::new( - 1, - base_b.to_string(), - Some("base-b".to_string()), - false, - )], - None, - ) - .await - .unwrap(); - let binding_after = dataset.manifest.base_paths.get(&1).unwrap(); - assert_eq!(binding_before.path, base_a.to_string()); - assert_eq!(binding_after.path, base_b.to_string()); - assert_eq!( - binding_before.is_dataset_root, - binding_after.is_dataset_root - ); - assert_eq!( - dataset.get_fragment(1).unwrap().metadata(), - &fragment_before, - "UpdateBases must leave the DataFile identity unchanged" - ); - - dataset - .commit_existing_index_segments("text_idx", "text", vec![segment]) - .await - .unwrap(); - let committed = dataset - .load_index_by_name("text_idx") - .await - .unwrap() - .unwrap(); - let coverage = committed.fragment_bitmap.as_ref().unwrap(); - assert!(coverage.contains(0)); assert!( !coverage.contains(1), - "rebinding a referenced base must prune stale logical coverage" + "the same path on a different registered base must be pruned" ); - let appended = - arrow_array::record_batch!(("text", Utf8, ["tail beta common"]), ("id", Int32, [2])) - .unwrap(); - dataset - .append( - RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), - None, - ) - .await - .unwrap(); let query: FtsQuery = BooleanQuery::new([ - (Occur::Must, compound_match_query("common", "text", 0.0)), - (Occur::Should, compound_match_query("alpha", "text", 0.0)), - (Occur::Should, compound_match_query("beta", "text", 0.0)), - (Occur::MustNot, compound_match_query("stale", "text", 0.0)), + (Occur::Must, compound_match_query("beta", "text", 1.0)), + (Occur::MustNot, compound_match_query("alpha", "text", 1.0)), ]) .into(); - let plan = compound_fts_plan(&dataset, query.clone(), 3).await; + let mut scanner = dataset.scan(); + scanner + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query)) + .unwrap(); + scanner.limit(Some(2), None).unwrap(); + let plan = scanner.explain_plan(false).await.unwrap(); assert!( - !plan.contains("HybridCompoundFtsScorer"), - "a rebound registered base makes physical BM25 stats unsafe:\n{plan}" + plan.contains("HybridCompoundFtsScorer"), + "the physically pruned fragment should use hybrid residual scoring:\n{plan}" ); - let actual = compound_fts_results(&dataset, query.clone(), Some(3)).await; - - let mut rebuilt = dataset.clone(); - rebuilt - .create_index( - &["text"], - IndexType::Inverted, - Some("text_idx".to_string()), - ¶ms, - true, - ) - .await - .unwrap(); - let expected = compound_fts_results(&rebuilt, query, Some(3)).await; + let results = scanner.try_into_batch().await.unwrap(); assert_eq!( - scored_row_bits(&actual), - scored_row_bits(&expected), - "rebound-base fallback must match a rebuilt index exactly" + results["id"].as_primitive::().values(), + &[1], + "the current beta row must be visible without leaking stale alpha membership" + ); + assert!( + results[SCORE_COL] + .as_primitive::() + .values() + .iter() + .all(|score| score.is_finite()) ); } #[tokio::test] -async fn test_partial_compound_hybrid_rejects_retired_physical_fragments() { +async fn test_partial_compound_hybrid_uses_committed_index_statistics() { let initial = arrow_array::record_batch!( - ("text", Utf8, ["retired alpha", "live alpha"]), + ("text", Utf8, ["fresh alpha", "blocked fresh alpha"]), ("id", Int32, [0, 1]) ) .unwrap(); let schema = initial.schema(); - let test_uri = TempStrDir::default(); let mut dataset = Dataset::write( RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema), - &test_uri, + "memory://", None, ) .await .unwrap(); create_fragmented_fts_index(&mut dataset, "text", true).await; - let initial_segment = dataset - .load_index_by_name("text_idx") - .await - .unwrap() - .unwrap(); - let retired_fragments = initial_segment.fragment_bitmap.clone().unwrap(); - let initial_index = dataset - .open_scalar_index( + let appended = arrow_array::record_batch!( + ( "text", - &initial_segment.uuid, - &lance_index::metrics::NoOpMetricsCollector, - ) - .await - .unwrap(); - - dataset.delete("id = 0").await.unwrap(); - let metrics = compact_files( - &mut dataset, - CompactionOptions { - target_rows_per_fragment: 10, - materialize_deletions_threshold: 0.0, - ..Default::default() - }, - None, - ) - .await - .unwrap(); - assert!(metrics.fragments_removed > 0); - assert!( - dataset - .get_fragments() - .iter() - .all(|fragment| fragment.metadata().deletion_file.is_none()) - ); - let indexed_fragments = dataset.fragment_bitmap.as_ref().clone(); - assert!(indexed_fragments.is_disjoint(&retired_fragments)); - - // Model an incremental FTS replacement that retains the old postings and - // records their now-retired fragment ids for merge-on-read filtering. - let resolved = crate::index::scalar::inverted::resolve_fts_field_by_id( - dataset.schema(), - initial_segment.fields[0], - DocumentGranularity::Row, - ) - .unwrap(); - let current_fragments = dataset - .get_fragments() - .iter() - .map(|fragment| fragment.metadata().clone()) - .collect(); - let update_criteria = initial_index.update_criteria(); - let new_data = crate::index::scalar::load_fts_training_data( - &dataset, - &resolved, - &update_criteria.data_criteria, - Some(current_fragments), - true, - None, + Utf8, + [ + "fresh alpha", + "fresh beta", + "fresh alpha", + "fresh gamma", + "fresh beta blocked" + ] + ), + ("id", Int32, [2, 3, 4, 5, 6]) ) - .await .unwrap(); - let updated_uuid = Uuid::new_v4(); - let updated_store = LanceIndexStore::from_dataset_for_new(&dataset, &updated_uuid).unwrap(); - let created = initial_index - .update( - new_data, - &updated_store, - Some(OldIndexDataFilter::Fragments { - to_keep: RoaringBitmap::new(), - to_remove: retired_fragments.clone(), - }), - ) - .await - .unwrap(); - let updated_segment = lance_table::format::IndexMetadata { - uuid: updated_uuid, - dataset_version: dataset.manifest.version, - fragment_bitmap: Some(indexed_fragments), - index_details: Some(Arc::new(created.index_details)), - index_version: created.index_version as i32, - created_at: Some(chrono::Utc::now()), - base_id: None, - files: Some(index_files_to_table(created.files)), - ..initial_segment - }; - dataset - .commit_existing_index_segments("text_idx", "text", vec![updated_segment]) - .await - .unwrap(); - - let committed = dataset - .load_index_by_name("text_idx") - .await - .unwrap() - .unwrap(); - let committed_index = dataset - .open_scalar_index( - "text", - &committed.uuid, - &lance_index::metrics::NoOpMetricsCollector, - ) - .await - .unwrap(); - let committed_index = committed_index - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(committed_index.deleted_fragments(), &retired_fragments); - - let appended = - arrow_array::record_batch!(("text", Utf8, ["tail alpha"]), ("id", Int32, [2])).unwrap(); let schema = appended.schema(); dataset .append( RecordBatchIterator::new(vec![appended].into_iter().map(Ok), schema), - None, + Some(WriteParams { + // Keep the residual rows in separate fragments; execution may + // rechunk their scan batches before query-local indexing. + max_rows_per_file: 1, + ..Default::default() + }), ) .await .unwrap(); - let query: FtsQuery = BooleanQuery::new([ - (Occur::Should, compound_match_query("retired", "text", 1.0)), - (Occur::Should, compound_match_query("tail", "text", 1.0)), + let positive: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("fresh", "text", 1.0)), + (Occur::Should, compound_match_query("alpha", "text", 1.0)), + (Occur::MustNot, compound_match_query("blocked", "text", 1.0)), ]) .into(); - let plan = compound_fts_plan(&dataset, query.clone(), 10).await; + let boost_query: FtsQuery = BoostQuery::new( + positive, + compound_match_query("alpha", "text", 1.0), + Some(0.25), + ) + .into(); + let partial_boost = compound_fts_results(&dataset, boost_query.clone(), Some(10)).await; + assert_eq!( + partial_boost.len(), + 5, + "MUST_NOT must exclude the blocked row" + ); + assert_eq!(partial_boost[0].1.to_bits(), partial_boost[1].1.to_bits()); assert!( - !plan.contains("HybridCompoundFtsScorer"), - "retired physical docs make unified hybrid BM25 stats unsafe:\n{plan}" + partial_boost[0].0 < partial_boost[1].0, + "equal-score rows must use ascending row id as the exact tie break" ); + let multimatch_query: FtsQuery = MultiMatchQuery::try_new( + "fresh alpha".to_string(), + vec!["text".to_string(), "text".to_string()], + ) + .unwrap() + .try_with_boosts(vec![1.0, 2.0]) + .unwrap() + .into(); + for (query_name, query) in [ + ("Boost", boost_query.clone()), + ("MultiMatch", multimatch_query.clone()), + ] { + let mut scanner = dataset.scan(); + scanner + .project(&["id"]) + .unwrap() + .full_text_search(FullTextSearchQuery::new_query(query)) + .unwrap(); + scanner.limit(Some(10), None).unwrap(); + let batch = scanner.try_into_batch().await.unwrap(); + let ids = batch["id"].as_primitive::().values(); + let scores = batch[SCORE_COL].as_primitive::().values(); + let score_bits = ids + .iter() + .copied() + .zip(scores.iter().map(|score| score.to_bits())) + .collect::>(); + for residual_id in [2, 4] { + assert_eq!( + score_bits.get(&residual_id), + score_bits.get(&0), + "{query_name} must score identical indexed and residual documents identically" + ); + } + } - let mut scanner = dataset.scan(); - scanner - .project(&["id"]) + let mut indexed_only_scanner = dataset.scan(); + indexed_only_scanner + .with_row_id() + .full_text_search(FullTextSearchQuery::new_query(boost_query.clone())) .unwrap() - .full_text_search(FullTextSearchQuery::new_query(query)) + .fast_search(); + indexed_only_scanner.limit(Some(10), None).unwrap(); + let indexed_only = + compound_fts_result_bits(&indexed_only_scanner.try_into_batch().await.unwrap()) + .into_iter() + .collect::>(); + let partial_boost_bits = scored_row_bits(&partial_boost) + .into_iter() + .collect::>(); + for (row_id, score) in indexed_only { + assert_eq!( + partial_boost_bits.get(&row_id), + Some(&score), + "hybrid scoring must preserve committed-index scores for indexed row {row_id}" + ); + } + + let residual_only_query: FtsQuery = BooleanQuery::new([ + (Occur::Must, compound_match_query("beta", "text", 1.0)), + (Occur::MustNot, compound_match_query("blocked", "text", 1.0)), + ]) + .into(); + let mut residual_only_scanner = dataset.scan(); + residual_only_scanner + .project(&["id"]) .unwrap() - .limit(Some(10), None) + .full_text_search(FullTextSearchQuery::new_query(residual_only_query)) .unwrap(); - let results = scanner.try_into_batch().await.unwrap(); + residual_only_scanner.limit(Some(10), None).unwrap(); + let residual_only_plan = residual_only_scanner.explain_plan(false).await.unwrap(); + assert!( + residual_only_plan.contains("HybridCompoundFtsScorer"), + "residual-only term membership must use the indexed-statistics hybrid path:\n{residual_only_plan}" + ); + let residual_only = residual_only_scanner.try_into_batch().await.unwrap(); assert_eq!( - results["id"].as_primitive::().values(), - &[2], - "fallback must prune stale postings from the retired fragment" + residual_only["id"].as_primitive::().values(), + &[3], + "the zero-weight beta leaf must retain membership while MUST_NOT excludes id=6" + ); + assert_eq!( + residual_only[SCORE_COL] + .as_primitive::() + .values(), + &[0.0], + "a residual-only term must contribute zero when committed df is zero" ); } @@ -4332,16 +3851,6 @@ async fn test_fts_v1_remains_queryable_after_append_optimize() { let schema = batch.schema(); let batches = RecordBatchIterator::new(vec![batch].into_iter().map(Ok), schema); dataset.append(batches, None).await.unwrap(); - let compound_query: FtsQuery = BooleanQuery::new([ - (Occur::Must, compound_match_query("alpha", "text", 1.0)), - (Occur::Should, compound_match_query("original", "text", 1.0)), - ]) - .into(); - let plan = compound_fts_plan(&dataset, compound_query, 2).await; - assert!( - !plan.contains("HybridCompoundFtsScorer"), - "a physical FTS v1 segment must not enter the modern hybrid path:\n{plan}" - ); dataset .optimize_indices(&OptimizeOptions::append()) .await diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index ab8835d5bdb..ac57366c46b 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -217,9 +217,9 @@ fn fragment_field_files<'a>( /// segment's carried columns can go stale independently of its keyed column, so /// checking only the keyed subtree would leave a fragment covered after a carried /// column was rewritten, and the segment would answer with the obsolete value. -fn indexed_field_ids(dataset: &Dataset, fields: &[i32]) -> Result> { +fn segment_indexed_field_ids(dataset: &Dataset, segment: &IndexSegment) -> Result> { let mut indexed_field_ids = HashSet::new(); - for field_id in fields { + for field_id in segment.fields() { let field = dataset.schema().field_by_id(*field_id).ok_or_else(|| { Error::invalid_input(format!( "CreateIndex: field id {field_id} does not exist in the current schema" @@ -230,106 +230,6 @@ fn indexed_field_ids(dataset: &Dataset, fields: &[i32]) -> Result> Ok(indexed_field_ids) } -fn segment_indexed_field_ids(dataset: &Dataset, segment: &IndexSegment) -> Result> { - indexed_field_ids(dataset, segment.fields()) -} - -/// Prove that every indexed-field file present at a segment's physical source -/// version is still the current file for the same fragment id. -/// -/// A committed segment's logical bitmap may already have pruned a rewritten -/// fragment even though its physical postings and corpus statistics still -/// contain that fragment. Since the original per-segment coverage is no longer -/// available after that pruning, this check deliberately considers every -/// fragment present at each physical source version. That is conservative, but -/// it makes the hybrid scorer available only when the history can be proven to -/// be a pure append with respect to the indexed field subtrees. -pub(crate) async fn has_append_only_indexed_field_history( - dataset: &Dataset, - segments: &[IndexMetadata], -) -> bool { - let current_version = dataset.manifest.version; - let current_fragments = dataset - .fragments() - .iter() - .filter_map(|fragment| u32::try_from(fragment.id).ok().map(|id| (id, fragment))) - .collect::>(); - if current_fragments.len() != dataset.fragments().len() { - return false; - } - - let mut physical_sources = Vec::with_capacity(segments.len()); - for segment in segments { - let Ok(Some(source_versions)) = - scalar::inverted::physical_source_dataset_versions(dataset, segment).await - else { - return false; - }; - // `dataset_version` is an independent mutable address-remap watermark. - // A merged segment can legitimately contain physical sources newer - // than its oldest address watermark, but neither may be in the future. - if source_versions - .iter() - .any(|source_version| *source_version > current_version) - || segment.dataset_version > current_version - { - return false; - } - physical_sources.push((segment, source_versions)); - } - let build_versions = physical_sources - .iter() - .flat_map(|(_, source_versions)| source_versions.iter().copied()) - .collect::>(); - for build_version in build_versions { - if build_version > current_version { - return false; - } - if build_version == current_version { - continue; - } - let Ok(historical) = dataset.checkout_version(build_version).await else { - return false; - }; - let mut indexed_field_ids_at_version = HashSet::new(); - for (segment, _) in physical_sources - .iter() - .filter(|(_, source_versions)| source_versions.contains(&build_version)) - { - let Ok(historical_field_ids) = indexed_field_ids(&historical, &segment.fields) else { - return false; - }; - let Ok(current_field_ids) = indexed_field_ids(dataset, &segment.fields) else { - return false; - }; - if historical_field_ids != current_field_ids { - return false; - } - indexed_field_ids_at_version.extend(historical_field_ids); - } - - for historical_fragment in historical.fragments().iter() { - let Ok(fragment_id) = u32::try_from(historical_fragment.id) else { - return false; - }; - let Some(current_fragment) = current_fragments.get(&fragment_id) else { - return false; - }; - let historical_files = fragment_field_files( - &historical, - historical_fragment, - &indexed_field_ids_at_version, - ); - let current_files = - fragment_field_files(dataset, current_fragment, &indexed_field_ids_at_version); - if historical_files.is_none() || historical_files != current_files { - return false; - } - } - } - true -} - async fn prune_stale_segment_coverage( dataset: &Dataset, segments: &mut [IndexSegment], @@ -1339,8 +1239,7 @@ pub(crate) async fn remap_index( .as_any() .downcast_ref::() .ok_or(Error::index("expected inverted index".to_string()))?; - let is_legacy = inverted_index.is_legacy(); - let mut created_index = if is_legacy { + if inverted_index.is_legacy() { log::warn!( "reindex because of legacy format, index_type: {}, index_id: {}, field: {}", scalar_index.index_type(), @@ -1366,24 +1265,7 @@ pub(crate) async fn remap_index( .await? } else { scalar_index.remap(row_id_map, &new_store).await? - }; - let source_versions = if is_legacy { - // Legacy remapping performs a full rebuild from the - // current dataset instead of retaining old postings. - Some(vec![dataset.manifest.version]) - } else { - scalar::inverted::physical_source_dataset_versions(dataset, matched).await? - }; - if let Some(source_versions) = source_versions { - created_index.files.push( - scalar::inverted::write_physical_source_dataset_versions( - &new_store, - source_versions, - ) - .await?, - ); } - created_index } _ => scalar_index.remap(row_id_map, &new_store).await?, } diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 148e7a06dc5..34f52c58718 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -1203,7 +1203,7 @@ pub async fn merge_indices_with_unindexed_frags<'a>( ) .await?; let new_uuid = Uuid::new_v4(); - let mut created_index = super::scalar::build_scalar_index( + let created_index = super::scalar::build_scalar_index( dataset.as_ref(), &resolved.canonical_path, new_uuid, @@ -1214,14 +1214,6 @@ pub async fn merge_indices_with_unindexed_frags<'a>( Arc::new(NoopIndexBuildProgress), ) .await?; - let new_store = LanceIndexStore::from_dataset_for_new(&dataset, &new_uuid)?; - created_index.files.push( - crate::index::scalar::inverted::write_physical_source_dataset_versions( - &new_store, - [dataset.manifest.version], - ) - .await?, - ); return Ok(Some(IndexMergeResults { new_uuid, removed_indices: old_indices.to_vec(), @@ -1291,66 +1283,38 @@ pub async fn merge_indices_with_unindexed_frags<'a>( let new_uuid = Uuid::new_v4(); let new_store = LanceIndexStore::from_dataset_for_new(&dataset, &new_uuid)?; - let new_dataset_version = if selected_indices.is_empty() { - dataset.manifest.version - } else { - selected_old_indices - .iter() - .map(|index| index.dataset_version) - .min() - .unwrap_or(dataset.manifest.version) - }; - let source_versions = if selected_indices.is_empty() { - Some(vec![dataset.manifest.version]) - } else { - let mut source_sets = Vec::with_capacity(selected_old_indices.len() + 1); - for segment in &selected_old_indices { - source_sets.push( - crate::index::scalar::inverted::physical_source_dataset_versions( - dataset.as_ref(), - segment, - ) - .await?, - ); - } - if !unindexed.is_empty() { - source_sets.push(Some(vec![dataset.manifest.version])); - } - crate::index::scalar::inverted::merge_physical_source_dataset_versions( - source_sets, - ) - }; - let mut created_index = if selected_indices.is_empty() { - super::scalar::build_scalar_index( - dataset.as_ref(), - &resolved.canonical_path, - new_uuid, - &reference_index.derive_index_params()?, - true, - None, - Some(new_data_stream), - Arc::new(NoopIndexBuildProgress), + let (created_index, new_dataset_version) = if selected_indices.is_empty() { + ( + super::scalar::build_scalar_index( + dataset.as_ref(), + &resolved.canonical_path, + new_uuid, + &reference_index.derive_index_params()?, + true, + None, + Some(new_data_stream), + Arc::new(NoopIndexBuildProgress), + ) + .await?, + dataset.manifest.version, ) - .await? } else { - InvertedIndex::merge_segments( - &selected_indices, - new_data_stream, - &new_store, - old_data_filter, - options.progress.clone(), - ) - .await? - }; - if let Some(source_versions) = source_versions { - created_index.files.push( - crate::index::scalar::inverted::write_physical_source_dataset_versions( + ( + InvertedIndex::merge_segments( + &selected_indices, + new_data_stream, &new_store, - source_versions, + old_data_filter, + options.progress.clone(), ) .await?, - ); - } + selected_old_indices + .iter() + .map(|index| index.dataset_version) + .min() + .unwrap_or(dataset.manifest.version), + ) + }; Ok(( new_uuid, diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index f52da5ee441..8168cd1eb2d 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -5,7 +5,6 @@ use crate::{ Error, Result, dataset::{ Dataset, - index::LanceIndexStoreExt, transaction::{Operation, TransactionBuilder}, }, index::{ @@ -331,7 +330,7 @@ impl<'a> CreateIndexBuilder<'a> { let index_id = self.index_uuid.unwrap_or_else(Uuid::new_v4); let mut output_index_uuid = index_id; - let mut created_index = match (self.index_type, self.params.index_name()) { + let created_index = match (self.index_type, self.params.index_name()) { ( IndexType::Bitmap | IndexType::BTree @@ -597,23 +596,6 @@ impl<'a> CreateIndexBuilder<'a> { ))); } }; - if created_index - .index_details - .type_url - .ends_with("InvertedIndexDetails") - { - let store = lance_index::scalar::lance_format::LanceIndexStore::from_dataset_for_new( - self.dataset, - &output_index_uuid, - )?; - created_index.files.push( - crate::index::scalar::inverted::write_physical_source_dataset_versions( - &store, - [self.dataset.manifest.version], - ) - .await?, - ); - } Ok(IndexMetadata { uuid: output_index_uuid, diff --git a/rust/lance/src/index/scalar/inverted.rs b/rust/lance/src/index/scalar/inverted.rs index 1483ca31ac7..c7c521050b4 100644 --- a/rust/lance/src/index/scalar/inverted.rs +++ b/rust/lance/src/index/scalar/inverted.rs @@ -3,10 +3,7 @@ #![allow(clippy::redundant_pub_crate)] -use std::{ - collections::{BTreeMap, BTreeSet}, - sync::Arc, -}; +use std::{collections::BTreeMap, sync::Arc}; use arrow_array::cast::AsArray; use arrow_array::{ @@ -24,12 +21,12 @@ use lance_core::{ }; use lance_index::metrics::NoOpMetricsCollector; use lance_index::pbold::InvertedIndexDetails; +use lance_index::scalar::index_files_to_table; use lance_index::scalar::inverted::{ DocumentGranularity, InvertedIndex, InvertedIndexParams, doc_index_storage_column, }; use lance_index::scalar::lance_format::LanceIndexStore; use lance_index::scalar::registry::VALUE_COLUMN_NAME; -use lance_index::scalar::{IndexFile, IndexStore, index_files_to_table}; use lance_table::format::IndexMetadata; use prost::Message; use roaring::RoaringBitmap; @@ -41,9 +38,6 @@ use crate::{ index::{DatasetIndexExt, scalar::fetch_index_details}, }; -const PHYSICAL_SOURCE_VERSIONS_FILE: &str = "physical_source_versions.lance"; -const PHYSICAL_SOURCE_VERSION_COLUMN: &str = "dataset_version"; - #[derive(Debug, Clone)] enum FtsTraversal { Text, @@ -813,11 +807,6 @@ pub(crate) async fn merge_segments( let document_granularity = DocumentGranularity::try_from(details.document_granularity)?; let resolved = resolve_fts_field_by_id(dataset.schema(), field_id, document_granularity)?; load_segment_details(dataset, &resolved.canonical_path, &segments).await?; - let mut source_sets = Vec::with_capacity(segments.len()); - for segment in &segments { - source_sets.push(physical_source_dataset_versions(dataset, segment).await?); - } - let merged_source_versions = merge_physical_source_dataset_versions(source_sets); let mut source_indices = Vec::with_capacity(segments.len()); let mut fragment_bitmap = RoaringBitmap::new(); @@ -857,7 +846,7 @@ pub(crate) async fn merge_segments( let new_uuid = Uuid::new_v4(); let new_store = LanceIndexStore::from_dataset_for_new(dataset, &new_uuid)?; - let mut created_index = InvertedIndex::merge_segments( + let created_index = InvertedIndex::merge_segments( &source_indices, empty_inverted_update_stream(dataset, &resolved)?, &new_store, @@ -865,11 +854,6 @@ pub(crate) async fn merge_segments( lance_index::progress::noop_progress(), ) .await?; - if let Some(source_versions) = merged_source_versions { - created_index - .files - .push(write_physical_source_dataset_versions(&new_store, source_versions).await?); - } Ok(IndexMetadata { uuid: new_uuid, @@ -952,134 +936,6 @@ pub(crate) async fn fts_index_fragment_bitmap( Ok(fragment_bitmap) } -/// Read the dataset versions whose physical documents contribute to an -/// inverted segment. -/// -/// Missing sidecars represent legacy segments with unknown provenance. A -/// malformed sidecar returns an error so corruption is never mistaken for -/// missing provenance. -pub(crate) async fn physical_source_dataset_versions( - dataset: &Dataset, - segment: &IndexMetadata, -) -> Result>> { - match &segment.files { - Some(files) => { - if !files - .iter() - .any(|file| file.path == PHYSICAL_SOURCE_VERSIONS_FILE) - { - return Ok(None); - } - } - None => { - let sidecar_path = dataset - .indice_files_dir(segment)? - .join(segment.uuid.to_string()) - .join(PHYSICAL_SOURCE_VERSIONS_FILE); - let object_store = dataset.object_store_for_index(segment).await?; - if !object_store.exists(&sidecar_path).await? { - return Ok(None); - } - } - } - - let store = LanceIndexStore::from_dataset_for_existing(dataset, segment).await?; - let reader = store.open_index_file(PHYSICAL_SOURCE_VERSIONS_FILE).await?; - let num_rows = reader.num_rows(); - if num_rows == 0 { - return Err(Error::io(format!( - "physical provenance sidecar for index {} is empty", - segment.uuid - ))); - } - let batch = reader.read_range(0..num_rows, None).await?; - let schema = batch.schema(); - if schema.fields().len() != 1 - || schema.field(0).name() != PHYSICAL_SOURCE_VERSION_COLUMN - || schema.field(0).data_type() != &DataType::UInt64 - || schema.field(0).is_nullable() - { - return Err(Error::io(format!( - "physical provenance sidecar for index {} has invalid schema {:?}", - segment.uuid, schema - ))); - } - if batch.num_rows() != num_rows { - return Err(Error::io(format!( - "physical provenance sidecar for index {} declared {} rows but read {}", - segment.uuid, - num_rows, - batch.num_rows() - ))); - } - let versions = batch - .column(0) - .as_any() - .downcast_ref::() - .ok_or_else(|| { - Error::io(format!( - "physical provenance sidecar for index {} has a non-UInt64 column", - segment.uuid - )) - })?; - if versions.null_count() != 0 { - return Err(Error::io(format!( - "physical provenance sidecar for index {} contains null versions", - segment.uuid - ))); - } - let mut versions = versions.values().to_vec(); - versions.sort_unstable(); - versions.dedup(); - Ok(Some(versions)) -} - -/// Union known physical provenance sets. If any input is unknown, the merged -/// corpus is also unknown because its retained postings cannot be attributed -/// exactly. -pub(crate) fn merge_physical_source_dataset_versions( - source_sets: impl IntoIterator>>, -) -> Option> { - let mut merged = BTreeSet::new(); - let mut has_input = false; - for source_versions in source_sets { - has_input = true; - let source_versions = source_versions?; - if source_versions.is_empty() { - return None; - } - merged.extend(source_versions); - } - (has_input && !merged.is_empty()).then(|| merged.into_iter().collect()) -} - -/// Write immutable physical provenance as an index-local Lance sidecar. -pub(crate) async fn write_physical_source_dataset_versions( - store: &dyn IndexStore, - dataset_versions: impl IntoIterator, -) -> Result { - let dataset_versions = dataset_versions.into_iter().collect::>(); - if dataset_versions.is_empty() { - return Err(Error::invalid_input( - "physical source dataset versions must not be empty".to_string(), - )); - } - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - PHYSICAL_SOURCE_VERSION_COLUMN, - DataType::UInt64, - false, - )])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![Arc::new(UInt64Array::from_iter_values(dataset_versions))], - )?; - let mut writer = store - .new_index_file(PHYSICAL_SOURCE_VERSIONS_FILE, schema) - .await?; - writer.write_record_batch(batch).await?; - writer.finish().await -} - /// Load and validate the shared [`InvertedIndexDetails`] across committed /// segments returned by [`load_segments`]. /// @@ -1136,8 +992,7 @@ fn canonicalize_inverted_index_details( /// /// `posting_format_version` records how a single segment physically stores /// postings, so mixed-version FTS segments may disagree on it without being -/// semantically incompatible. Every other field remains part of the equality -/// check. +/// incompatible. Every other field remains part of the equality check. fn inverted_index_details_semantically_equal( left: &InvertedIndexDetails, right: &InvertedIndexDetails, @@ -1328,26 +1183,6 @@ mod tests { ); } - #[test] - fn merge_physical_provenance_sidecar_sources() { - assert_eq!( - merge_physical_source_dataset_versions([Some(vec![19, 17, 19]), Some(vec![18, 19])]), - Some(vec![17, 18, 19]) - ); - assert_eq!( - merge_physical_source_dataset_versions([Some(vec![17]), None]), - None - ); - assert_eq!( - merge_physical_source_dataset_versions([Some(Vec::new())]), - None - ); - assert_eq!( - merge_physical_source_dataset_versions(std::iter::empty::>>()), - None - ); - } - #[test] fn inverted_details_equal_when_only_posting_format_version_differs() { let left = canonicalize_inverted_index_details( diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 84f5a575a9a..086df5ee323 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -38,6 +38,7 @@ use lance_core::{ use lance_datafusion::utils::{ExecutionPlanMetricsSetExt, MetricsExt, PARTITIONS_SEARCHED_METRIC}; use lance_select::RowAddrMask; use lance_table::format::IndexMetadata; +use rustc_hash::FxHashSet; use super::PreFilterSource; use super::utils::{IndexMetrics, PreFilterMasks, build_prefilter}; @@ -798,7 +799,7 @@ impl CompoundQueryExec { async fn index_query_local_residual_batch( residual: QueryLocalFtsIndex, batch: RecordBatch, - allowed_terms: Arc>, + allowed_terms: Arc>, ) -> Result { spawn_cpu(move || { let row_ids = batch @@ -823,10 +824,23 @@ async fn index_query_local_residual_batch( /// next batch, so the entire stream is never collected in memory and the /// number of live tokenizers/posting maps is bounded by the CPU pool size. async fn index_query_local_residual( - mut residual_input: SendableRecordBatchStream, + residual_input: SendableRecordBatchStream, seed: QueryLocalFtsIndex, - allowed_terms: Arc>, + allowed_terms: Arc>, ) -> DataFusionResult> { + // Match flat FTS's CPU-task sizing. Dataset scan batches are normally + // row-bounded (often 8,192 rows), which can leave a small residual with + // only one or two tokenizer tasks. Byte rechunking keeps tasks substantial + // while exposing enough parallelism for variable-width text. + const ACCUMULATE_BYTES: usize = 256 * 1024; + const SLICE_BYTES: usize = 512 * 1024; + let input_schema = residual_input.schema(); + let mut residual_input = Box::pin(lance_arrow::stream::rechunk_stream_by_size( + residual_input, + input_schema, + ACCUMULATE_BYTES, + SLICE_BYTES, + )); let parallelism = get_num_compute_intensive_cpus().max(1); let mut initial_batches = Vec::with_capacity(parallelism); let mut is_input_exhausted = false; @@ -881,22 +895,6 @@ async fn index_query_local_residual( Ok(shards) } -async fn query_local_residual_stats( - shards: Vec, - terms: Arc<[String]>, -) -> Result> { - stream::iter(shards.into_iter().map(|shard| { - let terms = terms.clone(); - spawn_cpu(move || { - let stats = shard.bm25_stats_for_terms(terms.as_ref()); - Ok::<_, Error>((shard, stats)) - }) - })) - .buffered(get_num_compute_intensive_cpus().max(1)) - .try_collect() - .await -} - async fn query_local_residual_leaves( shards: Vec, query: FtsQuery, @@ -927,9 +925,13 @@ async fn query_local_residual_leaves( Ok(merged) } -/// Exact compound FTS over committed postings plus an append-only residual -/// scan. The residual documents are tokenized once into query-local postings, -/// rather than once for every compound leaf. +/// Compound FTS over committed postings plus a small append-only residual scan. +/// +/// The residual documents are tokenized once into query-local postings, rather +/// than once for every compound leaf. Both arms use the committed segments' +/// BM25 corpus statistics. This intentionally assumes the residual rows follow +/// the indexed corpus distribution: scores can differ from a rebuilt index, +/// but remain comparable without paying to rebuild corpus statistics per query. #[derive(Debug)] pub(crate) struct HybridCompoundQueryExec { dataset: Arc, @@ -1040,74 +1042,43 @@ impl ExecutionPlan for HybridCompoundQueryExec { )) })?; let field_id = dataset.schema().field_id(&column)?; - let residual_seed = QueryLocalFtsIndex::try_with_params( + let tokenizer = first_index.tokenizer(); + let doc_type = tokenizer.doc_type(); + let residual_seed = QueryLocalFtsIndex::try_with_loaded_tokenizer( field_id, column.clone(), first_index.params().clone(), + tokenizer, )?; let terms = residual_seed.exact_query_terms(&query)?; if terms.is_empty() { metrics.baseline_metrics.record_output(0); return scored_documents_batch(schema, Vec::new()).map_err(DataFusionError::from); } - let allowed_terms = Arc::new(terms.iter().cloned().collect::>()); - let residual_input = residual_input.execute(partition, context.clone())?; - let residual_shards = - index_query_local_residual(residual_input, residual_seed, allowed_terms).await?; - - let query_tokens = Tokens::new(terms.clone(), first_index.tokenizer().doc_type()); + let allowed_terms = Arc::new(terms.iter().cloned().collect::>()); + let query_tokens = Tokens::new(terms.clone(), doc_type); let exact_params = params .clone() .with_fuzziness(Some(0)) .with_phrase_slop(None); - let mut scorer = build_global_bm25_scorer( - &indices, - &query_tokens, - &exact_params, - Some(metrics.as_ref()), - ) - .await?; - let residual_shards = query_local_residual_stats( - residual_shards, - Arc::from(terms.clone().into_boxed_slice()), - ) - .await?; - for (_, residual_stats) in &residual_shards { - scorer.total_tokens = scorer - .total_tokens - .checked_add(residual_stats.total_tokens) - .ok_or_else(|| { - DataFusionError::Execution( - "hybrid compound FTS total token count overflow".to_string(), - ) - })?; - scorer.num_docs = scorer - .num_docs - .checked_add(residual_stats.num_docs) - .ok_or_else(|| { - DataFusionError::Execution( - "hybrid compound FTS document count overflow".to_string(), - ) - })?; - for term in &terms { - let residual_df = residual_stats.num_docs_containing_token(term); - let df = scorer.token_docs.get_mut(term).ok_or_else(|| { - DataFusionError::Execution(format!( - "hybrid compound FTS scorer is missing query term '{term}'" - )) - })?; - *df = df.checked_add(residual_df).ok_or_else(|| { - DataFusionError::Execution(format!( - "hybrid compound FTS document frequency overflow for term '{term}'" - )) - })?; - } - } - let residual_shards = residual_shards - .into_iter() - .map(|(shard, _)| shard) - .collect::>(); - let scorer = Arc::new(scorer); + + let residual_context = context.clone(); + let residual_indexing = async move { + let residual_input = residual_input.execute(partition, residual_context)?; + index_query_local_residual(residual_input, residual_seed, allowed_terms).await + }; + let scorer_build = async { + let scorer = build_global_bm25_scorer( + &indices, + &query_tokens, + &exact_params, + Some(metrics.as_ref()), + ) + .await?; + DataFusionResult::>::Ok(Arc::new(scorer)) + }; + let (residual_shards, scorer) = + futures::future::try_join(residual_indexing, scorer_build).await?; let limit = params.limit.ok_or_else(|| { DataFusionError::Execution( "hybrid compound FTS requires a bounded result limit".to_string(), From 18a0ec4539f9925769c5627f0ff948d7628b8c61 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sun, 30 Aug 2026 17:52:39 +0800 Subject: [PATCH 20/22] fix(fts): preserve residual term scoring --- rust/lance/src/dataset/mem_wal/index.rs | 2 +- rust/lance/src/dataset/mem_wal/index/fts.rs | 96 +++++++++++++++++-- rust/lance/src/dataset/tests/dataset_index.rs | 62 ++++++++---- rust/lance/src/io/exec/fts.rs | 69 +++++++++---- 4 files changed, 181 insertions(+), 48 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index.rs b/rust/lance/src/dataset/mem_wal/index.rs index 7653891bcf2..2735511faba 100644 --- a/rust/lance/src/dataset/mem_wal/index.rs +++ b/rust/lance/src/dataset/mem_wal/index.rs @@ -49,8 +49,8 @@ pub type RowPosition = u64; // Re-export public types used externally pub use btree::{BTreeIndexConfig, BTreeMemIndex}; -pub(crate) use fts::QueryLocalFtsIndex; pub use fts::{FtsIndexConfig, FtsMemIndex, FtsQueryExpr, SearchOptions}; +pub(crate) use fts::{QueryLocalFtsIndex, QueryLocalFtsStats}; pub use hnsw::{HnswIndexConfig, HnswMemIndex}; pub use pk_key::encode_pk_tuple; diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 5f67cf78e11..6a188fc60cd 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1030,6 +1030,51 @@ pub struct QueryLocalFtsIndex { inner: FtsMemIndex, } +#[derive(Debug, Default)] +pub(crate) struct QueryLocalFtsStats { + doc_count: usize, + total_tokens: u64, + token_docs: FxHashMap, +} + +impl QueryLocalFtsStats { + pub(crate) fn checked_add_assign(&mut self, other: Self) -> Result<()> { + self.doc_count = self + .doc_count + .checked_add(other.doc_count) + .ok_or_else(|| Error::internal("query-local FTS document count overflow"))?; + self.total_tokens = self + .total_tokens + .checked_add(other.total_tokens) + .ok_or_else(|| Error::internal("query-local FTS total token count overflow"))?; + for (token, df) in other.token_docs { + let current = self.token_docs.entry(token).or_default(); + *current = current + .checked_add(df) + .ok_or_else(|| Error::internal("query-local FTS term document count overflow"))?; + } + Ok(()) + } + + pub(crate) fn add_to_scorer(&self, scorer: &mut MemBM25Scorer) -> Result<()> { + scorer.num_docs = scorer + .num_docs + .checked_add(self.doc_count) + .ok_or_else(|| Error::internal("residual BM25 document count overflow"))?; + scorer.total_tokens = scorer + .total_tokens + .checked_add(self.total_tokens) + .ok_or_else(|| Error::internal("residual BM25 total token count overflow"))?; + for (token, df) in &self.token_docs { + let current = scorer.token_docs.entry(token.clone()).or_default(); + *current = current + .checked_add(*df) + .ok_or_else(|| Error::internal("residual BM25 term document count overflow"))?; + } + Ok(()) + } +} + impl QueryLocalFtsIndex { #[cfg(test)] pub(crate) fn try_with_params( @@ -1103,7 +1148,7 @@ impl QueryLocalFtsIndex { batch: &RecordBatch, row_ids: &UInt64Array, terms: &FxHashSet, - ) -> Result<()> { + ) -> Result { self.inner .insert_with_row_ids_for_terms(batch, row_ids, terms) } @@ -1402,7 +1447,7 @@ impl FtsMemIndex { batch: &RecordBatch, row_ids: &UInt64Array, terms: &FxHashSet, - ) -> Result<()> { + ) -> Result { if row_ids.len() != batch.num_rows() || row_ids.null_count() != 0 { return Err(Error::invalid_input(format!( "MemWAL FTS explicit row ids require {} non-null values, got len={} nulls={}", @@ -1424,6 +1469,7 @@ impl FtsMemIndex { }, None, ) + .map(|_| ()) } fn insert_batch_with_keys( @@ -1431,7 +1477,7 @@ impl FtsMemIndex { batch: &RecordBatch, row_position: impl Fn(usize) -> Result, allowed_terms: Option<&FxHashSet>, - ) -> Result<()> { + ) -> Result { let st = self.state.load_full(); let document_position_start = st.tail.doc_count(); if self.resolved_field.get().is_none() { @@ -1465,6 +1511,8 @@ impl FtsMemIndex { Vec::with_capacity(batch.num_rows()) }; let mut total_tokens: u64 = 0; + let mut query_local_corpus_doc_count = 0usize; + let mut query_local_corpus_total_tokens = 0u64; let preserve_zero_token_documents = self.params.get_document_granularity().is_list_element(); let mut index_document = |key: DocumentKey, text: &str| -> Result<()> { @@ -1482,10 +1530,19 @@ impl FtsMemIndex { false, ), }; + let belongs_in_corpus = preserve_zero_token_documents || num_tokens > 0; + if allowed_terms.is_some() && belongs_in_corpus { + query_local_corpus_doc_count = query_local_corpus_doc_count + .checked_add(1) + .ok_or_else(|| Error::internal("query-local FTS document count overflow"))?; + query_local_corpus_total_tokens = query_local_corpus_total_tokens + .checked_add(num_tokens as u64) + .ok_or_else(|| Error::internal("query-local FTS total token count overflow"))?; + } let retain_document = if allowed_terms.is_some() { retained_term } else { - preserve_zero_token_documents || num_tokens > 0 + belongs_in_corpus }; if retain_document { documents.push(DocumentMetadata { key, num_tokens }); @@ -1504,8 +1561,21 @@ impl FtsMemIndex { )?; } + let query_local_stats = if allowed_terms.is_some() { + QueryLocalFtsStats { + doc_count: query_local_corpus_doc_count, + total_tokens: query_local_corpus_total_tokens, + token_docs: term_builders + .iter() + .map(|(term, builder)| (term.to_string(), builder.row_positions.len())) + .collect(), + } + } else { + QueryLocalFtsStats::default() + }; + if documents.is_empty() { - return Ok(()); + return Ok(query_local_stats); } // Drop the tokenizer guard before publishing so we don't hold it @@ -1525,7 +1595,7 @@ impl FtsMemIndex { if self.background_maintenance && st.tail.doc_count() >= self.freeze_threshold_rows as u64 { self.freeze(&st)?; } - Ok(()) + Ok(query_local_stats) } /// Analyze every exact leaf and return the deduplicated query terms in @@ -4629,16 +4699,24 @@ mod tests { .unwrap(); let full_index = FtsMemIndex::new(1, "description".to_string()); - index + let stats = index .insert_with_row_ids_for_terms(&batch, &row_ids, &terms) .unwrap(); full_index.insert(&batch, 0).unwrap(); - // The unmatched row (row id 42) contributes neither postings nor - // metadata; retained documents still keep their full token counts. + // The unmatched nonempty row (row id 42) contributes no postings or + // metadata, but remains part of the approximate residual BM25 corpus. assert_eq!(index.doc_count(), 2); assert_eq!(index.inner.entry_count(), 2); + assert_eq!(stats.doc_count, 3); + assert_eq!(stats.total_tokens, 6); + assert_eq!(stats.token_docs.get("hello"), Some(&2)); let committed_scorer = MemBM25Scorer::new(6, 3, HashMap::from([("hello".to_string(), 2)])); + let mut residual_scorer = committed_scorer.clone(); + stats.add_to_scorer(&mut residual_scorer).unwrap(); + assert_eq!(residual_scorer.num_docs, 6); + assert_eq!(residual_scorer.total_tokens, 12); + assert_eq!(residual_scorer.token_docs.get("hello"), Some(&4)); let query = FtsQuery::Match( lance_index::scalar::inverted::query::MatchQuery::new("hello".to_string()) diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index f378c57ae45..6b71e0f2b66 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2755,7 +2755,7 @@ async fn test_partial_compound_hybrid_prunes_same_path_different_base_rewrite() ) .unwrap(); let schema = initial.schema(); - let dataset = Dataset::write( + let mut dataset = Dataset::write( RecordBatchIterator::new(vec![initial].into_iter().map(Ok), schema.clone()), &primary, Some(WriteParams { @@ -2885,7 +2885,7 @@ async fn test_partial_compound_hybrid_prunes_same_path_different_base_rewrite() } #[tokio::test] -async fn test_partial_compound_hybrid_uses_committed_index_statistics() { +async fn test_partial_compound_hybrid_uses_mixed_approximate_statistics() { let initial = arrow_array::record_batch!( ("text", Utf8, ["fresh alpha", "blocked fresh alpha"]), ("id", Int32, [0, 1]) @@ -2909,7 +2909,7 @@ async fn test_partial_compound_hybrid_uses_committed_index_statistics() { "fresh alpha", "fresh beta", "fresh alpha", - "fresh gamma", + "fresh beta beta", "fresh beta blocked" ] ), @@ -2948,10 +2948,19 @@ async fn test_partial_compound_hybrid_uses_committed_index_statistics() { 5, "MUST_NOT must exclude the blocked row" ); - assert_eq!(partial_boost[0].1.to_bits(), partial_boost[1].1.to_bits()); + let partial_boost_positions = partial_boost + .iter() + .enumerate() + .map(|(position, (row_id, _))| (*row_id, position)) + .collect::>(); + let partial_boost_scores = partial_boost + .iter() + .map(|(row_id, score)| (*row_id, score.to_bits())) + .collect::>(); + assert_eq!(partial_boost_scores.get(&2), partial_boost_scores.get(&4)); assert!( - partial_boost[0].0 < partial_boost[1].0, - "equal-score rows must use ascending row id as the exact tie break" + partial_boost_positions[&2] < partial_boost_positions[&4], + "equal-score residual rows must use ascending row id as the exact tie break" ); let multimatch_query: FtsQuery = MultiMatchQuery::try_new( "fresh alpha".to_string(), @@ -2980,13 +2989,20 @@ async fn test_partial_compound_hybrid_uses_committed_index_statistics() { .copied() .zip(scores.iter().map(|score| score.to_bits())) .collect::>(); - for residual_id in [2, 4] { - assert_eq!( - score_bits.get(&residual_id), - score_bits.get(&0), - "{query_name} must score identical indexed and residual documents identically" - ); - } + let positions = ids + .iter() + .enumerate() + .map(|(position, row_id)| (*row_id, position)) + .collect::>(); + assert_eq!( + score_bits.get(&2), + score_bits.get(&4), + "{query_name} must preserve equal scores within the residual arm" + ); + assert!( + positions[&2] < positions[&4], + "{query_name} must preserve the row-id tie break within the residual arm" + ); } let mut indexed_only_scanner = dataset.scan(); @@ -3010,6 +3026,9 @@ async fn test_partial_compound_hybrid_uses_committed_index_statistics() { "hybrid scoring must preserve committed-index scores for indexed row {row_id}" ); } + // The residual arm intentionally uses committed + query-local statistics, + // so its scores are not expected to equal either indexed-arm scores or a + // fully rebuilt index's exact global scores. let residual_only_query: FtsQuery = BooleanQuery::new([ (Occur::Must, compound_match_query("beta", "text", 1.0)), @@ -3031,15 +3050,16 @@ async fn test_partial_compound_hybrid_uses_committed_index_statistics() { let residual_only = residual_only_scanner.try_into_batch().await.unwrap(); assert_eq!( residual_only["id"].as_primitive::().values(), - &[3], - "the zero-weight beta leaf must retain membership while MUST_NOT excludes id=6" + &[5, 3], + "residual beta TF must rank id=5 first while MUST_NOT excludes id=6" ); - assert_eq!( - residual_only[SCORE_COL] - .as_primitive::() - .values(), - &[0.0], - "a residual-only term must contribute zero when committed df is zero" + let residual_only_scores = residual_only[SCORE_COL] + .as_primitive::() + .values(); + assert!( + residual_only_scores.iter().all(|score| score.is_finite()) + && residual_only_scores[0] > residual_only_scores[1], + "residual-only terms must retain membership and use query-local TF/DF scoring" ); } diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 086df5ee323..3483a476d52 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -42,7 +42,7 @@ use rustc_hash::FxHashSet; use super::PreFilterSource; use super::utils::{IndexMetrics, PreFilterMasks, build_prefilter}; -use crate::dataset::mem_wal::index::QueryLocalFtsIndex; +use crate::dataset::mem_wal::index::{QueryLocalFtsIndex, QueryLocalFtsStats}; use crate::index::scalar::inverted::{ ResolvedFtsField, fts_document_schema, load_segment_details, load_segments, transform_fts_document_stream, @@ -796,11 +796,17 @@ impl CompoundQueryExec { } } +#[derive(Debug)] +struct QueryLocalResidualShard { + index: QueryLocalFtsIndex, + stats: QueryLocalFtsStats, +} + async fn index_query_local_residual_batch( - residual: QueryLocalFtsIndex, + mut residual: QueryLocalResidualShard, batch: RecordBatch, allowed_terms: Arc>, -) -> Result { +) -> Result { spawn_cpu(move || { let row_ids = batch .column_by_name(ROW_ID) @@ -810,7 +816,12 @@ async fn index_query_local_residual_batch( ) })? .as_primitive::(); - residual.insert_with_row_ids_for_terms(&batch, row_ids, allowed_terms.as_ref())?; + let stats = residual.index.insert_with_row_ids_for_terms( + &batch, + row_ids, + allowed_terms.as_ref(), + )?; + residual.stats.checked_add_assign(stats)?; Ok(residual) }) .await @@ -827,7 +838,7 @@ async fn index_query_local_residual( residual_input: SendableRecordBatchStream, seed: QueryLocalFtsIndex, allowed_terms: Arc>, -) -> DataFusionResult> { +) -> DataFusionResult> { // Match flat FTS's CPU-task sizing. Dataset scan batches are normally // row-bounded (often 8,192 rows), which can leave a small residual with // only one or two tokenizer tasks. Byte rechunking keeps tasks substantial @@ -854,16 +865,25 @@ async fn index_query_local_residual( } if initial_batches.is_empty() { - return Ok(vec![seed]); + return Ok(vec![QueryLocalResidualShard { + index: seed, + stats: QueryLocalFtsStats::default(), + }]); } // Construct every shard from the already-loaded seed before dispatching // CPU work. This keeps tokenizer model I/O out of `spawn_cpu` closures. let mut initial_shards = Vec::with_capacity(initial_batches.len()); for _ in 1..initial_batches.len() { - initial_shards.push(seed.empty_sibling()); + initial_shards.push(QueryLocalResidualShard { + index: seed.empty_sibling(), + stats: QueryLocalFtsStats::default(), + }); } - initial_shards.push(seed); + initial_shards.push(QueryLocalResidualShard { + index: seed, + stats: QueryLocalFtsStats::default(), + }); let mut in_flight = FuturesUnordered::new(); for (shard, batch) in initial_shards.into_iter().zip(initial_batches) { @@ -896,14 +916,14 @@ async fn index_query_local_residual( } async fn query_local_residual_leaves( - shards: Vec, + shards: Vec, query: FtsQuery, scorer: Arc, ) -> Result>> { let shard_leaves = stream::iter(shards.into_iter().map(|shard| { let query = query.clone(); let scorer = scorer.clone(); - spawn_cpu(move || shard.exact_leaf_results(&query, scorer.as_ref())) + spawn_cpu(move || shard.index.exact_leaf_results(&query, scorer.as_ref())) })) .buffered(get_num_compute_intensive_cpus().max(1)) .try_collect::>() @@ -925,13 +945,25 @@ async fn query_local_residual_leaves( Ok(merged) } +fn residual_bm25_scorer( + committed_scorer: &MemBM25Scorer, + shards: &[QueryLocalResidualShard], +) -> Result { + let mut scorer = committed_scorer.clone(); + for shard in shards { + shard.stats.add_to_scorer(&mut scorer)?; + } + Ok(scorer) +} + /// Compound FTS over committed postings plus a small append-only residual scan. /// /// The residual documents are tokenized once into query-local postings, rather -/// than once for every compound leaf. Both arms use the committed segments' -/// BM25 corpus statistics. This intentionally assumes the residual rows follow -/// the indexed corpus distribution: scores can differ from a rebuilt index, -/// but remain comparable without paying to rebuild corpus statistics per query. +/// than once for every compound leaf. The indexed arm uses committed-index +/// BM25 statistics. The residual arm extends those statistics with the +/// query-local materialized documents, which matches the established mixed +/// flat-search approximation without rescanning the residual input or rebuilding +/// exact corpus statistics. #[derive(Debug)] pub(crate) struct HybridCompoundQueryExec { dataset: Arc, @@ -1077,8 +1109,12 @@ impl ExecutionPlan for HybridCompoundQueryExec { .await?; DataFusionResult::>::Ok(Arc::new(scorer)) }; - let (residual_shards, scorer) = + let (residual_shards, committed_scorer) = futures::future::try_join(residual_indexing, scorer_build).await?; + let residual_scorer = Arc::new(residual_bm25_scorer( + committed_scorer.as_ref(), + &residual_shards, + )?); let limit = params.limit.ok_or_else(|| { DataFusionError::Execution( "hybrid compound FTS requires a bounded result limit".to_string(), @@ -1102,10 +1138,9 @@ impl ExecutionPlan for HybridCompoundQueryExec { ¶ms, prefilter, metrics.clone(), - scorer.clone(), + committed_scorer, ); let residual_query = query.clone(); - let residual_scorer = scorer.clone(); let residual_search = async move { let residual_leaves = query_local_residual_leaves( residual_shards, From f048d57c8a30b43b7e7ed851526a55f442657ff7 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sun, 30 Aug 2026 18:31:30 +0800 Subject: [PATCH 21/22] test(fts): fix residual scorer expectations --- rust/lance/src/dataset/mem_wal/index/fts.rs | 4 ++-- rust/lance/src/dataset/tests/dataset_index.rs | 14 -------------- 2 files changed, 2 insertions(+), 16 deletions(-) diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 6a188fc60cd..451089c07d7 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -4709,13 +4709,13 @@ mod tests { assert_eq!(index.doc_count(), 2); assert_eq!(index.inner.entry_count(), 2); assert_eq!(stats.doc_count, 3); - assert_eq!(stats.total_tokens, 6); + assert_eq!(stats.total_tokens, 5); assert_eq!(stats.token_docs.get("hello"), Some(&2)); let committed_scorer = MemBM25Scorer::new(6, 3, HashMap::from([("hello".to_string(), 2)])); let mut residual_scorer = committed_scorer.clone(); stats.add_to_scorer(&mut residual_scorer).unwrap(); assert_eq!(residual_scorer.num_docs, 6); - assert_eq!(residual_scorer.total_tokens, 12); + assert_eq!(residual_scorer.total_tokens, 11); assert_eq!(residual_scorer.token_docs.get("hello"), Some(&4)); let query = FtsQuery::Match( diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index 6b71e0f2b66..19673e83ae1 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -2948,20 +2948,6 @@ async fn test_partial_compound_hybrid_uses_mixed_approximate_statistics() { 5, "MUST_NOT must exclude the blocked row" ); - let partial_boost_positions = partial_boost - .iter() - .enumerate() - .map(|(position, (row_id, _))| (*row_id, position)) - .collect::>(); - let partial_boost_scores = partial_boost - .iter() - .map(|(row_id, score)| (*row_id, score.to_bits())) - .collect::>(); - assert_eq!(partial_boost_scores.get(&2), partial_boost_scores.get(&4)); - assert!( - partial_boost_positions[&2] < partial_boost_positions[&4], - "equal-score residual rows must use ascending row id as the exact tie break" - ); let multimatch_query: FtsQuery = MultiMatchQuery::try_new( "fresh alpha".to_string(), vec!["text".to_string(), "text".to_string()], From 7c30fa8d2296331576d1e6a10729f744bba62c55 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Sun, 30 Aug 2026 19:17:04 +0800 Subject: [PATCH 22/22] fix(fts): satisfy query-local stats visibility lint --- rust/lance/src/dataset/mem_wal/index/fts.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/mem_wal/index/fts.rs b/rust/lance/src/dataset/mem_wal/index/fts.rs index 451089c07d7..50ab8cc3b36 100644 --- a/rust/lance/src/dataset/mem_wal/index/fts.rs +++ b/rust/lance/src/dataset/mem_wal/index/fts.rs @@ -1031,7 +1031,7 @@ pub struct QueryLocalFtsIndex { } #[derive(Debug, Default)] -pub(crate) struct QueryLocalFtsStats { +pub struct QueryLocalFtsStats { doc_count: usize, total_tokens: u64, token_docs: FxHashMap,