diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index da8348d57cc..8acc5d02601 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -76,7 +76,10 @@ use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, DIST_COL, Query use lance_index::{metrics::NoOpMetricsCollector, scalar::inverted::FTS_SCHEMA}; use lance_io::stream::RecordBatchStream; use lance_linalg::distance::MetricType; -use lance_select::{IndexExprResult, RowAddrMask, RowAddrTreeMap}; +use lance_select::IndexExprResult; +// Re-exported so callers of `Scanner::with_row_addr_prefilter` can name the mask +// type without depending on `lance-select` directly. +pub use lance_select::{RowAddrMask, RowAddrTreeMap}; use lance_table::format::{Fragment, IndexMetadata}; use roaring::RoaringBitmap; use tracing::{Span, info_span, instrument}; @@ -101,7 +104,7 @@ use crate::io::exec::knn::MultivectorScoringExec; use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec}; use crate::io::exec::{ AddRowAddrExec, FilterPlan as ExprFilterPlan, KNNVectorDistanceExec, LancePushdownScanExec, - LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec, + LanceScanExec, Planner, PreFilterSource, RowAddrMaskFilterExec, ScanConfig, TakeExec, knn::{ KnnBatchParams, QUERY_INDEX_COL, knn_empty_result_schema, new_knn_exec, query_index_field, }, @@ -736,6 +739,14 @@ pub struct Scanner { /// If true then the filter will be applied before an index scan prefilter: bool, + /// Optional external allow/block mask keyed in `_rowid` space. On a vector + /// search it is combined with the index-side prefilter and applied to the + /// flat branch for fragments not covered by the index; on a plain scan it is + /// the row source (see `use_external_mask`). Held behind an Arc so cloning it + /// into the ANN sub-plans and the flat-branch filter is cheap regardless of + /// mask size. + external_row_mask: Option>, + /// Materialization style controls when columns are fetched materialization_style: MaterializationStyle, @@ -1034,6 +1045,7 @@ impl Scanner { projection_plan, blob_handling: BlobHandling::default(), prefilter: false, + external_row_mask: None, materialization_style: MaterializationStyle::Heuristic, filter: LanceFilter::default(), full_text_query: None, @@ -1198,6 +1210,47 @@ impl Scanner { self } + /// Set an external [`RowAddrMask`] allow/block prefilter. + /// + /// Build the mask with [`RowAddrMask::from_allowed`] to keep only the listed + /// rows or [`RowAddrMask::from_block`] to drop them. On a vector + /// ([`nearest`](Self::nearest)) search the mask is combined with any + /// filter-derived prefilter on the index branch and applied to the flat + /// branch for fragments not covered by the vector index. On a + /// [`full_text_search`](Self::full_text_search) (match or phrase query) the + /// mask is combined into the FTS prefilter so BM25 top-k is computed over + /// masked rows, and the flat branch that scores unindexed fragments + /// (plan_flat_match_query) is masked with RowAddrMaskFilterExec. On a plain + /// scan the mask is used directly as the row source, with any + /// [`filter`](Self::filter) applied as a refine on top. + /// + /// The mask is keyed in the dataset's `_rowid` space, so build it from the + /// same dataset you query. That space is the row address when stable row ids + /// are disabled and the stable row id when they are enabled; both are handled + /// (index prefilter and filtered read branch on `uses_stable_row_ids`), so no + /// caller-side translation is needed either way. + /// + /// # Example + /// + /// ```no_run + /// # use lance::dataset::Dataset; + /// # async fn example(dataset: &Dataset) -> lance::Result<()> { + /// use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap}; + /// + /// // Restrict the scan to rows whose _rowid is 0, 2, or 4. + /// let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([0u64, 2, 4])); + /// let mut scanner = dataset.scan(); + /// scanner.with_row_addr_prefilter(mask); + /// let batch = scanner.try_into_batch().await?; + /// # let _ = batch; + /// # Ok(()) + /// # } + /// ``` + pub fn with_row_addr_prefilter(&mut self, mask: RowAddrMask) -> &mut Self { + self.external_row_mask = Some(Arc::new(mask)); + self + } + /// Set the callback to be called after the scan with summary statistics pub fn scan_stats_callback(&mut self, callback: ExecutionStatsCallback) -> &mut Self { self.scan_stats_callback = Some(callback); @@ -2895,6 +2948,32 @@ impl Scanner { } } + // A plain-scan external row mask is fed as the FilteredReadExec row source so + // only masked rows are read, with any SQL filter applied as a refine on top. + // Vector and full-text searches apply the mask via their own prefilter paths + // (KNN external_mask / FTS build_prefilter), so this plain-scan source is + // scoped to scans that are neither. FTS in particular has nearest.is_none(), + // so excluding it here keeps the FTS prefilter's own filtered read unmasked. + fn use_external_mask(&self) -> bool { + self.nearest.is_none() && self.full_text_query.is_none() && self.external_row_mask.is_some() + } + + // The filter plan actually handed to the filtered read. With an external mask + // active the mask is the row source, so any SQL filter is demoted to a refine + // on top of it; otherwise the plan is used as-is. Projection and scan-range + // planning must be done against this, not the raw filter_plan, so refine + // columns are retained and limit/offset is not pushed down before masking. + fn effective_filter_plan(&self, filter_plan: &ExprFilterPlan) -> ExprFilterPlan { + if self.use_external_mask() { + match filter_plan.full_expr.clone() { + Some(expr) => ExprFilterPlan::new_refine_only(expr), + None => ExprFilterPlan::default(), + } + } else { + filter_plan.clone() + } + } + // Helper function for filtered_read // // Do not call this directly, use filtered_read instead @@ -2906,8 +2985,11 @@ impl Scanner { fragments: Option>>, scan_range: Option>, ) -> Result> { + let use_external_mask = self.use_external_mask(); + let effective_filter = self.effective_filter_plan(filter_plan); + let mut read_options = FilteredReadOptions::basic_full_read(&self.dataset) - .with_filter_plan(filter_plan.clone()) + .with_filter_plan(effective_filter) .with_projection(projection); if let Some(fragments) = fragments { @@ -2948,13 +3030,16 @@ impl Scanner { } let result_format = self.index_expr_result_format(); - let index_input = filter_plan.index_query.clone().map(|index_query| { - Arc::new(ScalarIndexExec::new( - self.dataset.clone(), - index_query, - result_format, - )) as Arc - }); + let index_input = match self.external_row_mask.as_deref() { + Some(mask) if use_external_mask => Some(self.mask_as_index_input(mask)?), + _ => filter_plan.index_query.clone().map(|index_query| { + Arc::new(ScalarIndexExec::new( + self.dataset.clone(), + index_query, + result_format, + )) as Arc + }), + }; Ok(Arc::new(FilteredReadExec::try_new( self.dataset.clone(), @@ -2977,6 +3062,19 @@ impl Scanner { ) -> Result { // Use legacy path if dataset uses legacy storage format if self.dataset.is_legacy_storage() { + // The plain-scan mask path lives in new_filtered_read; legacy_filtered_read + // has no equivalent, so a masked plain scan here would silently drop the + // mask and return every row. Fail loudly instead. Vector and full-text + // searches apply the mask via their own prefilter paths (ANN prefilter / + // FTS build_prefilter) plus the RowAddrMaskFilterExec flat wrap, so they + // are unaffected -- use_external_mask() is false for them. + if self.use_external_mask() { + return Err(Error::not_supported( + "with_row_addr_prefilter is not supported for plain scans on \ + legacy-storage datasets" + .to_string(), + )); + } self.legacy_filtered_read( filter_plan, projection, @@ -3005,10 +3103,12 @@ impl Scanner { } } - fn u64s_as_take_input(&self, u64s: Vec) -> Result> { - let row_addrs = RowAddrTreeMap::from_iter(u64s); - let row_addr_mask = RowAddrMask::from_allowed(row_addrs); - let index_result = IndexExprResult::exact(row_addr_mask); + // Wrap an evaluated index result as a one-shot index input for + // FilteredReadExec. + fn index_result_as_index_input( + &self, + index_result: IndexExprResult, + ) -> Result> { let fragments_covered = self.dataset.fragment_bitmap.as_ref().clone(); let format = self.index_expr_result_format(); let batch = index_result.serialize(&fragments_covered, format)?; @@ -3018,6 +3118,17 @@ impl Scanner { Ok(Arc::new(OneShotExec::new(stream))) } + // Wrap an external row-address mask as a one-shot index input for + // FilteredReadExec, so a plain scan reads only the masked rows. + fn mask_as_index_input(&self, mask: &RowAddrMask) -> Result> { + self.index_result_as_index_input(IndexExprResult::exact(mask.clone())) + } + + fn u64s_as_take_input(&self, u64s: Vec) -> Result> { + let row_addr_mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter(u64s)); + self.index_result_as_index_input(IndexExprResult::exact(row_addr_mask)) + } + async fn take_source(&self, take_op: TakeOperation) -> Result> { // We generally assume that late materialization does not make sense for take operations // so we can just use the physical projection @@ -3071,11 +3182,15 @@ impl Scanner { self.projection_plan.physical_projection.clone() }; - let mut projection = if filter_plan.has_refine() { + // Plan against the effective filter: with an external mask the SQL filter + // becomes a refine, so its columns must be retained even when the original + // plan resolved to an exact scalar-index query (has_refine() == false). + let effective_filter = self.effective_filter_plan(filter_plan); + let mut projection = if effective_filter.has_refine() { // If the filter plan has two steps (a scalar indexed portion and a refine portion) then // it makes sense to grab cheap columns during the first step to avoid taking them for // the second step. - self.calc_eager_projection(filter_plan, &effective_projection)? + self.calc_eager_projection(&effective_filter, &effective_projection)? .with_row_id() } else { // If the filter plan only has one step then we just do a filtered read of all the @@ -3089,7 +3204,11 @@ impl Scanner { projection.with_row_addr = true; } - let scan_range = if filter_plan.is_empty() { + // An external mask is applied as the row source inside new_filtered_read, so + // limit/offset must not be pushed down as a pre-mask range (that would limit + // rows before masking). Leaving scan_range None keeps limit_pushed_down false + // so the limit is applied by a node above the masked source instead. + let scan_range = if filter_plan.is_empty() && !self.use_external_mask() { log::trace!("pushing scan_range into filtered_read"); self.get_scan_range(filter_plan).await? } else { @@ -3519,12 +3638,15 @@ impl Scanner { .to_string())); } - Ok(Arc::new(PhraseQueryExec::new( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - ))) + Ok(Arc::new( + PhraseQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + ) + .with_external_mask(self.external_row_mask.clone()), + )) } async fn plan_match_query( @@ -3571,12 +3693,15 @@ impl Scanner { } // Mixed case: use index + flat search for unindexed - let match_plan: Arc = Arc::new(MatchQueryExec::new( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - )); + let match_plan: Arc = Arc::new( + MatchQueryExec::new( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + ) + .with_external_mask(self.external_row_mask.clone()), + ); if self.fast_search || unindexed_fragments.is_empty() { (Some(match_plan), None) @@ -3666,12 +3791,19 @@ impl Scanner { } plan = self.ensure_column_alias(plan, &column)?; - let flat_match_plan = Arc::new(FlatMatchQueryExec::new( + let flat_match_plan: Arc = Arc::new(FlatMatchQueryExec::new( self.dataset.clone(), query.clone(), params.clone(), plan, )); + // Unindexed fragments never reach the index-side prefilter, so apply the + // external row-address mask to the flat FTS results here (mirrors the ANN + // flat branch). Applied before the caller's top-k so masked-out rows do + // not consume result slots. + if let Some(mask) = self.external_row_mask.clone() { + return Ok(Arc::new(RowAddrMaskFilterExec::new(flat_match_plan, mask))); + } Ok(flat_match_plan) } @@ -3903,6 +4035,11 @@ impl Scanner { if let Some(refine_expr) = &filter_plan.refine_expr { plan = Arc::new(LanceFilterExec::try_new(refine_expr.clone(), plan)?); } + // The flat branch never reaches the index-side prefilter, so apply + // the external row-address mask here against the scanned _rowid. + if let Some(mask) = self.external_row_mask.clone() { + plan = Arc::new(RowAddrMaskFilterExec::new(plan, mask)); + } Ok(self.flat_knn(plan, &q)?) } } @@ -4052,6 +4189,11 @@ impl Scanner { // If there is a prefilter we need to manually apply it to the new data scan_node = Arc::new(LanceFilterExec::try_new(expr.clone(), scan_node)?); } + // Appended fragments are not covered by the index, so the external + // row-address mask must be applied to them here. + if let Some(mask) = self.external_row_mask.clone() { + scan_node = Arc::new(RowAddrMaskFilterExec::new(scan_node, mask)); + } // first we do flat search on just the new data let topk_appended = self.flat_knn(scan_node, &q)?; @@ -4702,7 +4844,13 @@ impl Scanner { let prefilter_source = self .prefilter_source(filter_plan, self.get_indexed_frags(index)) .await?; - let inner_fanout_search = new_knn_exec(self.dataset.clone(), index, q, prefilter_source)?; + let inner_fanout_search = new_knn_exec( + self.dataset.clone(), + index, + q, + prefilter_source, + self.external_row_mask.clone(), + )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, inner_fanout_search.schema().as_ref())?, options: SortOptions { @@ -4761,6 +4909,7 @@ impl Scanner { index, &query, prefilter_source.clone(), + self.external_row_mask.clone(), )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, ann_node.schema().as_ref())?, @@ -5535,6 +5684,259 @@ mod test { } } + fn batch_row_ids(batch: &RecordBatch) -> Vec { + batch + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values() + .to_vec() + } + + #[rstest] + #[case::without_stable_row_ids(false)] + #[case::with_stable_row_ids(true)] + #[tokio::test] + async fn row_addr_mask_plain_scan_allow_block_refine(#[case] stable_row_ids: bool) { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, stable_row_ids) + .await + .unwrap(); + let ds = &test_ds.dataset; + + let mut scan = ds.scan(); + scan.with_row_id(); + let all_ids = batch_row_ids(&scan.try_into_batch().await.unwrap()); + let all_set: BTreeSet = all_ids.iter().copied().collect(); + let allow: Vec = all_ids.iter().copied().step_by(2).collect(); + let allow_set: BTreeSet = allow.iter().copied().collect(); + + // Allow-mask plain scan returns exactly the allowed rows. + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + allow.iter().copied(), + ))); + scan.with_row_id(); + let got: BTreeSet = batch_row_ids(&scan.try_into_batch().await.unwrap()) + .into_iter() + .collect(); + assert_eq!(got, allow_set); + + // Block-mask plain scan returns every row except the blocked ones, which + // also exercises FilteredReadExec index-input serialization of a BlockList. + let block: Vec = all_ids.iter().copied().step_by(3).collect(); + let block_set: BTreeSet = block.iter().copied().collect(); + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter( + block.iter().copied(), + ))); + scan.with_row_id(); + let got: BTreeSet = batch_row_ids(&scan.try_into_batch().await.unwrap()) + .into_iter() + .collect(); + let expected: BTreeSet = all_set.difference(&block_set).copied().collect(); + assert_eq!(got, expected); + + // With a SQL refine, the result is the allowed rows that also match the filter. + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + allow.iter().copied(), + ))); + scan.filter("i >= 200").unwrap(); + scan.project(&["i"]).unwrap(); + scan.with_row_id(); + let refined = scan.try_into_batch().await.unwrap(); + let refined_ids: BTreeSet = batch_row_ids(&refined).into_iter().collect(); + assert!(refined_ids.is_subset(&allow_set) && !refined_ids.is_empty()); + let is = refined + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert!(is.values().iter().all(|v| *v >= 200)); + } + + #[tokio::test] + async fn row_addr_mask_plain_scan_rejected_on_legacy() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Legacy, false) + .await + .unwrap(); + let ds = &test_ds.dataset; + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([0u64]))); + let Err(err) = scan.try_into_stream().await else { + panic!("expected legacy-storage masked plain scan to be rejected"); + }; + assert!( + err.to_string().contains("legacy-storage"), + "unexpected: {err}" + ); + } + + #[rstest] + #[case::without_stable_row_ids(false)] + #[case::with_stable_row_ids(true)] + #[tokio::test] + async fn row_addr_mask_ann_search_only_allowed(#[case] stable_row_ids: bool) { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, stable_row_ids) + .await + .unwrap(); + test_ds.make_vector_index().await.unwrap(); + // Append after indexing so the appended fragment is unindexed (flat branch). + test_ds.append_new_data().await.unwrap(); + let ds = &test_ds.dataset; + + let mut scan = ds.scan(); + scan.with_row_id(); + let all_ids = batch_row_ids(&scan.try_into_batch().await.unwrap()); + let allow: Vec = all_ids.iter().copied().step_by(3).collect(); + let allow_set: BTreeSet = allow.iter().copied().collect(); + + let key: Float32Array = (0..32).map(|v| v as f32).collect(); + let mut scan = ds.scan(); + scan.nearest("vec", &key, 15).unwrap(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + allow.iter().copied(), + ))); + scan.with_row_id(); + let got = batch_row_ids(&scan.try_into_batch().await.unwrap()); + assert!(!got.is_empty()); + for id in got { + assert!( + allow_set.contains(&id), + "returned _rowid {id} not in allowlist" + ); + } + } + + #[tokio::test] + async fn row_addr_mask_plain_scan_with_limit() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + let ds = &test_ds.dataset; + + let mut scan = ds.scan(); + scan.with_row_id(); + let all_ids = batch_row_ids(&scan.try_into_batch().await.unwrap()); + let allow: Vec = all_ids.iter().copied().step_by(2).collect(); + let allow_set: BTreeSet = allow.iter().copied().collect(); + + // limit must apply AFTER masking: 5 rows, all from the allowlist. + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + allow.iter().copied(), + ))); + scan.limit(Some(5), None).unwrap(); + scan.with_row_id(); + let got = batch_row_ids(&scan.try_into_batch().await.unwrap()); + assert_eq!(got.len(), 5, "masked limit should yield 5 masked rows"); + for id in &got { + assert!(allow_set.contains(id), "returned {id} not allowed"); + } + } + + #[tokio::test] + async fn row_addr_mask_plain_scan_filter_unprojected_column() { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + let ds = &test_ds.dataset; + + let mut scan = ds.scan(); + scan.with_row_id(); + let all_ids = batch_row_ids(&scan.try_into_batch().await.unwrap()); + + // Allow everything; filter on `i` but project only `s` (unrelated column). + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + all_ids.iter().copied(), + ))); + scan.filter("i >= 200").unwrap(); + scan.project(&["s"]).unwrap(); + let out = scan.try_into_batch().await.unwrap(); + assert_eq!(out.num_rows(), 200, "expected 200 rows with i>=200"); + } + + #[tokio::test] + async fn row_addr_mask_plain_scan_exact_index_filter_unprojected_column() { + // A scalar index on `i` turns `i >= 200` into an exact index query with no + // refine. Under an external mask that predicate is demoted to a refine over + // the masked rows, so `i` must still be projected for the read even though + // the user only asked for `s`. + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_scalar_index().await.unwrap(); + let ds = &test_ds.dataset; + + let mut scan = ds.scan(); + scan.with_row_id(); + let all_ids = batch_row_ids(&scan.try_into_batch().await.unwrap()); + + let mut scan = ds.scan(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + all_ids.iter().copied(), + ))); + scan.filter("i >= 200").unwrap(); + scan.project(&["s"]).unwrap(); + let out = scan.try_into_batch().await.unwrap(); + assert_eq!(out.num_rows(), 200, "expected 200 rows with i>=200"); + } + + #[tokio::test] + async fn row_addr_mask_fts_search_only_allowed() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_fts_index().await.unwrap(); + // Re-append the low-i rows AFTER indexing so token "4" matches both an + // indexed row (index prefilter path) and an unindexed one (flat FTS branch). + test_ds.append_data_with_range(0, 10).await.unwrap(); + let ds = &test_ds.dataset; + + // Baseline: the row ids an unmasked FTS query matches. + let mut scan = ds.scan(); + scan.full_text_search(FullTextSearchQuery::new("4".into())) + .unwrap(); + scan.with_row_id(); + let base_ids = batch_row_ids(&scan.try_into_batch().await.unwrap()); + let base_set: BTreeSet = base_ids.iter().copied().collect(); + assert!( + base_ids.len() >= 2, + "expected indexed + unindexed matches for token 4, got {base_ids:?}" + ); + + // Allow only every other matching row; the mask must prefilter BM25 so the + // result is exactly the allowed subset of the baseline matches. + let allow: Vec = base_ids.iter().copied().step_by(2).collect(); + let allow_set: BTreeSet = allow.iter().copied().collect(); + + let mut scan = ds.scan(); + scan.full_text_search(FullTextSearchQuery::new("4".into())) + .unwrap(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter( + allow.iter().copied(), + ))); + scan.with_row_id(); + let got: BTreeSet = batch_row_ids(&scan.try_into_batch().await.unwrap()) + .into_iter() + .collect(); + let expected: BTreeSet = base_set.intersection(&allow_set).copied().collect(); + assert_eq!(got, expected, "masked FTS must return allowed matches only"); + assert!(!got.is_empty()); + + // Block every match -> empty, proving the mask actually filters FTS results + // on both the indexed and flat branches. + let mut scan = ds.scan(); + scan.full_text_search(FullTextSearchQuery::new("4".into())) + .unwrap(); + scan.with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter( + base_ids.iter().copied(), + ))); + scan.with_row_id(); + let blocked = batch_row_ids(&scan.try_into_batch().await.unwrap()); + assert!(blocked.is_empty(), "block-mask must drop all FTS matches"); + } + #[tokio::test] async fn test_strict_batch_size() { let dataset = lance_datagen::gen_batch() diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index a477d60d56d..54a0c310b45 100644 --- a/rust/lance/src/io/exec.rs +++ b/rust/lance/src/io/exec.rs @@ -18,6 +18,7 @@ pub(crate) mod knn; mod optimizer; mod projection; mod pushdown_scan; +pub(crate) mod row_addr_mask; mod rowids; pub mod scalar_index; mod scan; @@ -35,6 +36,7 @@ pub use lance_index::scalar::expression::FilterPlan; pub use optimizer::get_physical_optimizer; pub use projection::project; pub use pushdown_scan::{LancePushdownScanExec, ScanConfig}; +pub use row_addr_mask::RowAddrMaskFilterExec; pub use rowids::{AddRowAddrExec, AddRowOffsetExec}; pub use scan::{LanceScanConfig, LanceScanExec}; pub use take::TakeExec; diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 414ce601763..ff10d96f336 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -2314,7 +2314,12 @@ impl ExecutionPlan for FilteredReadExec { let mut updated_options = self.options.clone(); if self.options.full_filter.is_none() && self.options.refine_filter.is_none() { - if self.options.scan_range_before_filter.is_some() { + // A before-filter range trims raw scan positions, which is only valid for + // an unindexed full scan. With an index_input (e.g. an external row mask or + // a scalar-index result) the rows are selected by that input, so a pre-range + // would apply before selection and keep the wrong rows; leave the limit to a + // node above the read instead. + if self.options.scan_range_before_filter.is_some() || self.index_input.is_some() { return None; } updated_options.scan_range_before_filter = Some(0..(limit as u64)); @@ -3322,6 +3327,35 @@ mod tests { let result = plan.with_fetch(None); assert!(result.is_none()); } + + // Case 7: index_input present with no filter (the external-row-mask + // plain-scan shape) - with_fetch must reject before-filter pushdown, since + // the index_input selects the rows and a raw before-filter range would trim + // scan positions before that selection. + { + // Build a real scalar-index input, then attach it to options that carry + // no filter of their own. + let index_filter_plan = fixture.filter_plan("fully_indexed < 200", false).await; + let index_input = fixture + .index_input(&base_options.clone().with_filter_plan(index_filter_plan)) + .await; + assert!(index_input.is_some(), "expected a scalar-index input"); + + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + base_options.clone(), + index_input, + ) + .unwrap(); + assert!(plan.index_input().is_some()); + assert!(plan.options().full_filter.is_none() && plan.options().refine_filter.is_none()); + + let result = plan.with_fetch(Some(100)); + assert!( + result.is_none(), + "with_fetch must reject before-filter pushdown when index_input is present" + ); + } } #[tokio::test] diff --git a/rust/lance/src/io/exec/fts.rs b/rust/lance/src/io/exec/fts.rs index 2c1859211f6..084afe1382d 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; @@ -238,6 +239,9 @@ pub struct MatchQueryExec { /// When set, `execute()` skips `load_segments` and searches exactly these /// segments. preset_segments: Option>, + /// Optional external row-address mask combined (logical AND) with the BM25 + /// prefilter so only masked rows are scored (see [`Self::with_external_mask`]). + external_mask: Option>, properties: Arc, metrics: ExecutionPlanMetricsSet, @@ -296,6 +300,7 @@ impl MatchQueryExec { prefilter_source, base_scorer: None, preset_segments: None, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -331,6 +336,7 @@ impl MatchQueryExec { prefilter_source, base_scorer: None, preset_segments: Some(segments), + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -353,6 +359,15 @@ impl MatchQueryExec { self } + /// Restrict BM25 scoring to rows selected by an external row-address mask. + /// The mask is combined (logical AND) with the prefilter (see + /// [`build_prefilter`]) so top-k is computed over masked rows only. No-op + /// when `mask` is `None`. + pub fn with_external_mask(mut self, mask: Option>) -> Self { + self.external_mask = mask; + self + } + pub fn query(&self) -> &MatchQuery { &self.query } @@ -418,6 +433,7 @@ impl ExecutionPlan for MatchQueryExec { prefilter_source: PreFilterSource::None, base_scorer: self.base_scorer.clone(), preset_segments: self.preset_segments.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -445,6 +461,7 @@ impl ExecutionPlan for MatchQueryExec { prefilter_source, base_scorer: self.base_scorer.clone(), preset_segments: self.preset_segments.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -468,6 +485,7 @@ impl ExecutionPlan for MatchQueryExec { let params = self.params.clone(); let ds = self.dataset.clone(); let prefilter_source = self.prefilter_source.clone(); + let external_mask = self.external_mask.clone(); let preset_base_scorer = self.base_scorer.clone(); let preset_segments = self.preset_segments.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); @@ -490,8 +508,14 @@ impl ExecutionPlan for MatchQueryExec { let indices = open_fts_segments(&ds, &column, &segments, &metrics.index_metrics).await?; - let mut pre_filter = - build_prefilter(context.clone(), partition, &prefilter_source, ds, &segments)?; + let mut pre_filter = build_prefilter( + context.clone(), + partition, + &prefilter_source, + ds, + &segments, + external_mask, + )?; let deleted_fragments = indices .iter() @@ -1155,6 +1179,9 @@ pub struct PhraseQueryExec { /// Optional pre-resolved segment list. See /// [`MatchQueryExec::new_with_segments`]. preset_segments: Option>, + /// Optional external row-address mask combined (logical AND) with the BM25 + /// prefilter so only masked rows are scored (see [`MatchQueryExec::with_external_mask`]). + external_mask: Option>, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -1204,6 +1231,7 @@ impl PhraseQueryExec { prefilter_source, base_scorer: None, preset_segments: None, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -1232,6 +1260,7 @@ impl PhraseQueryExec { prefilter_source, base_scorer: None, preset_segments: Some(segments), + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -1243,6 +1272,12 @@ impl PhraseQueryExec { self } + /// See [`MatchQueryExec::with_external_mask`]. + pub fn with_external_mask(mut self, mask: Option>) -> Self { + self.external_mask = mask; + self + } + pub fn query(&self) -> &PhraseQuery { &self.query } @@ -1301,6 +1336,7 @@ impl ExecutionPlan for PhraseQueryExec { prefilter_source: PreFilterSource::None, base_scorer: self.base_scorer.clone(), preset_segments: self.preset_segments.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), }, @@ -1326,6 +1362,7 @@ impl ExecutionPlan for PhraseQueryExec { prefilter_source, base_scorer: self.base_scorer.clone(), preset_segments: self.preset_segments.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -1349,6 +1386,7 @@ impl ExecutionPlan for PhraseQueryExec { let params = self.params.clone(); let ds = self.dataset.clone(); let prefilter_source = self.prefilter_source.clone(); + let external_mask = self.external_mask.clone(); let preset_base_scorer = self.base_scorer.clone(); let preset_segments = self.preset_segments.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); @@ -1371,8 +1409,14 @@ impl ExecutionPlan for PhraseQueryExec { let indices = open_fts_segments(&ds, &column, &segments, &metrics.index_metrics).await?; - let mut pre_filter = - build_prefilter(context.clone(), partition, &prefilter_source, ds, &segments)?; + let mut pre_filter = build_prefilter( + context.clone(), + partition, + &prefilter_source, + ds, + &segments, + external_mask, + )?; let deleted_fragments = indices .iter() diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index fde0289621d..565a6ba3823 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -56,6 +56,7 @@ use lance_index::vector::{ }; use lance_linalg::distance::DistanceType; use lance_linalg::kernels::normalize_arrow; +use lance_select::RowAddrMask; use lance_table::format::IndexMetadata; use tokio::sync::Notify; use uuid::Uuid; @@ -67,6 +68,7 @@ use crate::index::vector::utils::{get_vector_type, validate_distance_type_for}; use crate::{Error, Result}; use lance_arrow::*; +use super::row_addr_mask::MaskAndLoader; use super::utils::{ FilteredRowIdsToPrefilter, IndexMetrics, InstrumentedRecordBatchStreamAdapter, PreFilterSource, SelectionVectorToPrefilter, @@ -1104,6 +1106,7 @@ pub fn new_knn_exec( indices: &[IndexMetadata], query: &Query, prefilter_source: PreFilterSource, + external_mask: Option>, ) -> Result> { let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), @@ -1117,6 +1120,7 @@ pub fn new_knn_exec( indices.to_vec(), query.clone(), prefilter_source, + external_mask, )?; Ok(Arc::new(sub_index)) @@ -1377,6 +1381,10 @@ pub struct ANNIvfSubIndexExec { /// Prefiltering input prefilter_source: PreFilterSource, + /// Optional external row-address allow/block mask, combined with the + /// prefilter using logical AND. + external_mask: Option>, + /// Datafusion Plan Properties properties: Arc, @@ -1390,6 +1398,7 @@ impl ANNIvfSubIndexExec { indices: Vec, query: Query, prefilter_source: PreFilterSource, + external_mask: Option>, ) -> Result { if input.schema().field_with_name(PART_ID_COLUMN).is_err() { return Err(Error::index(format!( @@ -1409,6 +1418,7 @@ impl ANNIvfSubIndexExec { indices, query, prefilter_source, + external_mask, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -1880,6 +1890,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { indices: self.indices.clone(), query: self.query.clone(), prefilter_source, + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -1960,6 +1971,14 @@ impl ExecutionPlan for ANNIvfSubIndexExec { PreFilterSource::None => None, }; + // AND the external row-address mask into whatever the filter produced. + let prefilter_loader = match self.external_mask.clone() { + Some(mask) => { + Some(Box::new(MaskAndLoader::new(mask, prefilter_loader)) as Box) + } + None => prefilter_loader, + }; + let pre_filter = Arc::new(DatasetPreFilter::new( ds.clone(), &indices, diff --git a/rust/lance/src/io/exec/row_addr_mask.rs b/rust/lance/src/io/exec/row_addr_mask.rs new file mode 100644 index 00000000000..eb7059098bc --- /dev/null +++ b/rust/lance/src/io/exec/row_addr_mask.rs @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! RowAddrMask prefilter wiring for vector search. +//! +//! An externally supplied [`RowAddrMask`] is applied to a KNN search through two +//! pieces, because the two search branches consume a prefilter differently: +//! - [`MaskAndLoader`] folds the mask into the index-side prefilter loader +//! (ANN / IVF branch). The mask, any filter-derived selection vector, and +//! the deletion vector are all combined (logical AND) by DatasetPreFilter. +//! - [`RowAddrMaskFilterExec`] applies the mask to the flat-KNN branch, which +//! scans fragments not covered by the vector index and so never reaches the +//! index-side prefilter. + +use std::sync::Arc; + +use arrow::datatypes::UInt64Type; +use arrow_array::cast::AsArray; +use arrow_array::{BooleanArray, RecordBatch}; +use async_trait::async_trait; +use datafusion::error::DataFusionError; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; +use futures::StreamExt; +use lance_core::error::DataFusionResult; +use lance_core::{ROW_ID, Result}; +use lance_index::prefilter::FilterLoader; +use lance_select::RowAddrMask; + +/// FilterLoader that combines an external RowAddrMask (logical AND) with an +/// optional inner loader. +/// +/// With an inner loader present the two masks are intersected; otherwise the +/// external mask is used alone. DatasetPreFilter later intersects the result +/// with the dataset deletion vector. +pub struct MaskAndLoader { + mask: Arc, + inner: Option>, +} + +impl MaskAndLoader { + pub fn new(mask: Arc, inner: Option>) -> Self { + Self { mask, inner } + } +} + +#[async_trait] +impl FilterLoader for MaskAndLoader { + async fn load(self: Box) -> Result { + match self.inner { + Some(inner) => Ok(Arc::unwrap_or_clone(self.mask) & inner.load().await?), + None => Ok(Arc::unwrap_or_clone(self.mask)), + } + } +} + +/// Execution node that drops rows whose `_rowid` is not selected by `mask`. +/// +/// The key is read from the `_rowid` column, and `mask` is keyed in that same +/// `_rowid` space, so this is consistent whether stable row ids are enabled (the +/// value is the stable row id) or disabled (it is the row address). Schema and +/// ordering are preserved; only the row count changes. +#[derive(Debug)] +pub struct RowAddrMaskFilterExec { + input: Arc, + mask: Arc, + properties: Arc, +} + +impl RowAddrMaskFilterExec { + pub fn new(input: Arc, mask: Arc) -> Self { + // Filtering preserves schema, partitioning and ordering, so the input's + // plan properties carry over unchanged. + let properties = input.properties().clone(); + Self { + input, + mask, + properties, + } + } +} + +impl DisplayAs for RowAddrMaskFilterExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "RowAddrMaskFilter") + } +} + +impl ExecutionPlan for RowAddrMaskFilterExec { + fn name(&self) -> &str { + "RowAddrMaskFilterExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "RowAddrMaskFilterExec must have exactly one child".to_string(), + )); + } + let child = children.pop().ok_or_else(|| { + DataFusionError::Internal("RowAddrMaskFilterExec child unavailable".to_string()) + })?; + Ok(Arc::new(Self::new(child, self.mask.clone()))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let input_stream = self.input.execute(partition, context)?; + let schema = input_stream.schema(); + let mask = self.mask.clone(); + let stream = input_stream.map(move |batch| apply_mask(&mask, batch?)); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +/// Keep rows whose `_rowid` is selected by the mask (the mask is keyed in the +/// same `_rowid` space). Null ids are dropped; they cannot be in any allow set. +fn apply_mask(mask: &RowAddrMask, batch: RecordBatch) -> DataFusionResult { + let row_id_column = batch.column_by_name(ROW_ID).ok_or_else(|| { + DataFusionError::Internal(format!( + "RowAddrMaskFilterExec input missing {ROW_ID} column" + )) + })?; + let row_ids = row_id_column + .as_primitive_opt::() + .ok_or_else(|| { + DataFusionError::Internal(format!( + "{ROW_ID} column must be UInt64 but was {:?}", + row_id_column.data_type() + )) + })?; + let keep = BooleanArray::from_iter( + row_ids + .iter() + .map(|addr| Some(addr.is_some_and(|addr| mask.selected(addr)))), + ); + arrow::compute::filter_record_batch(&batch, &keep) + .map_err(|e| DataFusionError::ArrowError(Box::new(e), None)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow_array::{Int32Array, UInt64Array}; + use lance_select::RowAddrTreeMap; + + fn batch_with_rowids(ids: Vec>) -> RecordBatch { + let n = ids.len() as i32; + let schema = Arc::new(Schema::new(vec![ + Field::new(ROW_ID, DataType::UInt64, true), + Field::new("v", DataType::Int32, false), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from(ids)), + Arc::new(Int32Array::from((0..n).collect::>())), + ], + ) + .unwrap() + } + + fn kept_rowids(batch: &RecordBatch) -> Vec> { + batch + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .iter() + .collect() + } + + #[test] + fn apply_mask_allow_keeps_only_selected() { + let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64, 3, 5])); + let batch = batch_with_rowids(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); + let out = apply_mask(&mask, batch).unwrap(); + assert_eq!(kept_rowids(&out), vec![Some(1), Some(3), Some(5)]); + } + + #[test] + fn apply_mask_block_drops_selected() { + let mask = RowAddrMask::from_block(RowAddrTreeMap::from_iter([2u64, 4])); + let batch = batch_with_rowids(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); + let out = apply_mask(&mask, batch).unwrap(); + assert_eq!(kept_rowids(&out), vec![Some(1), Some(3), Some(5)]); + } + + #[test] + fn apply_mask_drops_null_rowids() { + // A null id cannot be in any allow set, so it is dropped. + let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64, 2, 3])); + let batch = batch_with_rowids(vec![Some(1), None, Some(3)]); + let out = apply_mask(&mask, batch).unwrap(); + assert_eq!(kept_rowids(&out), vec![Some(1), Some(3)]); + } + + #[test] + fn apply_mask_missing_rowid_column_errs() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); + let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64])); + let err = apply_mask(&mask, batch).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(_)), "got {err:?}"); + let msg = err.to_string(); + assert!( + msg.contains(ROW_ID) && msg.contains("missing"), + "unexpected: {msg}" + ); + } + + #[test] + fn apply_mask_wrong_type_rowid_column_errs() { + // _rowid present but not UInt64 -> Internal error naming the actual type. + let schema = Arc::new(Schema::new(vec![Field::new( + ROW_ID, + DataType::Int32, + false, + )])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); + let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64])); + let err = apply_mask(&mask, batch).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(_)), "got {err:?}"); + let msg = err.to_string(); + assert!( + msg.contains("UInt64") && msg.contains("Int32"), + "unexpected: {msg}" + ); + } + + struct FixedLoader(RowAddrMask); + + #[async_trait] + impl FilterLoader for FixedLoader { + async fn load(self: Box) -> Result { + Ok(self.0) + } + } + + #[tokio::test] + async fn mask_and_loader_without_inner_returns_mask() { + let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64, 2, 3])); + let loaded = Box::new(MaskAndLoader::new(Arc::new(mask), None)) + .load() + .await + .unwrap(); + assert!(loaded.selected(2)); + assert!(!loaded.selected(4)); + } + + #[tokio::test] + async fn mask_and_loader_with_inner_intersects() { + // {1,2,3,4} AND inner {2,4,6} = {2,4}. + let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64, 2, 3, 4])); + let inner = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([2u64, 4, 6])); + let loaded = Box::new(MaskAndLoader::new( + Arc::new(mask), + Some(Box::new(FixedLoader(inner))), + )) + .load() + .await + .unwrap(); + assert!(loaded.selected(2)); + assert!(loaded.selected(4)); + assert!(!loaded.selected(1)); + assert!(!loaded.selected(6)); + } + + #[tokio::test] + async fn mask_and_loader_block_and_allow() { + // block{2} AND allow{1,2,3} = allow({1,2,3} - {2}) = allow{1,3}. + let mask = RowAddrMask::from_block(RowAddrTreeMap::from_iter([2u64])); + let inner = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64, 2, 3])); + let loaded = Box::new(MaskAndLoader::new( + Arc::new(mask), + Some(Box::new(FixedLoader(inner))), + )) + .load() + .await + .unwrap(); + assert!(loaded.selected(1)); + assert!(loaded.selected(3)); + assert!(!loaded.selected(2)); + assert!(!loaded.selected(4)); + } + + #[tokio::test] + async fn mask_and_loader_block_and_block() { + // block{1} AND block{2} = block{1,2}: everything except 1 and 2 is selected. + let mask = RowAddrMask::from_block(RowAddrTreeMap::from_iter([1u64])); + let inner = RowAddrMask::from_block(RowAddrTreeMap::from_iter([2u64])); + let loaded = Box::new(MaskAndLoader::new( + Arc::new(mask), + Some(Box::new(FixedLoader(inner))), + )) + .load() + .await + .unwrap(); + assert!(!loaded.selected(1)); + assert!(!loaded.selected(2)); + assert!(loaded.selected(3)); + } +} diff --git a/rust/lance/src/io/exec/utils.rs b/rust/lance/src/io/exec/utils.rs index 1af0bc3f4ef..d5e4324c056 100644 --- a/rust/lance/src/io/exec/utils.rs +++ b/rust/lance/src/io/exec/utils.rs @@ -32,6 +32,7 @@ use lance_core::{ROW_ID, Result}; use lance_index::prefilter::FilterLoader; use lance_select::{RowAddrMask, RowAddrTreeMap, result::IndexExprResult}; +use super::row_addr_mask::MaskAndLoader; use crate::Dataset; use crate::index::prefilter::DatasetPreFilter; @@ -51,6 +52,7 @@ pub(crate) fn build_prefilter( prefilter_source: &PreFilterSource, ds: Arc, index_meta: &[IndexMetadata], + external_mask: Option>, ) -> Result> { let prefilter_loader = match &prefilter_source { PreFilterSource::FilteredRowIds(src_node) => { @@ -63,6 +65,15 @@ pub(crate) fn build_prefilter( } PreFilterSource::None => None, }; + // Combine the external row-address mask (logical AND) with whatever the + // filter produced, so an FTS prefilter restricts BM25 scoring to masked rows + // (mirrors the ANN path). + let prefilter_loader = match external_mask { + Some(mask) => { + Some(Box::new(MaskAndLoader::new(mask, prefilter_loader)) as Box) + } + None => prefilter_loader, + }; Ok(Arc::new(DatasetPreFilter::new( ds, index_meta,