From 8f35e0d23bc6b4dd4ebcc1563f3e01d180483d61 Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Wed, 12 Aug 2026 03:01:52 -0700 Subject: [PATCH 1/3] feat(scanner): add external row address mask prefilter Lets callers pass a serialized RowAddrMask as an allow/block prefilter into vector, full-text, and plain scans, reusing the scanner's retrieval plan. The mask feeds the KNN prefilter source on the ANN branch, the FTS prefilter so BM25 top-k is computed over masked rows, and FilteredReadExec as the row source for plain scans; a new RowAddrMaskFilterExec honors the mask on the flat/unindexed-fragment branch. Addresses #6852. --- rust/lance-select/src/mask.rs | 158 ++++++- rust/lance/src/dataset/scanner.rs | 548 ++++++++++++++++++++++-- rust/lance/src/io/exec.rs | 2 + rust/lance/src/io/exec/filtered_read.rs | 36 +- rust/lance/src/io/exec/fts.rs | 64 ++- rust/lance/src/io/exec/knn.rs | 31 +- rust/lance/src/io/exec/row_addr_mask.rs | 326 ++++++++++++++ rust/lance/src/io/exec/utils.rs | 12 + 8 files changed, 1147 insertions(+), 30 deletions(-) create mode 100644 rust/lance/src/io/exec/row_addr_mask.rs diff --git a/rust/lance-select/src/mask.rs b/rust/lance-select/src/mask.rs index 492ef6ef608..ccad0bc6243 100644 --- a/rust/lance-select/src/mask.rs +++ b/rust/lance-select/src/mask.rs @@ -107,6 +107,51 @@ impl RowAddrMask { } } + /// Build a mask from serialized [`RowAddrTreeMap`] payloads. + /// + /// `allow` selects rows, `block` excludes them; each is the output of + /// [`RowAddrTreeMap::serialize_into`]. Returns `None` when neither is given, + /// which callers read as "no mask" rather than "select nothing". + /// + /// Bytes rather than treemaps on purpose: a caller living in a different + /// dynamically-linked extension module has its own copy of these Rust types + /// and cannot hand one over, but both sides agree on this encoding. + pub fn from_serialized_parts( + allow: Option<&[u8]>, + block: Option<&[u8]>, + ) -> Result> { + // Name the offending side: the underlying failure is a bare "failed to + // fill whole buffer", which tells a caller holding two blobs nothing. + fn decode(bytes: &[u8], which: &str) -> Result { + RowAddrTreeMap::deserialize_from(bytes).map_err(|e| { + Error::invalid_input(format!( + "row address {which} is not a serialized RowAddrTreeMap: {e}" + )) + }) + } + let allow = allow.map(|b| decode(b, "allowlist")).transpose()?; + let block = block.map(|b| decode(b, "blocklist")).transpose()?; + Ok(match (allow, block) { + (Some(allow), Some(block)) => Some(Self::from_allowed(allow).also_block(block)), + (Some(allow), None) => Some(Self::from_allowed(allow)), + (None, Some(block)) => Some(Self::from_block(block)), + (None, None) => None, + }) + } + + /// Intersect two masks: a row survives only if both select it. + /// + /// Lets a planner apply a caller-supplied mask at one boundary rather than + /// at every branch that produces rows, which is how branches get missed. + pub fn intersect(self, other: Self) -> Self { + match (self, other) { + (Self::AllowList(a), Self::AllowList(b)) => Self::AllowList(a & b), + (Self::AllowList(a), Self::BlockList(b)) => Self::AllowList(a).also_block(b), + (Self::BlockList(a), Self::AllowList(b)) => Self::AllowList(b).also_block(a), + (Self::BlockList(a), Self::BlockList(b)) => Self::BlockList(a | b), + } + } + /// Also allow the given addrs pub fn also_allow(self, allow_list: RowAddrTreeMap) -> Self { match self { @@ -623,8 +668,21 @@ impl RowAddrTreeMap { if bitmap_size == 0 { inner.insert(fragment, RowAddrSelection::Full); } else { - let mut buffer = vec![0; bitmap_size as usize]; - reader.read_exact(&mut buffer)?; + // Grow with the bytes that actually arrive instead of trusting the + // declared size. This is reachable from a public byte boundary, so + // a 12-byte payload could otherwise declare 4 GiB and abort the + // process on the allocation before any read fails. + let mut buffer = Vec::new(); + let read = reader + .by_ref() + .take(u64::from(bitmap_size)) + .read_to_end(&mut buffer)?; + if read != bitmap_size as usize { + return Err(Error::invalid_input(format!( + "row addr treemap declares a {bitmap_size} byte bitmap for \ + fragment {fragment} but only {read} bytes remain" + ))); + } let set = RoaringBitmap::deserialize_from(&buffer[..])?; inner.insert(fragment, RowAddrSelection::Partial(set)); } @@ -1294,6 +1352,102 @@ mod tests { assert!(mask.iter_addrs().is_none()); } + #[test] + fn test_row_addr_mask_intersect() { + let a = rows(&[1, 2, 3]); + let b = rows(&[3, 4]); + + // allow & allow -> only rows in both + assert_mask_selects( + &RowAddrMask::from_allowed(a.clone()).intersect(RowAddrMask::from_allowed(b.clone())), + &[3], + &[1, 2, 4, 100], + ); + // allow & block -> allowed minus blocked + assert_mask_selects( + &RowAddrMask::from_allowed(a.clone()).intersect(RowAddrMask::from_block(b.clone())), + &[1, 2], + &[3, 4, 100], + ); + // block & allow -> same, order independent + assert_mask_selects( + &RowAddrMask::from_block(b.clone()).intersect(RowAddrMask::from_allowed(a.clone())), + &[1, 2], + &[3, 4, 100], + ); + // block & block -> both exclusions apply + assert_mask_selects( + &RowAddrMask::from_block(a.clone()).intersect(RowAddrMask::from_block(b)), + &[100], + &[1, 2, 3, 4], + ); + // all_rows is the identity, and intersecting with itself changes nothing + let allow_a = RowAddrMask::from_allowed(a.clone()); + assert_eq!(allow_a.clone().intersect(RowAddrMask::all_rows()), allow_a); + assert_eq!(allow_a.clone().intersect(allow_a.clone()), allow_a); + // allow_nothing absorbs + assert_mask_selects( + &RowAddrMask::allow_nothing().intersect(RowAddrMask::from_allowed(a)), + &[], + &[1, 2, 3, 100], + ); + } + + #[test] + fn test_row_addr_mask_from_serialized_parts() { + fn ser(tm: &RowAddrTreeMap) -> Vec { + let mut buf = Vec::new(); + tm.serialize_into(&mut buf).unwrap(); + buf + } + let allow = ser(&rows(&[1, 2, 3])); + let block = ser(&rows(&[3, 4])); + + // Neither part means "no mask", which is not the same as "select nothing". + assert!( + RowAddrMask::from_serialized_parts(None, None) + .unwrap() + .is_none() + ); + + let m = RowAddrMask::from_serialized_parts(Some(&allow), None) + .unwrap() + .unwrap(); + assert_mask_selects(&m, &[1, 2, 3], &[4, 100]); + + let m = RowAddrMask::from_serialized_parts(None, Some(&block)) + .unwrap() + .unwrap(); + assert_mask_selects(&m, &[1, 2, 100], &[3, 4]); + + // Block wins on the overlap. + let m = RowAddrMask::from_serialized_parts(Some(&allow), Some(&block)) + .unwrap() + .unwrap(); + assert_mask_selects(&m, &[1, 2], &[3, 4, 100]); + + // Round trips through the same encoding the caller used. + let again = RowAddrMask::from_serialized_parts(Some(&ser(m.allow_list().unwrap())), None) + .unwrap() + .unwrap(); + assert_mask_selects(&again, &[1, 2], &[3, 4]); + + assert!(RowAddrMask::from_serialized_parts(Some(b"not a treemap"), None).is_err()); + + // A declared bitmap size must not be allocated before the bytes are + // known to exist: this 12-byte payload claims ~4 GiB. + let bomb = [ + 1u8, 0, 0, 0, // one entry + 0, 0, 0, 0, // fragment zero + 0xff, 0xff, 0xff, 0xff, // declared bitmap size + ]; + let err = RowAddrMask::from_serialized_parts(Some(&bomb), None).unwrap_err(); + assert!( + err.to_string().contains("only 0 bytes remain"), + "expected a length complaint, got: {err}" + ); + } + #[test] fn test_row_addr_mask_not() { let allow_list = RowAddrMask::from_allowed(rows(&[1, 2, 3])); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 7f988f8a8a3..edaa6d6a25f 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -82,7 +82,10 @@ use lance_index::scalar::registry::VALUE_COLUMN_NAME; use lance_index::vector::{ApproxMode, DEFAULT_QUERY_PARALLELISM, DIST_COL, Query}; 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 prost::Message; use roaring::RoaringBitmap; @@ -116,7 +119,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, }, @@ -920,6 +923,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, @@ -1219,6 +1230,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, @@ -1383,6 +1395,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); @@ -3155,6 +3208,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 @@ -3168,8 +3247,11 @@ impl Scanner { ) -> Result> { // Kept for the overlay stale-Take path below, which re-evaluates blocked stale rows. let user_projection = projection.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(filter_plan.clone()) + .with_filter_plan(effective_filter) .with_projection(projection); if let Some(fragments) = fragments { @@ -3228,13 +3310,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_take_input(mask.clone())?), + _ => filter_plan.index_query.clone().map(|index_query| { + Arc::new(ScalarIndexExec::new( + self.dataset.clone(), + index_query, + result_format, + )) as Arc + }), + }; let plan: Arc = Arc::new(FilteredReadExec::try_new( self.dataset.clone(), @@ -3281,6 +3366,25 @@ impl Scanner { scan_range: Option>, is_prefilter: bool, ) -> BoxFuture<'a, Result> { + // The plain-scan mask path lives in new_filtered_read; legacy_filtered_read + // has no equivalent, so a masked plain scan there 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. + let is_legacy = self + .dataset + .manifest() + .data_storage_format + .lance_file_format() + == lance_file::version::ConcreteFileVersion::V1; + if is_legacy && self.use_external_mask() { + return std::future::ready(Err(Error::not_supported( + "with_row_addr_prefilter is not supported for plain scans on \ + legacy-storage datasets", + ))) + .boxed(); + } versions::filtered_read( self.dataset .manifest() @@ -3298,8 +3402,24 @@ impl Scanner { } fn row_ids_as_take_input(&self, row_ids: RowAddrTreeMap) -> Result> { - let row_id_mask = RowAddrMask::from_allowed(row_ids); - let index_result = IndexExprResult::exact(row_id_mask); + self.mask_as_take_input(RowAddrMask::from_allowed(row_ids)) + } + + // Wrap a row-address mask as a one-shot index input for FilteredReadExec, so a + // plain scan reads only the rows the mask selects. + // + // Every take-shaped row source funnels through here: plain takes, the + // _rowid/_rowaddr predicate shortcut, and the overlay stale-row replay under + // both scan and ANN. Intersecting the caller's mask once at this boundary is + // what keeps the invariant on all of them; applying it per branch is how + // branches get missed. Idempotent, so the branch that passes the external + // mask itself is unaffected. + fn mask_as_take_input(&self, mask: RowAddrMask) -> Result> { + let mask = match self.external_row_mask.as_deref() { + Some(external) => mask.intersect(external.clone()), + None => mask, + }; + let index_result = IndexExprResult::exact(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)?; @@ -3370,11 +3490,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 @@ -3388,7 +3512,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 { @@ -4056,13 +4184,16 @@ impl Scanner { let (_, segments) = segment_groups.into_iter().next().ok_or_else(|| { Error::internal("compound scorer requires one column".to_string()) })?; - return Ok(Some(Arc::new(CompoundQueryExec::new_with_segments( - self.dataset.clone(), - query.clone(), - params.clone(), - prefilter_source.clone(), - segments, - )))); + return Ok(Some(Arc::new( + CompoundQueryExec::new_with_segments( + self.dataset.clone(), + query.clone(), + params.clone(), + prefilter_source.clone(), + segments, + ) + .with_external_mask(self.external_row_mask.clone()), + ))); } let exec = CrossColumnCompoundQueryExec::new_with_segments( @@ -4071,7 +4202,8 @@ impl Scanner { params.clone(), prefilter_source.clone(), segment_groups, - )?; + )? + .with_external_mask(self.external_row_mask.clone()); Ok(Some(Arc::new(exec))) } @@ -4400,6 +4532,7 @@ impl Scanner { if let Some(shared_scorer) = &shared_scorer { phrase_exec = phrase_exec.with_shared_scorer(shared_scorer.clone()); } + phrase_exec = phrase_exec.with_external_mask(self.external_row_mask.clone()); let phrase_plan = Some(Arc::new(phrase_exec) as Arc); let flat_phrase_plan = if has_flat_path { Some( @@ -4558,6 +4691,7 @@ impl Scanner { if let Some(shared_scorer) = &shared_scorer { match_exec = match_exec.with_shared_scorer(shared_scorer.clone()); } + match_exec = match_exec.with_external_mask(self.external_row_mask.clone()); let match_plan = Some(Arc::new(match_exec) as Arc); let flat_match_plan = if has_flat_path { Some( @@ -4734,7 +4868,15 @@ impl Scanner { if let Some(shared_scorer) = shared_scorer { flat_match_plan = flat_match_plan.with_shared_scorer(shared_scorer); } - Ok(Arc::new(flat_match_plan)) + let flat_match_plan: Arc = Arc::new(flat_match_plan); + // Unindexed fragments and stale rows 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) } // ANN/KNN search execution node with optional prefilter @@ -4997,6 +5139,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)?) } } @@ -5155,6 +5302,12 @@ impl Scanner { if let Some(expr) = filter_plan.full_expr.as_ref() { 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. + let scan_node = match self.external_row_mask.clone() { + Some(mask) => Arc::new(RowAddrMaskFilterExec::new(scan_node, mask)) as _, + None => scan_node, + }; let topk_fallback = self.flat_knn(scan_node, &q)?; let topk_fallback: Arc = Arc::new(project(topk_fallback, knn_node.schema().as_ref())?); @@ -6110,6 +6263,7 @@ impl Scanner { q, prefilter_source, overlay_block, + self.external_row_mask.clone(), )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, inner_fanout_search.schema().as_ref())?, @@ -6171,6 +6325,7 @@ impl Scanner { &query, prefilter_source.clone(), overlay_block.clone(), + self.external_row_mask.clone(), )?; let sort_expr = PhysicalSortExpr { expr: expressions::col(DIST_COL, ann_node.schema().as_ref())?, @@ -7130,6 +7285,351 @@ 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"); + } + + /// A `_rowid` predicate is recognized as a TakeOperation and short-circuits + /// straight to `take_source`, which used to skip the mask entirely. + #[tokio::test] + async fn row_addr_mask_take_shortcut_respects_mask() { + 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 target = all_ids[0]; + + // Sanity: unmasked, the shortcut returns the row. + let mut scan = ds.scan(); + scan.with_row_id(); + scan.filter(&format!("_rowid = {target}")).unwrap(); + assert_eq!(scan.try_into_batch().await.unwrap().num_rows(), 1); + + // Masked to nothing, it must return nothing. + let mut scan = ds.scan(); + scan.with_row_id(); + scan.filter(&format!("_rowid = {target}")).unwrap(); + scan.with_row_addr_prefilter(RowAddrMask::allow_nothing()); + assert_eq!( + scan.try_into_batch().await.unwrap().num_rows(), + 0, + "the take shortcut must not return rows the mask excludes" + ); + + // And an allow-list restricts it rather than being ignored. + let mut scan = ds.scan(); + scan.with_row_id(); + scan.filter(&format!("_rowid = {target}")).unwrap(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([ + target, + ]))); + assert_eq!(scan.try_into_batch().await.unwrap().num_rows(), 1); + } + + /// A same-column compound query (Boost here) is optimized into + /// CompoundFtsScorer, a scorer that built its prefilter without the mask. + #[tokio::test] + async fn row_addr_mask_compound_fts_respects_mask() { + let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false) + .await + .unwrap(); + test_ds.make_fts_index().await.unwrap(); + let ds = &test_ds.dataset; + + let compound = || { + let positive = MatchQuery::new("4".to_owned()).with_column(Some("s".to_owned())); + let negative = MatchQuery::new("9".to_owned()).with_column(Some("s".to_owned())); + FullTextSearchQuery::new_query( + BoostQuery::new(positive.into(), negative.into(), Some(1.0)).into(), + ) + }; + + let mut scan = ds.scan(); + scan.full_text_search(compound()).unwrap(); + scan.with_row_id(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("CompoundFtsScorer"), + "expected the compound scorer path, got:\n{plan}" + ); + let base = batch_row_ids(&scan.try_into_batch().await.unwrap()); + assert!(!base.is_empty(), "compound query matched nothing"); + + let mut scan = ds.scan(); + scan.full_text_search(compound()).unwrap(); + scan.with_row_id(); + scan.with_row_addr_prefilter(RowAddrMask::allow_nothing()); + assert_eq!( + scan.try_into_batch().await.unwrap().num_rows(), + 0, + "the compound scorer must not return rows the mask excludes" + ); + + // Allow exactly one baseline hit; only that one may come back. + let keep = base[0]; + let mut scan = ds.scan(); + scan.full_text_search(compound()).unwrap(); + scan.with_row_id(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([keep]))); + assert_eq!( + batch_row_ids(&scan.try_into_batch().await.unwrap()), + vec![keep] + ); + } + + #[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_batch_size_bytes_across_data_files() { let num_rows = 300; diff --git a/rust/lance/src/io/exec.rs b/rust/lance/src/io/exec.rs index 6923f09c34f..d37b58a238e 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(crate) use scan::LanceStream; pub use scan::{LanceScanConfig, LanceScanExec}; diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 3c5936e9dfd..1e60556761e 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -3278,7 +3278,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)); @@ -4697,6 +4702,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 c99e964a3e8..4ccc0bc3207 100644 --- a/rust/lance/src/io/exec/fts.rs +++ b/rust/lance/src/io/exec/fts.rs @@ -33,6 +33,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; @@ -67,7 +68,6 @@ use lance_index::scalar::inverted::{ flat_bm25_search_stream_with_options_and_scorer, fts_schema, }; use lance_index::{prefilter::PreFilter, scalar::inverted::query::BooleanQuery}; -use lance_select::RowAddrMask; use lance_tokenizer::{SimpleTokenizer, TextAnalyzer}; use tracing::instrument; use uuid::Uuid; @@ -541,6 +541,10 @@ pub struct CompoundQueryExec { /// searched segments — see [`MatchQueryExec::with_base_scorer`]. base_scorer: Option>, segment_selection: FtsSegmentSelection, + /// Caller-supplied row-address mask, intersected into the prefilter so the + /// compound scorer ranks only surviving rows (see + /// [`MatchQueryExec::with_external_mask`]). + external_mask: Option>, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -593,6 +597,7 @@ impl CompoundQueryExec { prefilter_source, base_scorer: None, segment_selection, + external_mask: None, properties: Arc::new(PlanProperties::new( EquivalenceProperties::new(FTS_SCHEMA.clone()), Partitioning::RoundRobinBatch(1), @@ -603,6 +608,12 @@ impl CompoundQueryExec { } } + /// See [`MatchQueryExec::with_external_mask`]. + pub fn with_external_mask(mut self, mask: Option>) -> Self { + self.external_mask = mask; + self + } + /// Override locally computed BM25 statistics with a corpus-wide scorer. /// /// The scorer must cover every token in every query leaf, including fuzzy @@ -712,6 +723,7 @@ impl ExecutionPlan for CompoundQueryExec { prefilter_source, base_scorer: self.base_scorer.clone(), segment_selection: self.segment_selection.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), })) @@ -730,6 +742,7 @@ impl ExecutionPlan for CompoundQueryExec { let prefilter_source = self.prefilter_source.clone(); let base_scorer = self.base_scorer.clone(); let segment_selection = self.segment_selection.clone(); + let external_mask = self.external_mask.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); let stream = stream::once(async move { @@ -767,6 +780,7 @@ impl ExecutionPlan for CompoundQueryExec { dataset, &segments, None, + external_mask, )?; let deleted_fragments = indices @@ -854,6 +868,9 @@ pub struct CrossColumnCompoundQueryExec { params: FtsSearchParams, prefilter_source: PreFilterSource, columns: Arc<[CompoundColumnSelection]>, + /// Combined into the prefilter so only masked rows are scored (see + /// [`MatchQueryExec::with_external_mask`]). + external_mask: Option>, properties: Arc, metrics: ExecutionPlanMetricsSet, } @@ -930,6 +947,7 @@ impl CrossColumnCompoundQueryExec { params, prefilter_source, columns: Arc::from(columns), + external_mask: None, properties: Arc::new(PlanProperties::new( EquivalenceProperties::new(FTS_SCHEMA.clone()), Partitioning::RoundRobinBatch(1), @@ -940,6 +958,12 @@ impl CrossColumnCompoundQueryExec { }) } + /// See [`MatchQueryExec::with_external_mask`]. + pub fn with_external_mask(mut self, mask: Option>) -> Self { + self.external_mask = mask; + self + } + pub fn dataset(&self) -> &Arc { &self.dataset } @@ -1032,6 +1056,7 @@ impl ExecutionPlan for CrossColumnCompoundQueryExec { params: self.params.clone(), prefilter_source, columns: self.columns.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), })) @@ -1053,6 +1078,7 @@ impl ExecutionPlan for CrossColumnCompoundQueryExec { let params = self.params.clone(); let prefilter_source = self.prefilter_source.clone(); let columns = self.columns.clone(); + let external_mask = self.external_mask.clone(); let metrics = Arc::new(FtsIndexMetrics::new(&self.metrics, partition)); let stream = stream::once(async move { @@ -1083,6 +1109,7 @@ impl ExecutionPlan for CrossColumnCompoundQueryExec { dataset.clone(), &selected_segments, None, + external_mask, )?; let opened_columns = try_join_all(columns.iter().cloned().map(|selection| { let dataset = dataset.clone(); @@ -1736,6 +1763,9 @@ pub struct MatchQueryExec { overlay_block: Option, document_granularity: DocumentGranularity, schema: SchemaRef, + /// 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, @@ -1821,6 +1851,7 @@ impl MatchQueryExec { overlay_block: None, document_granularity, schema, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -1883,6 +1914,7 @@ impl MatchQueryExec { overlay_block: None, document_granularity, schema, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -1923,6 +1955,7 @@ impl MatchQueryExec { shared_scorer: None, segment_selection: FtsSegmentSelection::exact_uuids(segment_uuids), overlay_block: None, + external_mask: None, document_granularity, schema, properties, @@ -1958,6 +1991,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 built by + /// `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 } @@ -2037,6 +2079,7 @@ impl ExecutionPlan for MatchQueryExec { overlay_block: self.overlay_block.clone(), document_granularity: self.document_granularity, schema: self.schema.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -2069,6 +2112,7 @@ impl ExecutionPlan for MatchQueryExec { overlay_block: self.overlay_block.clone(), document_granularity: self.document_granularity, schema: self.schema.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -2093,6 +2137,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 shared_scorer = self.shared_scorer.clone(); let segment_selection = self.segment_selection.clone(); @@ -2124,6 +2169,7 @@ impl ExecutionPlan for MatchQueryExec { ds, &segments, overlay_block, + external_mask, )?; let deleted_fragments = indices @@ -3053,6 +3099,9 @@ pub struct PhraseQueryExec { overlay_block: Option, document_granularity: DocumentGranularity, schema: SchemaRef, + /// 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, } @@ -3129,6 +3178,7 @@ impl PhraseQueryExec { overlay_block: None, document_granularity, schema, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), } @@ -3182,6 +3232,7 @@ impl PhraseQueryExec { shared_scorer: None, segment_selection: FtsSegmentSelection::ExactResolved(Arc::from(segments)), overlay_block: None, + external_mask: None, document_granularity, schema, properties, @@ -3227,6 +3278,7 @@ impl PhraseQueryExec { overlay_block: None, document_granularity, schema, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -3249,6 +3301,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 } @@ -3321,6 +3379,7 @@ impl ExecutionPlan for PhraseQueryExec { overlay_block: self.overlay_block.clone(), document_granularity: self.document_granularity, schema: self.schema.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), }, @@ -3351,6 +3410,7 @@ impl ExecutionPlan for PhraseQueryExec { overlay_block: self.overlay_block.clone(), document_granularity: self.document_granularity, schema: self.schema.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -3375,6 +3435,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 shared_scorer = self.shared_scorer.clone(); let segment_selection = self.segment_selection.clone(); @@ -3406,6 +3467,7 @@ impl ExecutionPlan for PhraseQueryExec { ds, &segments, overlay_block, + external_mask, )?; let deleted_fragments = indices diff --git a/rust/lance/src/io/exec/knn.rs b/rust/lance/src/io/exec/knn.rs index 52a2af01b19..04d24170211 100644 --- a/rust/lance/src/io/exec/knn.rs +++ b/rust/lance/src/io/exec/knn.rs @@ -56,13 +56,12 @@ 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 roaring::RoaringBitmap; use tokio::sync::Notify; use uuid::Uuid; -use lance_select::RowAddrMask; - use crate::dataset::Dataset; use crate::index::DatasetIndexInternalExt; use crate::index::prefilter::{DatasetPreFilter, FilterLoader}; @@ -70,6 +69,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,12 +1104,15 @@ pub static KNN_PARTITION_SCHEMA: LazyLock = LazyLock::new(|| { /// Create a new ANN execution node. `overlay_block`, when `Some`, excludes rows whose index /// entries may be stale due to a newer data overlay (see [`ANNIvfSubIndexExec::with_overlay_block`]). +/// `external_mask`, when `Some`, additionally restricts the scan to a caller-supplied +/// allow/block set (see [`ANNIvfSubIndexExec::with_external_mask`]). pub fn new_knn_exec( dataset: Arc, indices: &[IndexMetadata], query: &Query, prefilter_source: PreFilterSource, overlay_block: Option, + external_mask: Option>, ) -> Result> { let ivf_node = ANNIvfPartitionExec::try_new( dataset.clone(), @@ -1127,6 +1130,9 @@ pub fn new_knn_exec( if let Some(overlay_block) = overlay_block { sub_index = sub_index.with_overlay_block(overlay_block); } + if external_mask.is_some() { + sub_index = sub_index.with_external_mask(external_mask); + } Ok(Arc::new(sub_index)) } @@ -1390,6 +1396,10 @@ pub struct ANNIvfSubIndexExec { /// index results at execution time via [`DatasetPreFilter::with_overlay_block`]. overlay_block: Option, + /// Optional external row-address allow/block mask, combined with the + /// prefilter using logical AND. + external_mask: Option>, + /// Datafusion Plan Properties properties: Arc, @@ -1423,6 +1433,7 @@ impl ANNIvfSubIndexExec { query, prefilter_source, overlay_block: None, + external_mask: None, properties, metrics: ExecutionPlanMetricsSet::new(), }) @@ -1434,6 +1445,14 @@ impl ANNIvfSubIndexExec { self } + /// Restrict the ANN search to a caller-supplied row-address allow/block set. + /// Intersected with the prefilter, so top-k is computed over surviving rows + /// rather than filtered afterwards. No-op when `mask` is `None`. + pub fn with_external_mask(mut self, mask: Option>) -> Self { + self.external_mask = mask; + self + } + /// Returns a reference to the vector query. pub fn query(&self) -> &Query { &self.query @@ -1988,6 +2007,7 @@ impl ExecutionPlan for ANNIvfSubIndexExec { query: self.query.clone(), prefilter_source, overlay_block: self.overlay_block.clone(), + external_mask: self.external_mask.clone(), properties: self.properties.clone(), metrics: ExecutionPlanMetricsSet::new(), } @@ -2082,6 +2102,13 @@ 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 = { let mut pf = DatasetPreFilter::new(ds.clone(), &indices, prefilter_loader); if let Some(block) = self.overlay_block.clone() { 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 44092f75459..492b2f42d07 100644 --- a/rust/lance/src/io/exec/utils.rs +++ b/rust/lance/src/io/exec/utils.rs @@ -33,6 +33,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; @@ -53,6 +54,7 @@ pub(crate) fn build_prefilter( ds: Arc, index_meta: &[IndexMetadata], overlay_block: Option, + external_mask: Option>, ) -> Result> { let prefilter_loader = match &prefilter_source { PreFilterSource::FilteredRowIds(src_node) => { @@ -65,6 +67,16 @@ 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). Independent of `overlay_block`, which the prefilter + // applies separately to drop index entries staled by a data overlay. + let prefilter_loader = match external_mask { + Some(mask) => { + Some(Box::new(MaskAndLoader::new(mask, prefilter_loader)) as Box) + } + None => prefilter_loader, + }; let mut prefilter = DatasetPreFilter::new(ds, index_meta, prefilter_loader); if let Some(overlay_block) = overlay_block { prefilter = prefilter.with_overlay_block(overlay_block); From 4d7d156d081964099b4c7fe007f499115eea56eb Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Wed, 12 Aug 2026 03:01:52 -0700 Subject: [PATCH 2/3] feat(python): expose the row address prefilter on the scanner Adds `row_addr_allowlist` / `row_addr_blocklist` to `Dataset.scanner`, plus a `ScannerBuilder.row_addr_prefilter()` setter, so a caller can restrict a scan, a KNN search, or a full-text search to a precomputed set of row addresses. Both are serialized RowAddrTreeMap payloads rather than objects. Two Python extension modules each link their own copy of the lance crates and cannot share a Rust value, but they can agree on this encoding, so the mask may be built by a different module than the one that runs the scan. The bytes are decoded through RowAddrMask::from_serialized_parts, so no binding reimplements the allow/block combination. --- python/python/lance/dataset.py | 48 +++ .../python/tests/test_row_addr_prefilter.py | 312 ++++++++++++++++++ python/src/dataset.rs | 33 +- python/src/lib.rs | 2 + 4 files changed, 393 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_row_addr_prefilter.py diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index e1c5fa88d26..84354ac385a 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -69,6 +69,7 @@ _MergeInsertBuilder, _parse_field_path, _Scanner, + _serialize_row_addrs, _write_dataset, indices, ) @@ -1169,6 +1170,8 @@ def scanner( strict_batch_size: Optional[bool] = None, order_by: Optional[List[Union[ColumnOrdering, str]]] = None, disable_scoring_autoprojection: Optional[bool] = None, + row_addr_allowlist: Optional[bytes] = None, + row_addr_blocklist: Optional[bytes] = None, ) -> LanceScanner: """Return a Scanner that can support various pushdowns. @@ -1368,6 +1371,14 @@ def scanner( This parameter allows you to opt-in to the new behavior early, to avoid being subject to breaking changes in the future. + row_addr_allowlist: bytes, default None + Restrict the scan to these row addresses. A serialized roaring treemap + over ``_rowid`` (``RowAddrTreeMap::serialize_into`` output). Applied + before KNN / BM25 ranking, so top-k is computed over the surviving rows + rather than filtered afterwards. + row_addr_blocklist: bytes, default None + Exclude these row addresses, same encoding as ``row_addr_allowlist``. + Combined with it when both are given. .. note:: @@ -1410,6 +1421,8 @@ def setopt(opt, val): setopt(builder.filter, filter) setopt(builder.prefilter, prefilter) + if row_addr_allowlist is not None or row_addr_blocklist is not None: + builder.row_addr_prefilter(row_addr_allowlist, row_addr_blocklist) setopt(builder.limit, limit) setopt(builder.offset, offset) setopt(builder.batch_size, batch_size) @@ -6421,6 +6434,18 @@ def _needs_substrait_placeholder(t: pa.DataType) -> bool: return False +def serialize_row_addrs(addrs: Iterable[int]) -> bytes: + """Encode row addresses for ``row_addr_allowlist`` / ``row_addr_blocklist``. + + Those parameters take a serialized roaring treemap over ``_rowid``; this is + the way to produce one from Python. + + >>> blob = serialize_row_addrs([0, 2, 4]) # doctest: +SKIP + >>> ds.scanner(row_addr_allowlist=blob).to_table() # doctest: +SKIP + """ + return _serialize_row_addrs(list(addrs)) + + class ScannerBuilder: def __init__(self, ds: LanceDataset): self.ds = ds @@ -6429,6 +6454,8 @@ def __init__(self, ds: LanceDataset): self._search_filter = None self._substrait_filter = None self._prefilter = False + self._row_addr_allowlist: Optional[bytes] = None + self._row_addr_blocklist: Optional[bytes] = None self._late_materialization = None self._blob_handling = None self._offset = None @@ -6644,6 +6671,25 @@ def filter( return self + def row_addr_prefilter( + self, + allowlist: Optional[bytes] = None, + blocklist: Optional[bytes] = None, + ) -> ScannerBuilder: + """Restrict the scan to an externally supplied set of row addresses. + + allowlist / blocklist are serialized roaring treemaps over ``_rowid`` + (``RowAddrTreeMap::serialize_into`` output); passing neither clears the + mask. Applied before KNN / BM25 ranking, so top-k is computed over the + surviving rows rather than filtered afterwards. + + Bytes rather than an object so the mask can be produced by a different + extension module -- nothing Rust-typed crosses the boundary. + """ + self._row_addr_allowlist = allowlist + self._row_addr_blocklist = blocklist + return self + def prefilter(self, prefilter: bool) -> ScannerBuilder: self._prefilter = prefilter return self @@ -6936,6 +6982,8 @@ def to_scanner(self) -> LanceScanner: self._orderings, self._disable_scoring_autoprojection, self._substrait_aggregate, + self._row_addr_allowlist, + self._row_addr_blocklist, ) return LanceScanner(scanner, self.ds, _snapshot_scanner_builder(self)) diff --git a/python/python/tests/test_row_addr_prefilter.py b/python/python/tests/test_row_addr_prefilter.py new file mode 100644 index 00000000000..7d3381ebdce --- /dev/null +++ b/python/python/tests/test_row_addr_prefilter.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""End-to-end tests for the external row-address prefilter. + +``row_addr_allowlist`` / ``row_addr_blocklist`` restrict a scan to a set of row +addresses supplied by the caller, rather than to rows a filter expression +selects. The mask is applied before ranking, so a KNN or full-text search +computes top-k over the surviving rows instead of trimming the result +afterwards -- the two differ whenever k is smaller than the candidate set. + +Each test asserts against ``_rowid`` ground truth so a mask that is silently +dropped (which would return every row) fails rather than passing by accident. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING + +import lance +import numpy as np +import pyarrow as pa +import pytest +from lance.dataset import ScannerBuilder, serialize_row_addrs +from lance.file import LanceFileWriter + +if TYPE_CHECKING: + from pathlib import Path + +N = 256 +DIM = 8 + + +def _write(tmp_path: Path, with_index: bool = False) -> lance.LanceDataset: + rng = np.random.default_rng(1234) + vectors = rng.standard_normal((N, DIM)).astype(np.float32) + tbl = pa.table( + { + "id": pa.array(range(N), pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(vectors.reshape(-1), pa.float32()), DIM + ), + "text": pa.array([f"row {i} lorem ipsum" for i in range(N)]), + } + ) + ds = lance.write_dataset(tbl, str(tmp_path / "t.lance"), mode="overwrite") + if with_index: + # IVF_FLAT with nprobes == num_partitions is exact, so the masked result + # can be compared against brute force without recall slack. + ds.create_index("vector", index_type="IVF_FLAT", num_partitions=4, metric="l2") + return ds + + +def _rowids(ds: lance.LanceDataset) -> list[int]: + return ds.to_table(with_row_id=True)["_rowid"].to_pylist() + + +def test_serialize_row_addrs_round_trips_through_a_scan(tmp_path: Path) -> None: + ds = _write(tmp_path) + addrs = _rowids(ds) + want = addrs[3:9] + + got = ds.scanner( + with_row_id=True, row_addr_allowlist=serialize_row_addrs(want) + ).to_table() + assert got["_rowid"].to_pylist() == want + + +def test_allowlist_and_blocklist_combine(tmp_path: Path) -> None: + ds = _write(tmp_path) + addrs = _rowids(ds) + + allow, block = addrs[:10], addrs[5:15] + got = ds.scanner( + with_row_id=True, + row_addr_allowlist=serialize_row_addrs(allow), + row_addr_blocklist=serialize_row_addrs(block), + ).to_table() + assert got["_rowid"].to_pylist() == addrs[:5] + + # Block alone excludes and leaves everything else. + got = ds.scanner( + with_row_id=True, row_addr_blocklist=serialize_row_addrs(addrs[:5]) + ).to_table() + assert got["_rowid"].to_pylist() == addrs[5:] + + +def test_no_mask_reads_everything(tmp_path: Path) -> None: + # Guards the "no mask" vs "empty mask" distinction: omitting both must not + # be read as an allowlist of nothing. + ds = _write(tmp_path) + assert ds.scanner().to_table().num_rows == N + + +def test_empty_allowlist_selects_nothing(tmp_path: Path) -> None: + ds = _write(tmp_path) + got = ds.scanner(row_addr_allowlist=serialize_row_addrs([])).to_table() + assert got.num_rows == 0 + + +def test_mask_composes_with_a_filter(tmp_path: Path) -> None: + ds = _write(tmp_path) + addrs = _rowids(ds) + got = ds.scanner( + columns=["id"], + filter="id % 2 == 0", + row_addr_allowlist=serialize_row_addrs(addrs[:20]), + ).to_table() + assert got["id"].to_pylist() == [i for i in range(20) if i % 2 == 0] + + +def test_builder_setter_matches_the_kwarg(tmp_path: Path) -> None: + ds = _write(tmp_path) + blob = serialize_row_addrs(_rowids(ds)[2:7]) + from_kwarg = ds.scanner(with_row_id=True, row_addr_allowlist=blob).to_table() + from_builder = ( + ScannerBuilder(ds) + .with_row_id(True) + .row_addr_prefilter(allowlist=blob) + .to_scanner() + .to_table() + ) + assert from_kwarg["_rowid"].to_pylist() == from_builder["_rowid"].to_pylist() + + +@pytest.mark.parametrize("with_index", [False, True]) +def test_knn_topk_is_computed_over_masked_rows( + tmp_path: Path, with_index: bool +) -> None: + # The point of a prefilter: with k=5 and a 10-row mask, post-filtering a + # global top-5 would usually return fewer than 5 (often 0) rows. + ds = _write(tmp_path, with_index=with_index) + addrs = _rowids(ds) + allowed = addrs[100:110] + query = np.zeros(DIM, dtype=np.float32) + + got = ds.scanner( + nearest={"column": "vector", "q": query, "k": 5, "nprobes": 4}, + with_row_id=True, + row_addr_allowlist=serialize_row_addrs(allowed), + ).to_table() + + assert got.num_rows == 5 + assert set(got["_rowid"].to_pylist()) <= set(allowed) + + # Exactly the 5 nearest *within* the mask, not the global 5 intersected. + vectors = np.stack( + [np.asarray(v) for v in ds.to_table(columns=["vector"])["vector"].to_pylist()] + ) + by_addr = dict(zip(addrs, vectors)) + expect = sorted(allowed, key=lambda a: np.linalg.norm(by_addr[a] - query))[:5] + assert sorted(got["_rowid"].to_pylist()) == sorted(expect) + + +def test_knn_blocklist_excludes_the_nearest(tmp_path: Path) -> None: + ds = _write(tmp_path) + query = np.zeros(DIM, dtype=np.float32) + unmasked = ( + ds.scanner(nearest={"column": "vector", "q": query, "k": 3}, with_row_id=True) + .to_table()["_rowid"] + .to_pylist() + ) + + got = ds.scanner( + nearest={"column": "vector", "q": query, "k": 3}, + with_row_id=True, + row_addr_blocklist=serialize_row_addrs(unmasked[:1]), + ).to_table() + + assert got.num_rows == 3 # refilled, not truncated + assert unmasked[0] not in got["_rowid"].to_pylist() + + +def test_full_text_search_honors_the_mask(tmp_path: Path) -> None: + ds = _write(tmp_path) + ds.create_scalar_index("text", index_type="INVERTED") + addrs = _rowids(ds) + allowed = addrs[50:60] + + got = ds.scanner( + full_text_query="lorem", + with_row_id=True, + row_addr_allowlist=serialize_row_addrs(allowed), + limit=5, + ).to_table() + + assert got.num_rows == 5 + assert set(got["_rowid"].to_pylist()) <= set(allowed) + + +def test_rejects_a_malformed_mask(tmp_path: Path) -> None: + ds = _write(tmp_path) + with pytest.raises(Exception, match="(?i)row address mask|invalid"): + ds.scanner(row_addr_allowlist=b"not a treemap").to_table() + + +def _overlay( + ds, base_dir: Path, name: str, batch: pa.Table, fields: list[int], offsets +): + """Commit a data overlay covering `offsets` of fragment 0. + + An overlay committed after an index makes the indexed values stale, so the + planner replays those rows through a separate take. That replay is a second + row source, and it has to honor the caller's mask like every other one. + """ + path = base_dir / "data" / name + with LanceFileWriter(str(path)) as writer: + writer.write_batch(batch) + base_df = ds.get_fragments()[0].metadata.files[0] + data_file = lance.fragment.DataFile( + path=name, + fields=fields, + column_indices=list(range(len(fields))), + file_major_version=base_df.file_major_version, + file_minor_version=base_df.file_minor_version, + file_size_bytes=os.path.getsize(path), + ) + op = lance.LanceOperation.DataOverlay( + [ + lance.LanceOperation.DataOverlayGroup( + 0, [lance.LanceOperation.DataOverlayFile(data_file, offsets=offsets)] + ) + ] + ) + return lance.LanceDataset.commit(ds, op, read_version=ds.version) + + +def test_overlay_stale_replay_scan_respects_mask(tmp_path: Path) -> None: + base_dir = tmp_path / "ov_scan" + ds = lance.write_dataset( + pa.table( + { + "id": pa.array(range(10), pa.int32()), + "val": pa.array([i * 10 for i in range(10)], pa.int32()), + } + ), + base_dir, + ) + # Index first, then overlay: offset 1 now reads 999 while the index still + # says 10, so `val = 999` can only be answered by the stale replay. + ds.create_scalar_index("val", index_type="BTREE") + ds = _overlay( + ds, + base_dir, + "ov.lance", + pa.table({"val": pa.array([999], pa.int32())}), + fields=[1], + offsets=[1], + ) + + base = ds.scanner(filter="val = 999", with_row_id=True).to_table() + assert base.num_rows == 1, "fixture did not produce a stale replay" + stale_addr = base["_rowid"].to_pylist()[0] + + got = ds.scanner( + filter="val = 999", row_addr_allowlist=serialize_row_addrs([]) + ).to_table() + assert got.num_rows == 0, "stale replay must not return rows the mask excludes" + + got = ds.scanner( + filter="val = 999", + with_row_id=True, + row_addr_allowlist=serialize_row_addrs([stale_addr]), + ).to_table() + assert got["_rowid"].to_pylist() == [stale_addr] + + +def test_overlay_stale_replay_ann_respects_mask(tmp_path: Path) -> None: + base_dir = tmp_path / "ov_ann" + rng = np.random.default_rng(7) + vectors = rng.standard_normal((N, DIM)).astype(np.float32) + ds = lance.write_dataset( + pa.table( + { + "id": pa.array(range(N), pa.int64()), + "vector": pa.FixedSizeListArray.from_arrays( + pa.array(vectors.reshape(-1), pa.float32()), DIM + ), + } + ), + base_dir, + ) + ds.create_index("vector", index_type="IVF_FLAT", num_partitions=4, metric="l2") + + # Move two rows onto the query point after indexing. The ANN index still has + # their old vectors, so they can only surface through the stale replay. + query = np.zeros(DIM, dtype=np.float32) + moved = pa.FixedSizeListArray.from_arrays( + pa.array(np.zeros(2 * DIM, dtype=np.float32), pa.float32()), DIM + ) + ds = _overlay( + ds, + base_dir, + "ov_vec.lance", + pa.table({"vector": moved}), + fields=[1], + offsets=[3, 7], + ) + + base = ds.scanner( + nearest={"column": "vector", "q": query, "k": 5}, with_row_id=True + ).to_table() + assert base.num_rows > 0, "fixture did not produce ANN results" + + got = ds.scanner( + nearest={"column": "vector", "q": query, "k": 5}, + row_addr_allowlist=serialize_row_addrs([]), + ).to_table() + assert got.num_rows == 0, ( + "the ANN stale replay must not return rows the mask excludes" + ) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index be568ecd537..365b0ed390b 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -41,7 +41,7 @@ use lance::dataset::cleanup::{CleanupFileKind, CleanupPolicyBuilder}; use lance::dataset::refs::{Ref, TagContents}; use lance::dataset::scanner::{ AggregateExpr, ColumnOrdering, DatasetRecordBatchStream, ExecutionStatsCallback, - MaterializationStyle, QueryFilter, + MaterializationStyle, QueryFilter, RowAddrMask, RowAddrTreeMap, }; use lance::dataset::statistics::{DataStatistics, DatasetStatisticsExt}; use lance::dataset::{ @@ -1178,7 +1178,7 @@ impl Dataset { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature=(columns=None, columns_with_transform=None, filter=None, search_filter=None, prefilter=None, limit=None, offset=None, nearest=None, batch_size=None, batch_size_bytes=None, io_buffer_size=None, batch_readahead=None, fragment_readahead=None, scan_in_order=None, fragments=None, index_segments=None, with_row_id=None, with_row_address=None, use_stats=None, substrait_filter=None, fast_search=None, full_text_query=None, late_materialization=None, blob_handling=None, use_scalar_index=None, include_deleted_rows=None, scan_stats_callback=None, strict_batch_size=None, order_by=None, disable_scoring_autoprojection=None, substrait_aggregate=None))] + #[pyo3(signature=(columns=None, columns_with_transform=None, filter=None, search_filter=None, prefilter=None, limit=None, offset=None, nearest=None, batch_size=None, batch_size_bytes=None, io_buffer_size=None, batch_readahead=None, fragment_readahead=None, scan_in_order=None, fragments=None, index_segments=None, with_row_id=None, with_row_address=None, use_stats=None, substrait_filter=None, fast_search=None, full_text_query=None, late_materialization=None, blob_handling=None, use_scalar_index=None, include_deleted_rows=None, scan_stats_callback=None, strict_batch_size=None, order_by=None, disable_scoring_autoprojection=None, substrait_aggregate=None, row_addr_allowlist=None, row_addr_blocklist=None))] fn scanner( self_: PyRef<'_, Self>, columns: Option>, @@ -1212,6 +1212,8 @@ impl Dataset { order_by: Option>>, disable_scoring_autoprojection: Option, substrait_aggregate: Option>, + row_addr_allowlist: Option>, + row_addr_blocklist: Option>, ) -> PyResult { let mut scanner: LanceScanner = self_.ds.scan(); @@ -1347,6 +1349,18 @@ impl Dataset { if let Some(prefilter) = prefilter { scanner.prefilter(prefilter); } + // Serialized RowAddrTreeMap payloads rather than an object: a mask built by + // another extension module cannot hand over a Rust value, but both sides + // agree on this encoding. RowAddrMask::from_serialized_parts is the shared + // entry point, so no binding has to reimplement the allow/block combination. + if let Some(mask) = RowAddrMask::from_serialized_parts( + row_addr_allowlist.as_deref(), + row_addr_blocklist.as_deref(), + ) + .infer_error()? + { + scanner.with_row_addr_prefilter(mask); + } scanner .limit(limit, offset) @@ -4583,6 +4597,21 @@ impl Dataset { } } +/// Serialize row addresses into the payload the scanner's `row_addr_allowlist` / +/// `row_addr_blocklist` parameters accept. +/// +/// Without this those parameters are unusable from Python: they take the roaring +/// `RowAddrTreeMap` encoding, which nothing else exposed here can produce. The +/// result stays plain bytes, so a mask may equally be built by another extension +/// module and handed in. +#[pyfunction(name = "_serialize_row_addrs")] +pub fn serialize_row_addrs(py: Python<'_>, addrs: Vec) -> PyResult> { + let treemap = RowAddrTreeMap::from_iter(addrs); + let mut buf = Vec::with_capacity(treemap.serialized_size()); + treemap.serialize_into(&mut buf).infer_error()?; + Ok(PyBytes::new(py, &buf).unbind()) +} + #[pyfunction(name = "_write_dataset")] pub fn write_dataset( reader: &Bound<'_, PyAny>, diff --git a/python/src/lib.rs b/python/src/lib.rs index ca76149b413..a74015a8b00 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -94,6 +94,7 @@ pub use crate::tracing::{TraceGuard, trace_to_chrome}; use crate::utils::Hnsw; use crate::utils::KMeans; pub use dataset::Dataset; +pub use dataset::serialize_row_addrs; pub use dataset::write_dataset; use fragment::{FileFragment, PyDeletionFile, PyRowDatasetVersionMeta, PyRowIdMeta}; pub use indices::register_indices; @@ -320,6 +321,7 @@ fn lance(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_wrapped(wrap_pyfunction!(mem_wal::py_write_pk_sidecar))?; m.add_wrapped(wrap_pyfunction!(bfloat16_array))?; m.add_wrapped(wrap_pyfunction!(write_dataset))?; + m.add_wrapped(wrap_pyfunction!(serialize_row_addrs))?; m.add_wrapped(wrap_pyfunction!(write_fragments))?; m.add_wrapped(wrap_pyfunction!(write_fragments_transaction))?; m.add_wrapped(wrap_pyfunction!(schema_to_json))?; From c372d131669b3748e7c49b2ccc96fd022d7dfafe Mon Sep 17 00:00:00 2001 From: Yuan Gao Date: Fri, 21 Aug 2026 22:38:46 -0700 Subject: [PATCH 3/3] test(scanner): cover the cross-column FTS scorer with the external mask The cross-column compound scorer is a separate exec from the same-column one and builds its own prefilter, so the external row-address mask has to reach it independently. Without it the scorer returns rows the caller excluded. --- rust/lance/src/dataset/scanner.rs | 89 +++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index edaa6d6a25f..084a9601316 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -7575,6 +7575,95 @@ mod test { ); } + /// A cross-column boolean query plans into CrossColumnCompoundFtsScorer, + /// which is a different exec from the same-column CompoundFtsScorer and + /// builds its own prefilter, so it needs the mask threaded separately. + #[tokio::test] + async fn row_addr_mask_cross_column_fts_respects_mask() { + use lance_index::scalar::inverted::tokenizer::InvertedIndexParams; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("title", DataType::Utf8, true), + ArrowField::new("body", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from_iter_values( + (0..64).map(|v| format!("alpha title {v}")), + )), + Arc::new(StringArray::from_iter_values( + (0..64).map(|v| format!("alpha body {v}")), + )), + ], + ) + .unwrap(); + + let path = TempStrDir::default(); + let reader = RecordBatchIterator::new([Ok(batch)], schema.clone()); + let mut dataset = Dataset::write(reader, &path, None).await.unwrap(); + let params = InvertedIndexParams::default() + .with_position(true) + .remove_stop_words(false); + for column in ["title", "body"] { + dataset + .create_index(&[column], IndexType::Inverted, None, ¶ms, true) + .await + .unwrap(); + } + + // Two leaves on different columns is what selects the cross-column + // scorer; a bounded limit is required by that exec. + let cross_column = || { + FullTextSearchQuery::new_query(FtsQuery::Boolean(BooleanQuery::new([ + ( + Occur::Should, + MatchQuery::new("title".to_string()) + .with_column(Some("title".to_string())) + .into(), + ), + ( + Occur::Should, + MatchQuery::new("body".to_string()) + .with_column(Some("body".to_string())) + .into(), + ), + ]))) + .limit(Some(10)) + }; + + let mut scan = dataset.scan(); + scan.full_text_search(cross_column()).unwrap(); + scan.with_row_id(); + let plan = scan.explain_plan(true).await.unwrap(); + assert!( + plan.contains("CrossColumnCompoundFtsScorer"), + "expected the cross-column compound scorer path, got:\n{plan}" + ); + let base = batch_row_ids(&scan.try_into_batch().await.unwrap()); + assert!(!base.is_empty(), "cross-column query matched nothing"); + + let mut scan = dataset.scan(); + scan.full_text_search(cross_column()).unwrap(); + scan.with_row_id(); + scan.with_row_addr_prefilter(RowAddrMask::allow_nothing()); + assert_eq!( + scan.try_into_batch().await.unwrap().num_rows(), + 0, + "the cross-column scorer must not return rows the mask excludes" + ); + + let keep = base[0]; + let mut scan = dataset.scan(); + scan.full_text_search(cross_column()).unwrap(); + scan.with_row_id(); + scan.with_row_addr_prefilter(RowAddrMask::from_allowed(RowAddrTreeMap::from_iter([keep]))); + assert_eq!( + batch_row_ids(&scan.try_into_batch().await.unwrap()), + vec![keep] + ); + } + #[tokio::test] async fn row_addr_mask_fts_search_only_allowed() { let mut test_ds = TestVectorDataset::new(LanceFileVersion::Stable, false)