From 4a76e6ef637a48627d7bee57a06bd7a8fe0a9b9e Mon Sep 17 00:00:00 2001 From: zhangstar333 Date: Wed, 30 Sep 2026 16:27:18 +0800 Subject: [PATCH] feat: Support vector post-filtering and scoped FTS prefiltering --- src/scanner.rs | 74 ++++++++++++++++++++++++++++++++------------------ 1 file changed, 47 insertions(+), 27 deletions(-) diff --git a/src/scanner.rs b/src/scanner.rs index 71d2dd2..cc44612 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -24,7 +24,7 @@ use lance_core::datatypes::BlobHandling; use lance_index::scalar::FullTextSearchQuery; use lance_index::vector::ApproxMode; use lance_io::stream::RecordBatchStream; -use lance_table::format::IndexMetadata; +use lance_table::format::{Fragment, IndexMetadata}; use uuid::Uuid; use crate::async_dispatcher::{self, LanceCallback}; @@ -421,7 +421,10 @@ impl LanceScanner { if self.include_deleted_rows { scanner.include_deleted_rows(); } - self.apply_fragment_filter(&mut scanner)?; + let apply_fragment_filter_after_nearest = self.nearest.is_some() && !self.prefilter; + if !apply_fragment_filter_after_nearest { + self.apply_fragment_filter(&mut scanner)?; + } if self.index_segments.is_some() && self.nearest.is_none() { return Err(lance_core::Error::invalid_input_source( "index_segments requires nearest() to be configured".into(), @@ -437,8 +440,8 @@ impl LanceScanner { "fragment_ids cannot be combined with an FTS query context; split the query by FTS index segment UUID instead".into(), )); } - // nearest() checks the current prefilter setting before accepting a - // fragment-scoped search. Enable it before installing the query. + // nearest() checks this setting at configuration time. In postfilter mode, defer + // applying fragment IDs until after nearest() is installed. if self.prefilter { scanner.prefilter(true); } @@ -495,6 +498,9 @@ impl LanceScanner { scanner.with_index_segments(segments.clone())?; } } + if apply_fragment_filter_after_nearest { + self.apply_fragment_filter(&mut scanner)?; + } if let Some(fts) = &self.fts_query { scanner.full_text_search(fts.clone())?; } @@ -505,6 +511,10 @@ impl LanceScanner { Some(PreparedFtsExecution { context: Arc::clone(context), segments, + scope_prefilter_to_fts_segments: self.prefilter + && (self.filter.is_some() + || self.substrait_filter.is_some() + || !self.additional_sql_filters.is_empty()), batch_size: self.batch_size, scan_statistics_callback: self.scan_statistics_callback.clone(), }) @@ -558,6 +568,7 @@ impl LanceScanner { struct PreparedFtsExecution { context: Arc, segments: Vec, + scope_prefilter_to_fts_segments: bool, batch_size: Option, scan_statistics_callback: Option, } @@ -596,11 +607,21 @@ impl PreparedScanner { let Some(distributed_fts) = self.distributed_fts else { return self.scanner.try_into_stream().await; }; - let plan = self.scanner.create_plan().await?; - let selected_segments_have_current_fragments = segments_have_current_fragments( + let selected_fragments = selected_current_fts_fragments( &distributed_fts.context.dataset, &distributed_fts.segments, )?; + let selected_segments_have_current_fragments = !selected_fragments.is_empty(); + let mut scanner = self.scanner; + if distributed_fts.scope_prefilter_to_fts_segments + && selected_segments_have_current_fragments + { + // The scanner is already split by the selected FTS segment(s). Applying the same + // fragment scope before plan creation lets Lance restrict scalar-index segment loads + // for the TVF prefilter without changing unfiltered FTS scans. + scanner.with_fragments(selected_fragments); + } + let plan = scanner.create_plan().await?; let (plan, rewritten) = rewrite_prepared_fts_plan( plan, &distributed_fts.segments, @@ -661,22 +682,11 @@ fn select_fts_segments( Ok(selected) } -fn segments_have_current_fragments( +fn selected_current_fts_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::>>()?; +) -> Result> { + 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!( @@ -684,14 +694,22 @@ fn segments_have_current_fragments( segment.uuid )) })?; - if fragment_bitmap - .iter() - .any(|fragment_id| current_fragment_ids.contains(&fragment_id)) - { - return Ok(true); + selected_fragment_ids.extend(fragment_bitmap.iter()); + } + + let mut selected_fragments = Vec::new(); + for fragment in dataset.get_fragments() { + let fragment_id = u32::try_from(fragment.id()).map_err(|_| { + lance_core::Error::internal(format!( + "current fragment id {} exceeds the validated u32 FTS coverage range", + fragment.id() + )) + })?; + if selected_fragment_ids.contains(&fragment_id) { + selected_fragments.push(fragment.metadata().clone()); } } - Ok(false) + Ok(selected_fragments) } #[derive(Default)] @@ -3357,7 +3375,9 @@ mod tests { ); let has_current_fragments = - segments_have_current_fragments(&distributed.context.dataset, &segments).unwrap(); + !selected_current_fts_fragments(&distributed.context.dataset, &segments) + .unwrap() + .is_empty(); let (rewritten, counts) = rewrite_prepared_fts_plan( plan, &segments,