diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 804129ad690..698f0118921 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -3575,40 +3575,38 @@ impl Scanner { .to_string())); } - // Mask data overlay files: a fragment with an overlay committed after this FTS index can - // no longer be trusted to its inverted-index positions, so a stale phrase could still - // match. Exclude any segment covering such a fragment. Unlike match queries, phrase - // queries have no flat re-evaluation path, so — exactly as for unindexed fragments, which - // phrase queries already do not search — overlaid fragments are simply dropped from the - // phrase result; new phrase matches there surface once compaction folds the overlay into - // the base. This removes stale hits without introducing wrong ones. + // Mask data overlay files: a row with an overlay committed after this FTS index touched the + // indexed column can no longer be trusted to its inverted-index positions, so its stale + // phrase positions could still match. Block just those rows from the phrase result, leaving + // every other (fresh) row in their segments searchable. Blocking whole segments instead + // would also drop unrelated matches co-located in the same segment. Phrase queries have no + // flat re-evaluation path (unlike match queries — see `plan_match_query`), so a new phrase + // match on a blocked row's current value only surfaces once compaction folds the overlay + // into the base. This removes stale hits without introducing wrong ones. let target_fragments = self .fragments .clone() .unwrap_or_else(|| self.dataset.fragments().to_vec()); - let (_flat_frag_ids, fresh_segments) = self - .fts_stale_frags_and_fresh_segments(&column, &target_fragments) - .await?; + let overlaid_frags = overlaid_fragments(&target_fragments); + let mut stale_rows: HashMap = HashMap::new(); + if !overlaid_frags.is_empty() { + for segment in &segments { + collect_overlay_stale_rows_for_segment(segment, &overlaid_frags, &mut stale_rows)?; + } + } + let overlay_block = self.stale_rows_block_mask(&stale_rows).await?; - let exec: Arc = match fresh_segments { - // Every segment covers a stale fragment: with no flat phrase path there is nothing - // trustworthy left to search, so the phrase query returns no rows. - Some(segs) if segs.is_empty() => Arc::new(EmptyExec::new(FTS_SCHEMA.clone())), - Some(segs) => Arc::new(PhraseQueryExec::new_with_segments( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - segs, - )), - None => Arc::new(PhraseQueryExec::new( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - )), - }; - Ok(exec) + let mut exec = PhraseQueryExec::new_with_segments( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + segments, + ); + if let Some(block) = overlay_block { + exec = exec.with_overlay_block(block); + } + Ok(Arc::new(exec)) } async fn plan_match_query( diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index cfabdc13266..293b50cc51c 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -11,7 +11,9 @@ use futures::TryStreamExt; use arrow_array::cast::AsArray; use arrow_array::types::Int32Type; -use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StringArray}; +use arrow_array::{ + ArrayRef, Int32Array, RecordBatch, RecordBatchIterator, StringArray, record_batch, +}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_index::IndexType; use lance_index::scalar::FullTextSearchQuery; @@ -815,6 +817,77 @@ async fn test_fts_overlay_unrelated_field_not_excluded() { assert_eq!(fts_ids_matching(&dataset, "banana").await, vec![3, 999]); } +/// Two-fragment text dataset where several rows share a phrase, plus one non-matching row that +/// will be overlaid. `text` is field 1. Fragment 0 (ids 0..6) holds three "red car" phrases and +/// the overlay target (id=2); fragment 1 (ids 6..12) holds a fourth "red car" phrase. +async fn create_phrase_sibling_dataset() -> Dataset { + let batch = record_batch!( + ("id", Int32, (0..12).collect::>()), + ( + "text", + Utf8, + vec![ + "red car fast", // 0 — matches "red car" + "blue sky wide", // 1 + "old truck slow", // 2 — overlay target (does NOT match the phrase) + "red car shiny", // 3 — matches "red car" + "green boat calm", // 4 + "red car brand", // 5 — matches "red car" + "red car extra", // 6 — fragment 1, matches "red car" + "pear tart sweet", // 7 + "lemon curd tang", // 8 + "peach cobbler warm", // 9 + "plum pudding rich", // 10 + "fig newton crisp", // 11 + ] + ) + ) + .unwrap(); + let schema = batch.schema(); + let write_params = WriteParams { + max_rows_per_file: 6, + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap() +} + +/// A phrase query must keep matches from rows that are merely co-located (same fragment/segment) +/// with an overlaid row. Overlaying one row's `text` must drop only that row's stale phrase hit, +/// not every phrase hit in its segment. Regression for the FTS phrase path excluding whole +/// overlay-stale segments without re-evaluating the sibling rows they cover. +#[tokio::test] +async fn test_fts_phrase_overlay_keeps_sibling_matches() { + let mut dataset = create_phrase_sibling_dataset().await; + build_text_fts_index_with_positions(&mut dataset).await; + + // Before any overlay the phrase "red car" matches ids 0, 3, 5 (fragment 0) and 6 (fragment 1). + assert_eq!( + fts_phrase_ids_matching(&dataset, "red car").await, + vec![0, 3, 5, 6] + ); + + // Overlay id=2's text (field 1). Its old value "old truck slow" never matched the phrase, so + // the correct result is unchanged. The overlay must not evict its siblings 0, 3 and 5. + let dataset = commit_overlay( + dataset, + "phrase_sibling_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([2])), + vec![Arc::new(StringArray::from(vec![Some("zzoverlay")]))], + ) + .await; + + assert_eq!( + fts_phrase_ids_matching(&dataset, "red car").await, + vec![0, 3, 5, 6], + "overlaying a non-matching row dropped its co-located phrase matches" + ); +} + /// Benchmark: measure query latency for BTree, FTS, and vector ANN with 0/4/16 overlay layers. /// /// Run with: cargo test -p lance --lib --release -- overlay_index_masking::bench --ignored --nocapture diff --git a/rust/lance/src/index/prefilter.rs b/rust/lance/src/index/prefilter.rs index 1e02b6f807a..179c5fda97c 100644 --- a/rust/lance/src/index/prefilter.rs +++ b/rust/lance/src/index/prefilter.rs @@ -238,6 +238,12 @@ impl DatasetPreFilter { self } + /// Like [`Self::with_overlay_block`] but mutates in place, for callers holding the prefilter + /// behind an `Arc` before `wait_for_ready` (e.g. `PhraseQueryExec`). + pub fn set_overlay_block(&mut self, block: RowAddrMask) { + self.overlay_block = Some(block); + } + /// Creates a task to load a mask that filters out deleted rows and, /// when `restrict_to_fragments` is true, also restricts results to only /// the given `fragments`. diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 2c1859211f6..7f2ba5211fb 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -31,6 +31,7 @@ use lance_core::{ utils::{tokio::get_num_compute_intensive_cpus, tracing::StreamTracingExt}, }; use lance_datafusion::utils::{ExecutionPlanMetricsSetExt, MetricsExt, PARTITIONS_SEARCHED_METRIC}; +use lance_select::RowAddrMask; use lance_table::format::IndexMetadata; use super::PreFilterSource; @@ -1155,6 +1156,12 @@ pub struct PhraseQueryExec { /// Optional pre-resolved segment list. See /// [`MatchQueryExec::new_with_segments`]. preset_segments: Option>, + /// Row addresses whose inverted-index entries are stale because a data overlay committed after + /// the index touched the indexed column. Blocked from the phrase result so the pre-overlay + /// positions never match; unlike match queries there is no flat phrase path to re-score them, + /// so a new phrase match on the current value only surfaces after compaction folds the overlay + /// into the base. + overlay_block: Option, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -1204,6 +1211,7 @@ impl PhraseQueryExec { prefilter_source, base_scorer: None, preset_segments: None, + overlay_block: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -1232,6 +1240,7 @@ impl PhraseQueryExec { prefilter_source, base_scorer: None, preset_segments: Some(segments), + overlay_block: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -1243,6 +1252,12 @@ impl PhraseQueryExec { self } + /// Block the given stale row addresses from the phrase result (see the `overlay_block` field). + pub fn with_overlay_block(mut self, overlay_block: RowAddrMask) -> Self { + self.overlay_block = Some(overlay_block); + self + } + pub fn query(&self) -> &PhraseQuery { &self.query } @@ -1301,6 +1316,7 @@ impl ExecutionPlan for PhraseQueryExec { prefilter_source: PreFilterSource::None, base_scorer: self.base_scorer.clone(), preset_segments: self.preset_segments.clone(), + overlay_block: self.overlay_block.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), }, @@ -1326,6 +1342,7 @@ impl ExecutionPlan for PhraseQueryExec { prefilter_source, base_scorer: self.base_scorer.clone(), preset_segments: self.preset_segments.clone(), + overlay_block: self.overlay_block.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -1351,6 +1368,7 @@ impl ExecutionPlan for PhraseQueryExec { let prefilter_source = self.prefilter_source.clone(); let preset_base_scorer = self.base_scorer.clone(); let preset_segments = self.preset_segments.clone(); + let overlay_block = self.overlay_block.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); let stream = stream::once(async move { let _timer = metrics.baseline_metrics.elapsed_compute().timer(); @@ -1385,6 +1403,11 @@ impl ExecutionPlan for PhraseQueryExec { .expect("prefilter just created") .set_deleted_fragments(deleted_fragments); } + if let Some(overlay_block) = overlay_block { + Arc::get_mut(&mut pre_filter) + .expect("prefilter just created") + .set_overlay_block(overlay_block); + } metrics .record_parts_searched(indices.iter().map(|index| index.partition_count()).sum());