From a4d6e489c627fe4b0e49d9a2991c9436313debc0 Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 1 Sep 2026 11:15:00 +0800 Subject: [PATCH 1/2] update --- src/fts_query.rs | 2 ++ src/scanner.rs | 44 ++++++++++++++++++++++--- tests/c_api_test.rs | 78 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 120 insertions(+), 4 deletions(-) diff --git a/src/fts_query.rs b/src/fts_query.rs index cd194c7..e85ed99 100644 --- a/src/fts_query.rs +++ b/src/fts_query.rs @@ -54,6 +54,7 @@ pub(crate) struct FtsQueryContextInner { pub(crate) query: FullTextSearchQuery, pub(crate) segments: Vec, pub(crate) scorer: Arc, + pub(crate) has_unindexed_fragments: bool, } impl FtsQueryContextInner { @@ -216,6 +217,7 @@ async fn prepare_fts_query_context( query, segments, scorer, + has_unindexed_fragments: !unindexed_fragment_ids.is_empty(), }) } diff --git a/src/scanner.rs b/src/scanner.rs index 5b8c34e..f60f0c5 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -179,6 +179,45 @@ impl LanceScanner { Ok(()) } + /// Restrict an INDEX_ONLY prepared FTS scan to fragments covered by the + /// selected committed segments. This is deliberately separate from + /// `fast_search`: that option is scanner-wide, changes unrelated scalar + /// index fallback behavior, and also forces `_rowid` into the output. + fn apply_prepared_fts_fragment_filter( + &self, + scanner: &mut lance::dataset::scanner::Scanner, + context: &FtsQueryContextInner, + segments: &[IndexMetadata], + ) -> Result<()> { + if !context.has_unindexed_fragments { + return Ok(()); + } + + let mut selected_fragment_ids = std::collections::HashSet::new(); + for segment in segments { + let fragment_bitmap = segment.fragment_bitmap.as_ref().ok_or_else(|| { + lance_core::Error::internal(format!( + "prepared FTS segment {} lost its validated fragment coverage", + segment.uuid + )) + })?; + selected_fragment_ids.extend(fragment_bitmap.iter()); + } + + let selected_fragments = self + .dataset + .get_fragments() + .into_iter() + .filter(|fragment| { + u32::try_from(fragment.id()) + .is_ok_and(|fragment_id| selected_fragment_ids.contains(&fragment_id)) + }) + .map(|fragment| fragment.metadata().clone()) + .collect(); + scanner.with_fragments(selected_fragments); + Ok(()) + } + fn apply_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(substrait) = &self.substrait_filter { scanner.filter_substrait(substrait)?; @@ -282,11 +321,8 @@ impl LanceScanner { let distributed_fts = if let Some(context) = &self.fts_context { context.validate_dataset_identity(&self.dataset)?; let segments = select_fts_segments(context, self.fts_index_segments.as_deref())?; + self.apply_prepared_fts_fragment_filter(&mut scanner, context, &segments)?; scanner.full_text_search(context.query.clone())?; - // Both STRICT and INDEX_ONLY context scans must use only the - // committed segments pinned in the context. In STRICT mode all - // current fragments were already proven covered during prepare. - scanner.fast_search(); Some(PreparedFtsExecution { context: Arc::clone(context), segments, diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 74b9f85..3627c4e 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -5762,6 +5762,84 @@ fn load_fts_segment_uuids(uri: &str, column: &str) -> Vec<[u8; 16]> { }) } +#[test] +fn test_prepared_fts_row_id_output_is_explicit() { + let (_tmp, uri) = create_test_dataset(); + let uri_c = c_str(&uri); + let column = c_str("name"); + let query = c_str("alice"); + let inverted_params = c_str(r#"{"base_tokenizer":"simple","language":"English"}"#); + + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + assert_eq!( + unsafe { + lance_dataset_create_scalar_index( + dataset, + column.as_ptr(), + ptr::null(), + LanceScalarIndexType::Inverted as i32, + inverted_params.as_ptr(), + false, + ) + }, + 0 + ); + let context = unsafe { + lance_dataset_prepare_fts_query( + dataset, + column.as_ptr(), + query.as_ptr(), + 0, + LanceFtsCoverageMode::Strict as i32, + ) + }; + assert!(!context.is_null(), "{}", unsafe { + std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() + }); + + let id = c_str("id"); + let columns = [id.as_ptr(), ptr::null()]; + let scan_schema = |with_row_id: bool| { + let scanner = unsafe { lance_scanner_new(dataset, columns.as_ptr(), ptr::null()) }; + assert!(!scanner.is_null()); + if with_row_id { + assert_eq!(unsafe { lance_scanner_with_row_id(scanner, true) }, 0); + } + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0, + "{}", + unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() } + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; + let schema = reader.schema(); + let rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum(); + assert!(rows > 0); + unsafe { lance_scanner_close(scanner) }; + schema + }; + + let without_row_id = scan_schema(false); + assert_eq!(without_row_id.fields().len(), 2); + assert!(without_row_id.field_with_name("id").is_ok()); + assert!(without_row_id.field_with_name("_score").is_ok()); + assert!(without_row_id.field_with_name("_rowid").is_err()); + + let with_row_id = scan_schema(true); + assert_eq!(with_row_id.fields().len(), 3); + assert!(with_row_id.field_with_name("id").is_ok()); + assert!(with_row_id.field_with_name("_score").is_ok()); + assert!(with_row_id.field_with_name("_rowid").is_ok()); + + unsafe { lance_fts_query_context_close(context) }; + unsafe { lance_dataset_close(dataset) }; +} + #[test] fn test_prepare_fts_query_index_only_allows_unindexed_fragment() { let (_tmp, uri) = create_test_dataset(); From 6f0fae4cc51bf144685564b1d2e9f6f7afc71f8a Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Tue, 1 Sep 2026 12:34:05 +0800 Subject: [PATCH 2/2] update --- src/fts_query.rs | 2 - src/scanner.rs | 253 ++++++++++++++++++++++++++++++++++---------- tests/c_api_test.rs | 79 ++++++++++++++ 3 files changed, 279 insertions(+), 55 deletions(-) diff --git a/src/fts_query.rs b/src/fts_query.rs index e85ed99..cd194c7 100644 --- a/src/fts_query.rs +++ b/src/fts_query.rs @@ -54,7 +54,6 @@ pub(crate) struct FtsQueryContextInner { pub(crate) query: FullTextSearchQuery, pub(crate) segments: Vec, pub(crate) scorer: Arc, - pub(crate) has_unindexed_fragments: bool, } impl FtsQueryContextInner { @@ -217,7 +216,6 @@ async fn prepare_fts_query_context( query, segments, scorer, - has_unindexed_fragments: !unindexed_fragment_ids.is_empty(), }) } diff --git a/src/scanner.rs b/src/scanner.rs index f60f0c5..0c29b17 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -12,13 +12,13 @@ use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; use arrow::ffi_stream::FFI_ArrowArrayStream; use arrow_schema::SchemaRef; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::{ExecutionPlan, empty::EmptyExec}; use futures::{FutureExt, Stream, StreamExt}; use lance::Dataset; use lance::dataset::scanner::{ DatasetRecordBatchStream, ExecutionStatsCallback, ExecutionSummaryCounts, }; -use lance::io::exec::fts::MatchQueryExec; +use lance::io::exec::fts::{FlatMatchQueryExec, MatchQueryExec}; use lance_core::Result; use lance_index::scalar::FullTextSearchQuery; use lance_io::stream::RecordBatchStream; @@ -179,45 +179,6 @@ impl LanceScanner { Ok(()) } - /// Restrict an INDEX_ONLY prepared FTS scan to fragments covered by the - /// selected committed segments. This is deliberately separate from - /// `fast_search`: that option is scanner-wide, changes unrelated scalar - /// index fallback behavior, and also forces `_rowid` into the output. - fn apply_prepared_fts_fragment_filter( - &self, - scanner: &mut lance::dataset::scanner::Scanner, - context: &FtsQueryContextInner, - segments: &[IndexMetadata], - ) -> Result<()> { - if !context.has_unindexed_fragments { - return Ok(()); - } - - let mut selected_fragment_ids = std::collections::HashSet::new(); - for segment in segments { - let fragment_bitmap = segment.fragment_bitmap.as_ref().ok_or_else(|| { - lance_core::Error::internal(format!( - "prepared FTS segment {} lost its validated fragment coverage", - segment.uuid - )) - })?; - selected_fragment_ids.extend(fragment_bitmap.iter()); - } - - let selected_fragments = self - .dataset - .get_fragments() - .into_iter() - .filter(|fragment| { - u32::try_from(fragment.id()) - .is_ok_and(|fragment_id| selected_fragment_ids.contains(&fragment_id)) - }) - .map(|fragment| fragment.metadata().clone()) - .collect(); - scanner.with_fragments(selected_fragments); - Ok(()) - } - fn apply_filter(&self, scanner: &mut lance::dataset::scanner::Scanner) -> Result<()> { if let Some(substrait) = &self.substrait_filter { scanner.filter_substrait(substrait)?; @@ -321,7 +282,6 @@ impl LanceScanner { let distributed_fts = if let Some(context) = &self.fts_context { context.validate_dataset_identity(&self.dataset)?; let segments = select_fts_segments(context, self.fts_index_segments.as_deref())?; - self.apply_prepared_fts_fragment_filter(&mut scanner, context, &segments)?; scanner.full_text_search(context.query.clone())?; Some(PreparedFtsExecution { context: Arc::clone(context), @@ -361,14 +321,24 @@ impl PreparedScanner { return self.scanner.try_into_stream().await; }; let plan = self.scanner.create_plan().await?; - let (plan, replaced) = replace_match_query_exec( + let selected_segments_have_current_fragments = segments_have_current_fragments( + &distributed_fts.context.dataset, + &distributed_fts.segments, + )?; + let (plan, rewritten) = rewrite_prepared_fts_plan( plan, &distributed_fts.segments, &distributed_fts.context.scorer, + selected_segments_have_current_fragments, )?; - if replaced != 1 { + if rewritten.match_query_execs > 1 + || rewritten.flat_match_query_execs > 1 + || rewritten.match_query_execs + rewritten.flat_match_query_execs == 0 + || (selected_segments_have_current_fragments && rewritten.match_query_execs != 1) + { return Err(lance_core::Error::internal(format!( - "expected exactly one MatchQueryExec in prepared FTS plan, replaced {replaced}" + "unexpected prepared FTS plan for selected segments with current fragment coverage {selected_segments_have_current_fragments}: rewrote {} MatchQueryExec node(s) and removed {} FlatMatchQueryExec node(s)", + rewritten.match_query_execs, rewritten.flat_match_query_execs ))); } let stream = lance_datafusion::exec::execute_plan( @@ -415,22 +385,81 @@ fn select_fts_segments( Ok(selected) } -fn replace_match_query_exec( +fn segments_have_current_fragments( + dataset: &lance::Dataset, + segments: &[IndexMetadata], +) -> Result { + let current_fragment_ids = dataset + .get_fragments() + .into_iter() + .map(|fragment| { + u32::try_from(fragment.id()).map_err(|_| { + lance_core::Error::internal(format!( + "current fragment id {} exceeds the validated u32 FTS coverage range", + fragment.id() + )) + }) + }) + .collect::>>()?; + for segment in segments { + let fragment_bitmap = segment.fragment_bitmap.as_ref().ok_or_else(|| { + lance_core::Error::internal(format!( + "prepared FTS segment {} lost its validated fragment coverage", + segment.uuid + )) + })?; + if fragment_bitmap + .iter() + .any(|fragment_id| current_fragment_ids.contains(&fragment_id)) + { + return Ok(true); + } + } + Ok(false) +} + +#[derive(Default)] +struct PreparedFtsPlanRewriteCounts { + match_query_execs: usize, + flat_match_query_execs: usize, +} + +fn rewrite_prepared_fts_plan( plan: Arc, segments: &[IndexMetadata], scorer: &Arc, -) -> Result<(Arc, usize)> { + selected_segments_have_current_fragments: bool, +) -> Result<(Arc, PreparedFtsPlanRewriteCounts)> { + // Lance's ordinary FTS planner adds a flat-search branch for fragments not + // covered by the logical index. A prepared INDEX_ONLY scan must omit that + // branch, but using Scanner::with_fragments to do so would turn an + // otherwise unfiltered index search into a full row-id prefilter scan. + if plan.downcast_ref::().is_some() { + return Ok(( + Arc::new(EmptyExec::new(plan.schema())), + PreparedFtsPlanRewriteCounts { + match_query_execs: 0, + flat_match_query_execs: 1, + }, + )); + } + let children = plan.children(); - let mut replaced = 0; + let mut rewritten = PreparedFtsPlanRewriteCounts::default(); let rebuilt = if children.is_empty() { plan } else { let mut new_children = Vec::with_capacity(children.len()); for child in children { - let (new_child, child_replaced) = - replace_match_query_exec(Arc::clone(child), segments, scorer)?; + let (new_child, child_rewritten) = rewrite_prepared_fts_plan( + Arc::clone(child), + segments, + scorer, + selected_segments_have_current_fragments, + )?; new_children.push(new_child); - replaced += child_replaced; + rewritten.match_query_execs += child_rewritten.match_query_execs; + rewritten.flat_match_query_execs += child_rewritten.flat_match_query_execs; } plan.with_new_children(new_children).map_err(|error| { lance_core::Error::internal(format!( @@ -440,6 +469,10 @@ fn replace_match_query_exec( }; if let Some(exec) = rebuilt.downcast_ref::() { + rewritten.match_query_execs += 1; + if !selected_segments_have_current_fragments { + return Ok((Arc::new(EmptyExec::new(rebuilt.schema())), rewritten)); + } let replacement = MatchQueryExec::new_with_segments( Arc::clone(exec.dataset()), exec.query().clone(), @@ -448,9 +481,9 @@ fn replace_match_query_exec( segments.to_vec(), ) .with_base_scorer(Arc::clone(scorer)); - return Ok((Arc::new(replacement), replaced + 1)); + return Ok((Arc::new(replacement), rewritten)); } - Ok((rebuilt, replaced)) + Ok((rebuilt, rewritten)) } /// Type of a dynamically named scan metric. @@ -2118,6 +2151,9 @@ mod tests { use super::*; use crate::dataset::{lance_dataset_close, lance_dataset_open}; use crate::error::{lance_last_error_code, lance_last_error_message}; + use crate::fts_query::{ + LanceFtsCoverageMode, lance_dataset_prepare_fts_query, lance_fts_query_context_close, + }; use std::ffi::{CStr, CString}; use std::sync::atomic::{AtomicI32, AtomicUsize}; use std::sync::{Barrier, mpsc}; @@ -2125,6 +2161,9 @@ mod tests { use arrow_array::{Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; + use lance::index::DatasetIndexExt; + use lance::io::exec::PreFilterSource; + use lance_index::{IndexType, scalar::InvertedIndexParams}; /// Write a 3-row dataset to a tempdir, returning (tempdir, uri). fn create_test_dataset() -> (tempfile::TempDir, String) { @@ -2169,6 +2208,114 @@ mod tests { .store(true, Ordering::SeqCst); } + fn prepared_fts_plan_shape(plan: &Arc) -> (usize, usize, usize) { + let mut match_query_execs = 0; + let mut flat_match_query_execs = 0; + let mut filtered_row_id_prefilters = 0; + if let Some(exec) = plan.downcast_ref::() { + match_query_execs += 1; + if matches!(exec.prefilter_source(), PreFilterSource::FilteredRowIds(_)) { + filtered_row_id_prefilters += 1; + } + } + if plan.downcast_ref::().is_some() { + flat_match_query_execs += 1; + } + for child in plan.children() { + let (child_match, child_flat, child_filtered) = + prepared_fts_plan_shape(&Arc::clone(child)); + match_query_execs += child_match; + flat_match_query_execs += child_flat; + filtered_row_id_prefilters += child_filtered; + } + ( + match_query_execs, + flat_match_query_execs, + filtered_row_id_prefilters, + ) + } + + #[test] + fn prepared_fts_index_only_plan_does_not_scan_indexed_fragment_row_ids() { + let (_tmp, uri) = create_test_dataset(); + block_on(async { + let mut dataset = Dataset::open(&uri).await.unwrap(); + dataset + .create_index( + &["name"], + IndexType::Inverted, + None, + &InvertedIndexParams::default(), + false, + ) + .await + .unwrap(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![4])), + Arc::new(StringArray::from(vec!["a"])), + ], + ) + .unwrap(); + dataset + .append( + arrow::record_batch::RecordBatchIterator::new(vec![Ok(batch)], schema), + None, + ) + .await + .unwrap(); + }); + + let (dataset, scanner) = open_dataset_and_scanner(&uri); + let column = CString::new("name").unwrap(); + let query = CString::new("a").unwrap(); + let context = unsafe { + lance_dataset_prepare_fts_query( + dataset, + column.as_ptr(), + query.as_ptr(), + 0, + LanceFtsCoverageMode::IndexOnly as i32, + ) + }; + assert!(!context.is_null()); + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + + let prepared = unsafe { &*scanner }.build_scanner().unwrap(); + let distributed = prepared.distributed_fts.as_ref().unwrap(); + let segments = distributed.segments.clone(); + let scorer = Arc::clone(&distributed.context.scorer); + let plan = block_on(prepared.scanner.create_plan()).unwrap(); + assert_eq!( + prepared_fts_plan_shape(&plan), + (1, 1, 0), + "an unfiltered prepared FTS plan must not materialize selected fragment row IDs" + ); + + let has_current_fragments = + segments_have_current_fragments(&distributed.context.dataset, &segments).unwrap(); + let (rewritten, counts) = + rewrite_prepared_fts_plan(plan, &segments, &scorer, has_current_fragments).unwrap(); + assert_eq!(counts.match_query_execs, 1); + assert_eq!(counts.flat_match_query_execs, 1); + assert_eq!(prepared_fts_plan_shape(&rewritten), (1, 0, 0)); + + unsafe { + lance_scanner_close(scanner); + lance_fts_query_context_close(context); + lance_dataset_close(dataset); + } + } + /// Assert the pending thread-local error is `Panic` carrying the poison /// message; consumes it so the next assertion starts from a clean slate. fn assert_poison_error_pending() { diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 3627c4e..8805764 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -5943,6 +5943,85 @@ fn test_prepare_fts_query_index_only_allows_unindexed_fragment() { unsafe { lance_dataset_close(dataset) }; } +#[test] +fn test_prepared_fts_index_only_empty_segment_returns_empty_shard() { + use lance::index::DatasetIndexExt; + use lance_index::{IndexType, scalar::InvertedIndexParams}; + + let (_tmp, uri) = create_test_dataset(); + lance_c::runtime::block_on(async { + let mut dataset = Dataset::open(&uri).await.unwrap(); + let params = InvertedIndexParams::default(); + dataset + .create_index_builder(&["name"], IndexType::Inverted, ¶ms) + .name("empty_name_fts".to_string()) + .train(false) + .await + .unwrap(); + let segments = dataset + .load_indices_by_name("empty_name_fts") + .await + .unwrap(); + assert_eq!(segments.len(), 1); + assert!( + segments[0] + .fragment_bitmap + .as_ref() + .is_some_and(|fragment_bitmap| fragment_bitmap.is_empty()) + ); + }); + + let uri_c = c_str(&uri); + let column = c_str("name"); + let query = c_str("alice"); + let dataset = unsafe { lance_dataset_open(uri_c.as_ptr(), ptr::null(), 0) }; + let context = unsafe { + lance_dataset_prepare_fts_query( + dataset, + column.as_ptr(), + query.as_ptr(), + 0, + LanceFtsCoverageMode::IndexOnly as i32, + ) + }; + assert!(!context.is_null(), "{}", unsafe { + std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() + }); + let segment_uuids = load_fts_segment_uuids(&uri, "name"); + assert_eq!(segment_uuids.len(), 1); + + let scanner = unsafe { lance_scanner_new(dataset, ptr::null(), ptr::null()) }; + assert_eq!( + unsafe { lance_scanner_set_fts_query_context(scanner, context) }, + 0 + ); + assert_eq!( + unsafe { + lance_scanner_set_fts_index_segments( + scanner, + segment_uuids.as_ptr().cast::(), + segment_uuids.len(), + ) + }, + 0 + ); + + let mut stream = FFI_ArrowArrayStream::empty(); + assert_eq!( + unsafe { lance_scanner_to_arrow_stream(scanner, &mut stream) }, + 0, + "{}", + unsafe { std::ffi::CStr::from_ptr(lance_last_error_message()).to_string_lossy() } + ); + let reader = unsafe { ArrowArrayStreamReader::from_raw(&mut stream).unwrap() }; + let total_rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum(); + assert_eq!(total_rows, 0); + + unsafe { lance_scanner_close(scanner) }; + unsafe { lance_fts_query_context_close(context) }; + unsafe { lance_dataset_close(dataset) }; +} + #[test] fn test_prepared_fts_global_scorer_is_shared_across_segment_splits() { use lance::index::DatasetIndexExt;