From 1958a59139d5abdeea2ccce7a5cc049206644832 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Mon, 15 Jun 2026 22:25:38 -0700 Subject: [PATCH 1/8] feat(scanner): add external row-address mask prefilter Addresses lance-format/lance#6852. --- rust/lance/src/dataset/scanner.rs | 85 +++++++++++-- rust/lance/src/io/exec.rs | 2 + rust/lance/src/io/exec/knn.rs | 18 +++ rust/lance/src/io/exec/row_addr_mask.rs | 159 ++++++++++++++++++++++++ 4 files changed, 254 insertions(+), 10 deletions(-) create mode 100644 rust/lance/src/io/exec/row_addr_mask.rs diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index da8348d57cc..5668a06d1d4 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -101,7 +101,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 +736,11 @@ pub struct Scanner { /// If true then the filter will be applied before an index scan prefilter: bool, + /// Optional external row-address allow/block mask used as a vector-search + /// prefilter. ANDed into the index-side prefilter and applied to the flat + /// branch for fragments not covered by the index. + external_row_mask: Option, + /// Materialization style controls when columns are fetched materialization_style: MaterializationStyle, @@ -1034,6 +1039,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 +1204,17 @@ impl Scanner { self } + /// Set an external row-address allow/block mask as a vector-search prefilter. + /// + /// The mask is ANDed into any filter-derived prefilter on the index branch + /// and applied to the flat branch for fragments not covered by the vector + /// index. The mask is keyed by row address; with stable row ids disabled the + /// row address equals the row id. Only affects `nearest` queries. + pub fn with_row_addr_prefilter(&mut self, mask: RowAddrMask) -> &mut Self { + self.external_row_mask = Some(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); @@ -2906,8 +2923,22 @@ impl Scanner { fragments: Option>>, scan_range: Option>, ) -> Result> { + // 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 scans apply the mask via the KNN prefilter path, + // so this is scoped to non-nearest scans. + let use_external_mask = self.nearest.is_none() && self.external_row_mask.is_some(); + let effective_filter = if use_external_mask { + match filter_plan.full_expr.clone() { + Some(expr) => ExprFilterPlan::new_refine_only(expr), + None => ExprFilterPlan::default(), + } + } else { + filter_plan.clone() + }; + 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 +2979,17 @@ 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 = if use_external_mask { + Some(self.mask_as_index_input(self.external_row_mask.as_ref().unwrap())?) + } else { + 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(), @@ -3005,6 +3040,19 @@ impl Scanner { } } + // 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> { + let index_result = IndexExprResult::exact(mask.clone()); + 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)?; + let schema = batch.schema(); + let stream = futures::stream::once(async move { Ok(batch) }); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + Ok(Arc::new(OneShotExec::new(stream))) + } + 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); @@ -3903,6 +3951,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, Arc::new(mask))); + } Ok(self.flat_knn(plan, &q)?) } } @@ -4052,6 +4105,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, Arc::new(mask))); + } // first we do flat search on just the new data let topk_appended = self.flat_knn(scan_node, &q)?; @@ -4702,7 +4760,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 +4825,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())?, 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/knn.rs b/rust/lance/src/io/exec/knn.rs index fde0289621d..3cdaab8902c 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,9 @@ pub struct ANNIvfSubIndexExec { /// Prefiltering input prefilter_source: PreFilterSource, + /// Optional external row-address allow/block mask, ANDed into the prefilter. + external_mask: Option, + /// Datafusion Plan Properties properties: Arc, @@ -1390,6 +1397,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 +1417,7 @@ impl ANNIvfSubIndexExec { indices, query, prefilter_source, + external_mask, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -1880,6 +1889,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 +1970,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..fc0d736c94e --- /dev/null +++ b/rust/lance/src/io/exec/row_addr_mask.rs @@ -0,0 +1,159 @@ +// 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 all get ANDed together 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::any::Any; +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 ANDs an external RowAddrMask into an optional inner loader. +/// +/// With an inner loader present the two masks are intersected; otherwise the +/// external mask is used alone. DatasetPreFilter later ANDs the result with the +/// dataset deletion vector. +pub struct MaskAndLoader { + mask: RowAddrMask, + inner: Option>, +} + +impl MaskAndLoader { + pub fn new(mask: RowAddrMask, inner: Option>) -> Self { + Self { mask, inner } + } +} + +#[async_trait] +impl FilterLoader for MaskAndLoader { + async fn load(self: Box) -> Result { + match self.inner { + Some(inner) => Ok(self.mask & inner.load().await?), + None => Ok(self.mask), + } + } +} + +/// Execution node that drops rows whose row address is not selected by `mask`. +/// +/// The row address is read from the `_rowid` column, which equals the row +/// address when stable row ids are disabled (the case this prefilter targets). +/// 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 as_any(&self) -> &dyn Any { + self + } + + 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(), + )); + } + Ok(Arc::new(Self::new( + children.pop().expect("length checked"), + 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` (== row address with stable row ids disabled) is +/// selected by the mask. Null ids are dropped; they cannot be in any allow set. +fn apply_mask(mask: &RowAddrMask, batch: RecordBatch) -> DataFusionResult { + let row_ids = batch + .column_by_name(ROW_ID) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "RowAddrMaskFilterExec input missing {ROW_ID} column" + )) + })? + .as_primitive_opt::() + .ok_or_else(|| DataFusionError::Internal(format!("{ROW_ID} column is not UInt64")))?; + 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)) +} From 9d2f76529615f850c5559f8f8cbfb744cf60aba9 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 04:55:06 -0700 Subject: [PATCH 2/8] fix: applied arc for rowaddrmask; fix with_row_addr_prefilter docs and wording --- rust/lance/src/dataset/scanner.rs | 35 ++++++++++++++++--------- rust/lance/src/io/exec/knn.rs | 9 ++++--- rust/lance/src/io/exec/row_addr_mask.rs | 35 +++++++++++-------------- 3 files changed, 43 insertions(+), 36 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 5668a06d1d4..3559aacc667 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -736,10 +736,13 @@ pub struct Scanner { /// If true then the filter will be applied before an index scan prefilter: bool, - /// Optional external row-address allow/block mask used as a vector-search - /// prefilter. ANDed into the index-side prefilter and applied to the flat - /// branch for fragments not covered by the index. - external_row_mask: Option, + /// 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, @@ -1204,14 +1207,20 @@ impl Scanner { self } - /// Set an external row-address allow/block mask as a vector-search prefilter. + /// Set an external row allow/block mask prefilter. /// - /// The mask is ANDed into any filter-derived prefilter on the index branch - /// and applied to the flat branch for fragments not covered by the vector - /// index. The mask is keyed by row address; with stable row ids disabled the - /// row address equals the row id. Only affects `nearest` queries. + /// On a vector (`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 plain scan the + /// mask is used directly as the row source. + /// + /// 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. pub fn with_row_addr_prefilter(&mut self, mask: RowAddrMask) -> &mut Self { - self.external_row_mask = Some(mask); + self.external_row_mask = Some(Arc::new(mask)); self } @@ -2980,7 +2989,7 @@ impl Scanner { let result_format = self.index_expr_result_format(); let index_input = if use_external_mask { - Some(self.mask_as_index_input(self.external_row_mask.as_ref().unwrap())?) + Some(self.mask_as_index_input(self.external_row_mask.as_deref().unwrap())?) } else { filter_plan.index_query.clone().map(|index_query| { Arc::new(ScalarIndexExec::new( @@ -3954,7 +3963,7 @@ impl Scanner { // 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, Arc::new(mask))); + plan = Arc::new(RowAddrMaskFilterExec::new(plan, mask)); } Ok(self.flat_knn(plan, &q)?) } @@ -4108,7 +4117,7 @@ impl Scanner { // 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, Arc::new(mask))); + 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)?; diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 3cdaab8902c..565a6ba3823 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -1106,7 +1106,7 @@ pub fn new_knn_exec( indices: &[IndexMetadata], query: &Query, prefilter_source: PreFilterSource, - external_mask: Option, + external_mask: Option>, ) -> Result> { let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), @@ -1381,8 +1381,9 @@ pub struct ANNIvfSubIndexExec { /// Prefiltering input prefilter_source: PreFilterSource, - /// Optional external row-address allow/block mask, ANDed into the prefilter. - external_mask: Option, + /// Optional external row-address allow/block mask, combined with the + /// prefilter using logical AND. + external_mask: Option>, /// Datafusion Plan Properties properties: Arc, @@ -1397,7 +1398,7 @@ impl ANNIvfSubIndexExec { indices: Vec, query: Query, prefilter_source: PreFilterSource, - external_mask: Option, + external_mask: Option>, ) -> Result { if input.schema().field_with_name(PART_ID_COLUMN).is_err() { return Err(Error::index(format!( diff --git a/rust/lance/src/io/exec/row_addr_mask.rs b/rust/lance/src/io/exec/row_addr_mask.rs index fc0d736c94e..cf362768e41 100644 --- a/rust/lance/src/io/exec/row_addr_mask.rs +++ b/rust/lance/src/io/exec/row_addr_mask.rs @@ -7,12 +7,11 @@ //! 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 all get ANDed together by DatasetPreFilter. +//! 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::any::Any; use std::sync::Arc; use arrow::datatypes::UInt64Type; @@ -31,18 +30,19 @@ use lance_core::{ROW_ID, Result}; use lance_index::prefilter::FilterLoader; use lance_select::RowAddrMask; -/// FilterLoader that ANDs an external RowAddrMask into an optional inner loader. +/// 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 ANDs the result with the -/// dataset deletion vector. +/// external mask is used alone. DatasetPreFilter later intersects the result +/// with the dataset deletion vector. pub struct MaskAndLoader { - mask: RowAddrMask, + mask: Arc, inner: Option>, } impl MaskAndLoader { - pub fn new(mask: RowAddrMask, inner: Option>) -> Self { + pub fn new(mask: Arc, inner: Option>) -> Self { Self { mask, inner } } } @@ -51,17 +51,18 @@ impl MaskAndLoader { impl FilterLoader for MaskAndLoader { async fn load(self: Box) -> Result { match self.inner { - Some(inner) => Ok(self.mask & inner.load().await?), - None => Ok(self.mask), + 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 row address is not selected by `mask`. +/// Execution node that drops rows whose `_rowid` is not selected by `mask`. /// -/// The row address is read from the `_rowid` column, which equals the row -/// address when stable row ids are disabled (the case this prefilter targets). -/// Schema and ordering are preserved; only the row count changes. +/// 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, @@ -93,10 +94,6 @@ impl ExecutionPlan for RowAddrMaskFilterExec { "RowAddrMaskFilterExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &Arc { &self.properties } @@ -137,8 +134,8 @@ impl ExecutionPlan for RowAddrMaskFilterExec { } } -/// Keep rows whose `_rowid` (== row address with stable row ids disabled) is -/// selected by the mask. Null ids are dropped; they cannot be in any allow set. +/// 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_ids = batch .column_by_name(ROW_ID) From bc541476e3c8fb0721e4e64fdb0cd550a23a5a8e Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 08:05:07 -0700 Subject: [PATCH 3/8] address review: legacy-scan mask guard, dedup index-input helper, mask unit tests --- rust/lance/src/dataset/scanner.rs | 40 +++++---- rust/lance/src/io/exec/row_addr_mask.rs | 105 ++++++++++++++++++++++++ 2 files changed, 131 insertions(+), 14 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 3559aacc667..aaa6e263596 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -3021,6 +3021,18 @@ 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 searches + // (nearest.is_some()) apply the mask via the ANN prefilter and the + // RowAddrMaskFilterExec flat wrap, so they are unaffected. + if self.nearest.is_none() && self.external_row_mask.is_some() { + 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, @@ -3049,10 +3061,12 @@ impl Scanner { } } - // 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> { - let index_result = IndexExprResult::exact(mask.clone()); + // 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)?; @@ -3062,17 +3076,15 @@ 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_addrs = RowAddrTreeMap::from_iter(u64s); - let row_addr_mask = RowAddrMask::from_allowed(row_addrs); - let index_result = IndexExprResult::exact(row_addr_mask); - 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)?; - let schema = batch.schema(); - let stream = futures::stream::once(async move { Ok(batch) }); - let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); - Ok(Arc::new(OneShotExec::new(stream))) + 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> { diff --git a/rust/lance/src/io/exec/row_addr_mask.rs b/rust/lance/src/io/exec/row_addr_mask.rs index cf362768e41..900dd09181b 100644 --- a/rust/lance/src/io/exec/row_addr_mask.rs +++ b/rust/lance/src/io/exec/row_addr_mask.rs @@ -154,3 +154,108 @@ fn apply_mask(mask: &RowAddrMask, batch: RecordBatch) -> DataFusionResult>) -> 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])); + assert!(apply_mask(&mask, batch).is_err()); + } + + 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)); + } +} From 4198c4e5af6deb163482fc5206538c7c79097502 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 08:43:37 -0700 Subject: [PATCH 4/8] test(scanner): row-addr mask regression tests; return errors instead of panicking --- rust/lance/src/dataset/scanner.rs | 104 ++++++++++++++++++++++-- rust/lance/src/io/exec/row_addr_mask.rs | 79 ++++++++++++++---- 2 files changed, 164 insertions(+), 19 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index aaa6e263596..d548aabc5dc 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2988,16 +2988,15 @@ impl Scanner { } let result_format = self.index_expr_result_format(); - let index_input = if use_external_mask { - Some(self.mask_as_index_input(self.external_row_mask.as_deref().unwrap())?) - } else { - filter_plan.index_query.clone().map(|index_query| { + 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( @@ -5621,6 +5620,101 @@ mod test { } } + fn batch_row_ids(batch: &RecordBatch) -> Vec { + batch + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::() + .values() + .to_vec() + } + + #[tokio::test] + async fn row_addr_mask_plain_scan_allow_with_refine() { + 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(); + + // 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); + + // 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}" + ); + } + + #[tokio::test] + async fn row_addr_mask_ann_search_only_allowed() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .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 test_strict_batch_size() { let dataset = lance_datagen::gen_batch() diff --git a/rust/lance/src/io/exec/row_addr_mask.rs b/rust/lance/src/io/exec/row_addr_mask.rs index 900dd09181b..f1b43bfc09d 100644 --- a/rust/lance/src/io/exec/row_addr_mask.rs +++ b/rust/lance/src/io/exec/row_addr_mask.rs @@ -115,10 +115,10 @@ impl ExecutionPlan for RowAddrMaskFilterExec { "RowAddrMaskFilterExec must have exactly one child".to_string(), )); } - Ok(Arc::new(Self::new( - children.pop().expect("length checked"), - self.mask.clone(), - ))) + 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( @@ -137,15 +137,15 @@ impl ExecutionPlan for RowAddrMaskFilterExec { /// 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_ids = batch - .column_by_name(ROW_ID) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "RowAddrMaskFilterExec input missing {ROW_ID} column" - )) - })? - .as_primitive_opt::() - .ok_or_else(|| DataFusionError::Internal(format!("{ROW_ID} column is not UInt64")))?; + 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() @@ -218,7 +218,23 @@ mod tests { let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); let mask = RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([1u64])); - assert!(apply_mask(&mask, batch).is_err()); + 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); @@ -258,4 +274,39 @@ mod tests { 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)); + } } From 4ec81869da2f17b225798b88a002d262a7076309 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 09:04:58 -0700 Subject: [PATCH 5/8] test(scanner): extend the row-address scanner tests around row_addr_mask_plain_scan_allow_with_refine and row_addr_mask_ann_search_only_allowed --- rust/lance/src/dataset/scanner.rs | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index d548aabc5dc..14021aab7d0 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -5629,9 +5629,12 @@ mod test { .to_vec() } + #[rstest] #[tokio::test] - async fn row_addr_mask_plain_scan_allow_with_refine() { - let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + async fn row_addr_mask_plain_scan_allow_block_refine( + #[values(false, true)] stable_row_ids: bool, + ) { + let test_ds = TestVectorDataset::new(LanceFileVersion::Stable, stable_row_ids) .await .unwrap(); let ds = &test_ds.dataset; @@ -5639,6 +5642,7 @@ mod test { 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(); @@ -5653,6 +5657,21 @@ mod test { .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( @@ -5685,9 +5704,10 @@ mod test { ); } + #[rstest] #[tokio::test] - async fn row_addr_mask_ann_search_only_allowed() { - let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + async fn row_addr_mask_ann_search_only_allowed(#[values(false, true)] stable_row_ids: bool) { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, stable_row_ids) .await .unwrap(); test_ds.make_vector_index().await.unwrap(); From f85ecdad2f7a2bcc1d142d5c57891c2889beb5f1 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 17:58:54 -0700 Subject: [PATCH 6/8] address review: masked-scan projection/limit, with_fetch guard, docs --- rust/lance/src/dataset/scanner.rs | 179 ++++++++++++++++++++---- rust/lance/src/io/exec/filtered_read.rs | 7 +- rust/lance/src/io/exec/row_addr_mask.rs | 30 ++-- 3 files changed, 181 insertions(+), 35 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 14021aab7d0..2afc95ebb08 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}; @@ -1207,18 +1210,38 @@ impl Scanner { self } - /// Set an external row allow/block mask prefilter. + /// Set an external [`RowAddrMask`] allow/block prefilter. /// - /// On a vector (`nearest`) search the mask is combined with any + /// 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 plain scan the - /// mask is used directly as the row source. + /// 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 @@ -2921,6 +2944,30 @@ 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 scans apply the mask via the KNN prefilter path, so this is scoped to + // non-nearest scans. + fn use_external_mask(&self) -> bool { + self.nearest.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 @@ -2932,19 +2979,8 @@ impl Scanner { fragments: Option>>, scan_range: Option>, ) -> Result> { - // 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 scans apply the mask via the KNN prefilter path, - // so this is scoped to non-nearest scans. - let use_external_mask = self.nearest.is_none() && self.external_row_mask.is_some(); - let effective_filter = if use_external_mask { - match filter_plan.full_expr.clone() { - Some(expr) => ExprFilterPlan::new_refine_only(expr), - None => ExprFilterPlan::default(), - } - } else { - filter_plan.clone() - }; + 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(effective_filter) @@ -3139,11 +3175,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 @@ -3157,7 +3197,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 { @@ -5630,10 +5674,10 @@ mod test { } #[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( - #[values(false, true)] stable_row_ids: bool, - ) { + 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(); @@ -5683,7 +5727,10 @@ mod test { 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::(); + let is = refined + .column_by_name("i") + .unwrap() + .as_primitive::(); assert!(is.values().iter().all(|v| *v >= 200)); } @@ -5705,8 +5752,10 @@ mod test { } #[rstest] + #[case::without_stable_row_ids(false)] + #[case::with_stable_row_ids(true)] #[tokio::test] - async fn row_addr_mask_ann_search_only_allowed(#[values(false, true)] stable_row_ids: bool) { + 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(); @@ -5731,10 +5780,88 @@ mod test { 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"); + 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 test_strict_batch_size() { let dataset = lance_datagen::gen_batch() diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 414ce601763..6ca418ae74a 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)); diff --git a/rust/lance/src/io/exec/row_addr_mask.rs b/rust/lance/src/io/exec/row_addr_mask.rs index f1b43bfc09d..eb7059098bc 100644 --- a/rust/lance/src/io/exec/row_addr_mask.rs +++ b/rust/lance/src/io/exec/row_addr_mask.rs @@ -138,14 +138,18 @@ impl ExecutionPlan for RowAddrMaskFilterExec { /// 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() + "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() @@ -221,20 +225,30 @@ mod tests { 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}"); + 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 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}"); + assert!( + msg.contains("UInt64") && msg.contains("Int32"), + "unexpected: {msg}" + ); } struct FixedLoader(RowAddrMask); From cb2a120411ee544f449feebc047dcc2504eee927 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 20:23:15 -0700 Subject: [PATCH 7/8] feat(scanner): apply external row mask to full-text search --- rust/lance/src/dataset/scanner.rs | 102 +++++++++++++++++++++++++----- rust/lance/src/io/exec/fts.rs | 52 +++++++++++++-- rust/lance/src/io/exec/utils.rs | 11 ++++ 3 files changed, 145 insertions(+), 20 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 2afc95ebb08..7e7b41da71d 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2946,10 +2946,12 @@ 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 scans apply the mask via the KNN prefilter path, so this is scoped to - // non-nearest scans. + // 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.external_row_mask.is_some() + 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 @@ -3631,12 +3633,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( @@ -3683,12 +3688,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) @@ -3778,12 +3786,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) } @@ -5862,6 +5877,61 @@ mod test { 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/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/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, From b72c16fb2b17f67a4d30cd91806c8a7c9b5ac834 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Sun, 19 Jul 2026 21:23:02 -0700 Subject: [PATCH 8/8] address review: FTS docs, legacy mask guard, with_fetch coverage - with_row_addr_prefilter doc now describes full-text search behavior (mask combined into the FTS prefilter; unindexed flat branch masked via RowAddrMaskFilterExec) and links full_text_search. - Legacy-storage mask guard keys on use_external_mask() instead of nearest.is_none(), so FTS and vector masked scans (handled via their own prefilter wrapping) are no longer wrongly rejected on legacy storage; only the unsupported plain-scan mask path is rejected. - test_with_fetch_limit_pushdown gains a case asserting with_fetch returns None when an index_input is present with no filters (the external-mask shape). --- rust/lance/src/dataset/scanner.rs | 19 ++++++++++------ rust/lance/src/io/exec/filtered_read.rs | 29 +++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 7e7b41da71d..8acc5d02601 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -1216,9 +1216,13 @@ impl Scanner { /// 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 plain scan the - /// mask is used directly as the row source, with any [`filter`](Self::filter) - /// applied as a refine on top. + /// 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 @@ -3060,10 +3064,11 @@ impl Scanner { 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 searches - // (nearest.is_some()) apply the mask via the ANN prefilter and the - // RowAddrMaskFilterExec flat wrap, so they are unaffected. - if self.nearest.is_none() && self.external_row_mask.is_some() { + // 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" diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 6ca418ae74a..ff10d96f336 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -3327,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]