diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index fa597c5cd18..2bd497da61b 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -508,11 +508,29 @@ pub struct SearchOptions { /// Callers may disable this only when NULL rows cannot affect the final /// result, such as a top-level filter whose NULL results will be discarded. track_nulls: bool, + /// Best-effort cap on how many matching rows the search needs to find. + /// + /// An index may stop once it has `n` matches and may still return more, so + /// this narrows work without narrowing the contract. It applies only to + /// positive lookups that preserve matches, such as equality, range, and `IsIn`. + /// Negating and combining operators ignore it, because a partial match set + /// cannot be complemented or intersected soundly. + /// + /// A search that acts on the limit reports [`SearchResult::AtLeast`], never + /// [`SearchResult::Exact`], so a short-circuited scan can never be mistaken + /// for the complete match set. An index may also decline the hint and search + /// in full, in which case it still reports [`SearchResult::Exact`]. Notably + /// [`track_nulls`](Self::track_nulls) forces this, because a partial scan + /// cannot report a complete null set. + limit: Option, } impl Default for SearchOptions { fn default() -> Self { - Self { track_nulls: true } + Self { + track_nulls: true, + limit: None, + } } } @@ -528,6 +546,17 @@ impl SearchOptions { pub fn track_nulls(&self) -> bool { self.track_nulls } + + /// Set a best-effort cap on the number of matching rows the search must find. + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + /// The best-effort cap on matching rows, if any. + pub fn limit(&self) -> Option { + self.limit + } } /// A trait for a scalar index, a structure that can determine row ids that satisfy scalar queries diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index cb5f8818b59..5a875ce835a 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2245,13 +2245,43 @@ impl ScalarIndex for BTreeIndex { .collect::>(); debug!("Searching {} btree pages", page_tasks.len()); - // Collect both matching row IDs and null row IDs from all pages - let results: Vec = stream::iter(page_tasks) - // I/O and compute mixed here but important case is index in cache so - // use compute intensive thread count - .buffered(get_num_compute_intensive_cpus()) - .try_collect() - .await?; + // Null tracking and an early stop are incompatible. The null pages above are appended + // after the matching pages, and an unread matching page can hold nulls of its own, so a + // short-circuited scan cannot report a complete null set. Decline the hint rather than + // the search: the limit is only an optimization and must never turn an otherwise valid + // search into a failure or a wrong answer. + let limit = if options.track_nulls() { + None + } else { + options.limit() + }; + + // Collect both matching row IDs and null row IDs from all pages. + // + // With a limit, read one page at a time and stop as soon as enough TRUE matches + // have been seen. Fanning out would defeat the point by doing the work anyway. + // Without one, fan out across CPUs as usual. + let parallelism = if limit.is_some() { + 1 + } else { + get_num_compute_intensive_cpus() + }; + let mut page_stream = stream::iter(page_tasks).buffered(parallelism); + let mut results: Vec = Vec::new(); + let mut matches_found: u64 = 0; + while let Some(page_result) = page_stream.try_next().await? { + if let Some(limit) = limit { + // A limit implies nulls are not tracked, so every selected row is a TRUE match + // and `selected_rows()` counts them without the set arithmetic `len()` does. + matches_found += page_result.selected_rows().len().unwrap_or(0); + results.push(page_result); + if matches_found >= limit as u64 { + break; + } + } else { + results.push(page_result); + } + } let selection = if options.track_nulls() { NullableRowAddrSet::union_all(&results) @@ -2266,7 +2296,17 @@ impl ScalarIndex for BTreeIndex { ) }; - Ok(SearchResult::Exact(selection)) + // A limited search may stop before reading every matching page, so the returned set + // is a lower bound rather than the complete match set. Reporting `Exact` would let a + // downstream consumer treat a partial set as complete, so report `AtLeast` whenever a + // limit was in play, even if this particular scan happened to read every page, + // since the caller cannot tell the difference and must not rely on it. A declined + // limit read every page as usual, so it stays `Exact`. + Ok(if limit.is_some() { + SearchResult::AtLeast(selection) + } else { + SearchResult::Exact(selection) + }) } fn can_remap(&self) -> bool { @@ -5605,6 +5645,168 @@ mod tests { } } + /// A limited search returns at least `limit` matches but stops reading pages early, so + /// for a multi-page range it reads fewer pages than an unlimited search and must report + /// `AtLeast` rather than `Exact`. + #[tokio::test] + async fn test_search_limited_short_circuits() { + use arrow_array::{Int32Array, UInt64Array}; + + let tmpdir = TempObjDir::default(); + let test_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + // Five btree pages of `DEFAULT_BTREE_BATCH_SIZE` rows, no nulls, so every row matches. + let num_rows = 5 * DEFAULT_BTREE_BATCH_SIZE; + let values: Int32Array = (0..num_rows).map(|i| Some(i as i32)).collect(); + let row_ids = UInt64Array::from_iter_values(0..num_rows); + let data = arrow_array::RecordBatch::try_from_iter(vec![ + ("value", Arc::new(values) as arrow_array::ArrayRef), + ("_rowid", Arc::new(row_ids) as arrow_array::ArrayRef), + ]) + .unwrap(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + data.schema(), + stream::iter(vec![Ok(data)]), + )); + train_btree_index( + stream, + test_store.as_ref(), + DEFAULT_BTREE_BATCH_SIZE, + None, + None, + ) + .await + .unwrap(); + + let index = BTreeIndex::load(test_store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + let metrics = NoOpMetricsCollector; + let everything = + SargableQuery::Range(std::ops::Bound::Unbounded, std::ops::Bound::Unbounded); + + // Baseline: an unlimited search returns every row and is exact. + let full = index.search(&everything, &metrics).await.unwrap(); + assert!( + matches!(full, SearchResult::Exact(_)), + "an unlimited search must stay Exact, got {full:?}" + ); + let full_len = full.row_addrs().len().unwrap(); + assert_eq!(full_len, num_rows); + + // A limit reaching into the second page: satisfied, but stops before all five. + let limit = (DEFAULT_BTREE_BATCH_SIZE + 100) as usize; + let limited = index + .search_with_options( + &everything, + SearchOptions::default() + .with_track_nulls(false) + .with_limit(Some(limit)), + &metrics, + ) + .await + .unwrap(); + // A short-circuited search is not the complete match set, so it must be reported as + // `AtLeast` (a lower bound), never `Exact`. + assert!( + matches!(limited, SearchResult::AtLeast(_)), + "limited search must return AtLeast, got {limited:?}" + ); + let limited_len = limited.row_addrs().len().unwrap(); + assert!( + limited_len >= limit as u64, + "expected at least {limit} matches, got {limited_len}" + ); + assert!( + limited_len < full_len, + "expected the search to short-circuit, but it returned all {full_len} rows" + ); + } + + /// Null tracking and an early stop cannot both be honored: null pages are searched after + /// the matching pages, and an unread matching page can hold nulls too, so a short-circuited + /// scan could not report a complete null set. The limit is declined instead of the search, + /// so the result stays `Exact` and its nulls match an unlimited search exactly. + #[tokio::test] + async fn test_limit_is_declined_when_nulls_are_tracked() { + use arrow_array::{Int32Array, UInt64Array}; + + let tmpdir = TempObjDir::default(); + let test_store = Arc::new(LanceIndexStore::new( + Arc::new(ObjectStore::local()), + tmpdir.clone(), + Arc::new(LanceCache::no_cache()), + )); + + // Five pages where every fifth row is NULL, so nulls are spread across every page and + // a search that stopped early would miss the ones on the pages it never read. + let num_rows = 5 * DEFAULT_BTREE_BATCH_SIZE; + let values: Int32Array = (0..num_rows) + .map(|i| if i % 5 == 0 { None } else { Some(i as i32) }) + .collect(); + let row_ids = UInt64Array::from_iter_values(0..num_rows); + let data = arrow_array::RecordBatch::try_from_iter_with_nullable(vec![ + ("value", Arc::new(values) as arrow_array::ArrayRef, true), + ("_rowid", Arc::new(row_ids) as arrow_array::ArrayRef, false), + ]) + .unwrap(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + data.schema(), + stream::iter(vec![Ok(data)]), + )); + train_btree_index( + stream, + test_store.as_ref(), + DEFAULT_BTREE_BATCH_SIZE, + None, + None, + ) + .await + .unwrap(); + + let index = BTreeIndex::load(test_store.clone(), None, &LanceCache::no_cache()) + .await + .unwrap(); + let metrics = NoOpMetricsCollector; + let everything = + SargableQuery::Range(std::ops::Bound::Unbounded, std::ops::Bound::Unbounded); + + let track_nulls = |limit| { + index.search_with_options( + &everything, + SearchOptions::default() + .with_track_nulls(true) + .with_limit(limit), + &metrics, + ) + }; + + let unlimited = track_nulls(None).await.unwrap(); + // A limit small enough that acting on it would stop on the first page. + let limited = track_nulls(Some(10)).await.unwrap(); + + assert!( + matches!(limited, SearchResult::Exact(_)), + "a declined limit searched every page, so the result must stay Exact, got {limited:?}" + ); + let unlimited = unlimited.row_addrs(); + let limited = limited.row_addrs(); + assert_eq!( + limited.null_rows(), + unlimited.null_rows(), + "tracked nulls must be complete even when a limit was requested" + ); + assert!( + !limited.null_rows().is_empty(), + "the fixture must produce nulls, otherwise this asserts nothing" + ); + assert_eq!(limited.len(), unlimited.len()); + } + /// Regression test: disabling NULL tracking also omits NULL rows from a /// candidate page that contains both matching and NULL values. #[tokio::test] diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index 5662dd4ef7e..39ea8fc4f68 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1942,8 +1942,12 @@ impl ScalarIndexExpr { index_loader: &dyn ScalarIndexLoader, metrics: &dyn MetricsCollector, ) -> Result { - self.evaluate_with_options(index_loader, metrics, true) - .await + self.evaluate_with_options( + index_loader, + metrics, + SearchOptions::default().with_track_nulls(true), + ) + .await } #[async_recursion] @@ -1951,26 +1955,37 @@ impl ScalarIndexExpr { &self, index_loader: &dyn ScalarIndexLoader, metrics: &dyn MetricsCollector, - track_nulls: bool, + options: SearchOptions, ) -> Result { match self { Self::Not(inner) => { // NOT needs the child's NULL rows to preserve SQL three-valued // logic. Once enabled, keep tracking through the whole subtree. + // The limit is dropped for the same reason it cannot cross any + // combining operator: a partial match set cannot be complemented. let result = inner - .evaluate_with_options(index_loader, metrics, true) + .evaluate_with_options( + index_loader, + metrics, + options.with_track_nulls(true).with_limit(None), + ) .await?; Ok(!result) } Self::And(lhs, rhs) => { - let lhs_result = lhs.evaluate_with_options(index_loader, metrics, track_nulls); - let rhs_result = rhs.evaluate_with_options(index_loader, metrics, track_nulls); + // A limit must not cross AND/OR: each side would stop early at a + // different subset, and the intersection or union of two partial + // sets is not a prefix of the full result. + let options = options.with_limit(None); + let lhs_result = lhs.evaluate_with_options(index_loader, metrics, options); + let rhs_result = rhs.evaluate_with_options(index_loader, metrics, options); let (lhs_result, rhs_result) = try_join!(lhs_result, rhs_result)?; Ok(lhs_result & rhs_result) } Self::Or(lhs, rhs) => { - let lhs_result = lhs.evaluate_with_options(index_loader, metrics, track_nulls); - let rhs_result = rhs.evaluate_with_options(index_loader, metrics, track_nulls); + let options = options.with_limit(None); + let lhs_result = lhs.evaluate_with_options(index_loader, metrics, options); + let rhs_result = rhs.evaluate_with_options(index_loader, metrics, options); let (lhs_result, rhs_result) = try_join!(lhs_result, rhs_result)?; Ok(lhs_result | rhs_result) } @@ -1979,11 +1994,7 @@ impl ScalarIndexExpr { .load_index(&search.column, &search.index_name, metrics) .await?; let search_result = index - .search_with_options( - search.query.as_ref(), - SearchOptions::default().with_track_nulls(track_nulls), - metrics, - ) + .search_with_options(search.query.as_ref(), options, metrics) .await?; let result = search_result_to_nullable(search_result); if index.results_are_row_addresses() { @@ -1998,6 +2009,33 @@ impl ScalarIndexExpr { } } + /// Like [`Self::evaluate`] but pushes a best-effort `limit` into the index search so + /// it can stop once it has found at least `limit` matches. + /// + /// The limit reaches only a top-level positive lookup. It is dropped at any `Not`, + /// `And`, or `Or` node, because a partial match set cannot be complemented, + /// intersected, or unioned soundly. A limited search reports + /// [`SearchResult::AtLeast`], so callers must still enforce the exact limit + /// downstream. + #[instrument(level = "debug", skip_all)] + pub async fn evaluate_limited( + &self, + index_loader: &dyn ScalarIndexLoader, + metrics: &dyn MetricsCollector, + limit: Option, + ) -> Result { + Ok(self + .evaluate_with_options( + index_loader, + metrics, + SearchOptions::default() + .with_track_nulls(false) + .with_limit(limit), + ) + .await? + .drop_nulls()) + } + #[instrument(level = "debug", skip_all)] pub async fn evaluate( &self, @@ -2005,7 +2043,11 @@ impl ScalarIndexExpr { metrics: &dyn MetricsCollector, ) -> Result { Ok(self - .evaluate_with_options(index_loader, metrics, false) + .evaluate_with_options( + index_loader, + metrics, + SearchOptions::default().with_track_nulls(false), + ) .await? .drop_nulls()) } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index b16f2a6f3b4..04521bd85b6 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -3265,6 +3265,24 @@ impl Scanner { scan_range: Option>, session: Option<&dyn Session>, ) -> Result> { + // Decide whether a limit can be pushed into the index search. The fragments the read + // covers are the requested subset, or the whole dataset when none was given. The index + // is only allowed to stop early when every fragment it covers is in this scanned set + // and free of deletions (see `index_search_limit`). + // + // `scan_range` is a separate limit/offset pushdown that only applies when there is no + // filter (see the `filter_plan.is_empty()` guard at its only call site). With no filter + // there is no index query, so `index_search_limit` returns `None`. The two pushdowns are + // mutually exclusive and need no extra coordination here. + let all_fragments = self.dataset.fragments(); + let scanned_fragments: &[Fragment] = fragments + .as_ref() + .map(|frags| frags.as_slice()) + .unwrap_or_else(|| all_fragments.as_slice()); + let mut pushdown_limit = self + .index_search_limit(filter_plan, scanned_fragments) + .await?; + // 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(); @@ -3330,6 +3348,11 @@ impl Scanner { .await?; if let Some(block) = self.stale_rows_block_mask(&overlay_stale_rows).await? { read_options = read_options.with_overlay_block(block); + // The block drops index-matched rows *after* the search, exactly like a + // deletion file. An early stop would then spend its budget on rows that are + // subsequently blocked and could yield fewer than `limit` live rows, so the + // limit cannot be pushed once anything is masked. + pushdown_limit = None; } } @@ -3337,11 +3360,10 @@ impl Scanner { 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 + Arc::new( + ScalarIndexExec::new(self.dataset.clone(), index_query, result_format) + .with_limit(pushdown_limit), + ) as Arc }), }; @@ -5648,6 +5670,95 @@ impl Scanner { )?)) } + /// Compute the limit hint that can be safely pushed into a scalar index search. + /// + /// Pushing a limit is only an optimization. A `GlobalLimitExec` still applies the + /// exact limit and offset, so the index only needs to return at least `limit + offset` + /// rows. The first N matches are as good as any N matches only when all of these hold. + /// + /// - There is a positive row limit. + /// - The scan is unordered (`scan_in_order(false)`). The default ordered mode returns + /// matches in storage order, but a B-tree stops in index-value order, so pushing the + /// limit would change which rows `LIMIT`/`OFFSET` returns. + /// - The rows are not reordered before the limit (no `ORDER BY`, vector or FTS search). + /// - There is no aggregate (the limit applies after aggregation). + /// - The index result is used as is, with no refine filter and no recheck. Either of + /// those re-filters rows later and could drop matches. + /// - The index query is a single lookup, not a combination. A limit only stops a lone + /// `Query` early. `And`/`Or`/`Not` always compute the full result, so a pushed limit + /// has no effect there. + /// - Every fragment the index covers is in the scanned set and has no deletion file. + /// The index returns row addresses for every fragment it covers, and any that do not + /// survive into the result are pruned after the search. An early stop would then spend + /// its budget on rows that get dropped, leaving fewer than `limit` live rows. Rows are + /// dropped when their fragment has deletions or is not scanned (a retired fragment the + /// index still has stale entries for, or one excluded by `with_fragments`). + /// + /// Returns `None` when no limit can be pushed. + async fn index_search_limit( + &self, + filter_plan: &ExprFilterPlan, + scanned_fragments: &[Fragment], + ) -> Result> { + let Some(limit) = self.limit else { + return Ok(None); + }; + if limit <= 0 { + return Ok(None); + } + // Ordered scans return storage-order matches, while a B-tree stops in index-value order, + // so the two would return different subsets. The other modes reorder or re-filter the + // rows after the index search, which can also drop early-collected matches. + if self.ordered + || self.ordering.is_some() + || self.nearest.is_some() + || self.full_text_query.is_some() + || self.aggregate.is_some() + || filter_plan.has_refine() + { + return Ok(None); + } + let Some(index_query) = filter_plan.index_query.as_ref() else { + return Ok(None); + }; + if index_query.needs_recheck() { + return Ok(None); + } + // Only a single top-level index lookup stops early. `evaluate_nullable` drops the limit for + // And/Or/Not, so those always compute the full result. Pushing a limit there would have no + // effect and would rely on fragment coverage that is unsound for Or, since an Or can read + // the union of both sides' fragments while coverage intersects them. Restrict the pushdown + // to a lone `Query`, whose coverage is just that one index's fragment bitmap. + let ScalarIndexExpr::Query(search) = index_query else { + return Ok(None); + }; + // Every row address the index covers must survive into the result, otherwise an early + // stop could leave fewer than `limit` live rows. That requires every index-covered + // fragment to be both scanned and free of deletions. Fragments that are scanned but not + // covered by the index are fine. They only add rows via a separate scan of the missing + // fragments, and never remove index hits. + let live_undeleted: RoaringBitmap = scanned_fragments + .iter() + .filter(|fragment| fragment.deletion_file.is_none()) + .map(|fragment| fragment.id as u32) + .collect(); + // A legacy or unsupported index reports no fragment bitmap, so its coverage is unknown and + // an early stop cannot be proven safe. Treat unknown coverage as "do not push" rather than + // an error: this is only an optimization and must never turn an otherwise valid scan into a + // failure. + let Some(covered_frags) = + scalar_index_fragment_bitmap(self.dataset.as_ref(), &search.column, &search.index_name) + .await? + else { + return Ok(None); + }; + if !covered_frags.is_subset(&live_undeleted) { + return Ok(None); + } + let offset = self.offset.unwrap_or(0) as usize; + Ok(Some((limit as usize).saturating_add(offset))) + } + // First perform a lookup in a scalar index for ids and then perform a take on the // target fragments with those ids async fn scalar_indexed_scan( @@ -5673,6 +5784,14 @@ impl Scanner { .partition_frags_by_coverage(index_expr, fragments) .await?; + // A limit can be pushed into the index search only when safe. See index_search_limit. + // `relevant_frags` is the intersection of covered and scanned fragments, so requiring the + // index's covered fragments to be a subset of it rejects both retired or uncovered scanned + // fragments and `with_fragments` subsets that drop covered fragments. + let pushdown_limit = self + .index_search_limit(filter_plan, &relevant_frags) + .await?; + // Build the MaterializeIndexExec, blocking stale row addresses so the index never // emits them. Stale rows are re-scored separately via a targeted take below. let mat_exec = MaterializeIndexExec::new( @@ -5680,11 +5799,21 @@ impl Scanner { index_expr.clone(), Arc::new(relevant_frags), ); - let mat_exec = match self.stale_rows_block_mask(&stale_rows).await? { + let overlay_block = self.stale_rows_block_mask(&stale_rows).await?; + // An overlay block drops index-matched rows *after* the search, exactly like a + // deletion file does. An early stop would then spend its budget on rows that are + // subsequently blocked and could return fewer than `limit` live rows, so refuse to + // push a limit whenever any stale row is being masked. + let pushdown_limit = if overlay_block.is_some() { + None + } else { + pushdown_limit + }; + let mat_exec = match overlay_block { Some(block) => mat_exec.with_overlay_block(block), None => mat_exec, }; - let mut plan: Arc = Arc::new(mat_exec); + let mut plan: Arc = Arc::new(mat_exec.with_limit(pushdown_limit)); let refine_expr = filter_plan.refine_expr.as_ref(); @@ -8374,6 +8503,463 @@ mod test { assert_eq!(ids, &(10..20).collect::>()); } + #[rstest] + #[tokio::test] + async fn test_limit_pushed_into_scalar_index( + // Legacy storage routes through `scalar_indexed_scan` (MaterializeIndexExec), Stable + // through `new_filtered_read` (ScalarIndexExec). Both push the limit via the same gate. + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + ) { + // A scalar-index limit can be pushed only for an unordered scan, since the B-tree stops + // in index-value order while an ordered scan returns storage-order matches. + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + // Span several btree pages, with ids in descending order so storage order is the reverse + // of index-value order. + let num_rows = 20_000i32; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values((0..num_rows).rev()))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let write_params = WriteParams { + data_storage_version: Some(data_storage_version), + ..Default::default() + }; + let mut dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + let limit = 100i64; + let scan_ids = |dataset: Arc, ordered: bool, offset: Option| async move { + let mut scan = dataset.scan(); + scan.filter("id >= 5") + .unwrap() + .scan_in_order(ordered) + .limit(Some(limit), offset) + .unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }; + + // Row counts alone cannot tell a pushed limit from a full scan that happens to + // return the right rows, so assert on the plan too. Both index exec nodes render + // the pushed limit, and exactly one of them runs depending on storage version. + let plan_for = |dataset: Arc, ordered: bool, offset: Option| async move { + let mut scan = dataset.scan(); + scan.filter("id >= 5") + .unwrap() + .scan_in_order(ordered) + .limit(Some(limit), offset) + .unwrap(); + scan.explain_plan(true).await.unwrap() + }; + + // Ordered scan (the default): limit not pushed, so the first matches are the largest ids + // (descending storage). + let ids = scan_ids(Arc::new(dataset.clone()), true, None).await; + assert_eq!(ids.len(), limit as usize); + assert!( + ids.iter().all(|&id| id >= num_rows - limit as i32), + "ordered scan must return the storage-order subset, got {:?}", + &ids[..ids.len().min(5)] + ); + let ordered_plan = plan_for(Arc::new(dataset.clone()), true, None).await; + assert!( + !ordered_plan.contains("limit="), + "an ordered scan must not push a limit into the index, plan was:\n{ordered_plan}" + ); + + // Unordered scan: limit pushed into the index, but still exactly `limit` matching rows. + let ids = scan_ids(Arc::new(dataset.clone()), false, None).await; + assert_eq!(ids.len(), limit as usize); + assert!( + ids.iter().all(|&id| id >= 5), + "every returned row must satisfy the filter" + ); + let unordered_plan = plan_for(Arc::new(dataset.clone()), false, None).await; + assert!( + unordered_plan.contains("limit=100"), + "an unordered scan must push limit=100 into the index, plan was:\n{unordered_plan}" + ); + + // With an offset the pushed limit must cover `limit + offset` rows, otherwise the + // downstream skip would leave fewer than `limit`. This guards the `saturating_add(offset)`. + let ids = scan_ids(Arc::new(dataset.clone()), false, Some(50)).await; + assert_eq!( + ids.len(), + limit as usize, + "offset must not reduce the returned row count" + ); + assert!(ids.iter().all(|&id| id >= 5)); + let offset_plan = plan_for(Arc::new(dataset.clone()), false, Some(50)).await; + assert!( + offset_plan.contains("limit=150"), + "an offset of 50 must widen the pushed limit to 150, plan was:\n{offset_plan}" + ); + + // With deletions the limit must not be pushed even when unordered, since deleted rows are + // pruned after the index search. + dataset.delete("id >= 5 AND id < 10000").await.unwrap(); + let ids = scan_ids(Arc::new(dataset.clone()), false, None).await; + assert_eq!(ids.len(), limit as usize); + assert!( + ids.iter().all(|&id| id >= 10000), + "deleted rows must not be returned" + ); + let deleted_plan = plan_for(Arc::new(dataset), false, None).await; + assert!( + !deleted_plan.contains("limit="), + "deletions must suppress the pushdown, plan was:\n{deleted_plan}" + ); + } + + #[rstest] + #[tokio::test] + async fn test_limit_not_pushed_with_fragment_subset( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + ) { + // The scalar-index search runs over the whole dataset. A `with_fragments` subset is + // applied only afterwards. If the limit were pushed, an unordered scan restricted to a + // fragment whose ids sort last would early-stop on matches in the other fragment and + // return too few (here zero) rows. The limit must therefore not be pushed for a subset. + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + // Two fragments with ascending ids (index order == storage order): frag 0 holds the + // smallest ids, so the first matches in index order all live outside the second fragment. + let num_rows = 20_000i32; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..num_rows))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let write_params = WriteParams { + data_storage_version: Some(data_storage_version), + max_rows_per_file: 10_000, + ..Default::default() + }; + let mut dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + let fragments = dataset.fragments().as_ref().clone(); + assert_eq!(fragments.len(), 2, "expected two fragments"); + // Restrict to the second fragment (the largest ids). + let second_fragment = fragments[1].clone(); + + let limit = 100i64; + let mut scan = dataset.scan(); + scan.with_fragments(vec![second_fragment]) + .filter("id >= 0") + .unwrap() + .scan_in_order(false) + .limit(Some(limit), None) + .unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec(); + + assert_eq!( + ids.len(), + limit as usize, + "a fragment-subset scan must still return `limit` rows" + ); + assert!( + ids.iter().all(|&id| id >= 10_000), + "must only return rows from the requested fragment" + ); + } + + #[rstest] + #[tokio::test] + async fn test_limit_not_pushed_with_retired_fragment( + #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] + data_storage_version: LanceFileVersion, + ) { + // A scalar index keeps entries for every fragment it was trained on. When a fragment is + // retired (here by deleting all of its rows) the index still has stale entries for it, but + // those rows are dropped after the search. If the limit were pushed, an unordered scan + // could early-stop on the retired fragment's (smallest) ids and return fewer than `limit` + // live rows. The limit must therefore not be pushed when the index covers a retired + // fragment, even though no live fragment has a deletion file. + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + // Fragment 0 holds the smallest ids (0..10_000), fragment 1 the rest. Index order == + // storage order, so the first matches in index order all live in the soon-retired fragment. + let num_rows = 20_000i32; + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..num_rows))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let write_params = WriteParams { + data_storage_version: Some(data_storage_version), + max_rows_per_file: 10_000, + ..Default::default() + }; + let mut dataset = Dataset::write(reader, "memory://", Some(write_params)) + .await + .unwrap(); + assert_eq!(dataset.fragments().len(), 2, "expected two fragments"); + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + // Delete every row in fragment 0. Depending on the storage version this either removes the + // fragment from the manifest entirely (Stable) or leaves it behind with a deletion file + // (Legacy). Either way the single index segment was trained on both fragments and still has + // stale entries for fragment 0 that are dropped after the search. + dataset.delete("id < 10000").await.unwrap(); + // The set of fragments whose rows all survive the search: live fragments with no deletion + // file. This is exactly the set `index_search_limit` requires the index coverage to be a + // subset of. + let live_undeleted: std::collections::HashSet = dataset + .fragments() + .iter() + .filter(|fragment| fragment.deletion_file.is_none()) + .map(|fragment| fragment.id as u32) + .collect(); + assert_eq!( + live_undeleted.len(), + 1, + "only the second fragment should have surviving rows, got {live_undeleted:?}" + ); + // Precondition for what this test exercises: the index covers a fragment whose rows do not + // all survive (a retired fragment, or one masked by a deletion file), so an early stop could + // spend its budget on rows that get dropped afterwards. + let covered = scalar_index_fragment_bitmap(&dataset, "id", "id_idx") + .await + .unwrap() + .expect("index must report fragment coverage"); + assert!( + covered + .iter() + .any(|frag_id| !live_undeleted.contains(&frag_id)), + "test precondition: index must cover a retired/deleted fragment, covered={covered:?}, live_undeleted={live_undeleted:?}" + ); + + let limit = 100i64; + let mut scan = dataset.scan(); + scan.filter("id >= 0") + .unwrap() + .scan_in_order(false) + .limit(Some(limit), None) + .unwrap(); + let batch = scan.try_into_batch().await.unwrap(); + let ids = batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec(); + + assert_eq!( + ids.len(), + limit as usize, + "a scan over a stale-index dataset must still return `limit` live rows" + ); + assert!( + ids.iter().all(|&id| id >= 10_000), + "retired rows must not be returned" + ); + } + + #[tokio::test] + async fn test_limit_not_pushed_for_compound_index_query() { + // A limit only stops a single top-level index lookup early, since `evaluate_nullable` + // drops it for And/Or/Not. So `index_search_limit` must refuse to push a limit for a + // compound (here `Or`) index query even when every covered fragment is scanned and + // undeleted, while a lone `Query` on the same data must still push. This pins the gate + // decision directly, since a compound query returns identical rows either way and no + // black-box scan could tell the two apart. + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, false), + ])); + let num_rows = 1_000i32; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..num_rows)), + Arc::new(Int32Array::from_iter_values(0..num_rows)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + for col in ["a", "b"] { + dataset + .create_index( + &[col], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + } + let scanned = dataset.fragments().as_ref().clone(); + + // A compound `Or` over two indexed columns, both covering the single clean fragment. + let mut or_scan = dataset.scan(); + or_scan + .filter("a >= 500 OR b >= 500") + .unwrap() + .scan_in_order(false) + .limit(Some(100), None) + .unwrap(); + let or_plan = or_scan.create_filter_plan(true, None, None).await.unwrap(); + assert!( + matches!( + or_plan.expr_filter_plan.index_query, + Some(ScalarIndexExpr::Or(..)) + ), + "test precondition: the filter must plan to an Or index query, got {:?}", + or_plan.expr_filter_plan.index_query + ); + let pushed = or_scan + .index_search_limit(&or_plan.expr_filter_plan, &scanned) + .await + .unwrap(); + assert_eq!( + pushed, None, + "a compound Or index query must not push a limit" + ); + + // A lone `Query` on the same data still pushes, so the guard is not over-broad. + let mut single_scan = dataset.scan(); + single_scan + .filter("a >= 500") + .unwrap() + .scan_in_order(false) + .limit(Some(100), None) + .unwrap(); + let single_plan = single_scan + .create_filter_plan(true, None, None) + .await + .unwrap(); + assert!( + matches!( + single_plan.expr_filter_plan.index_query, + Some(ScalarIndexExpr::Query(_)) + ), + "test precondition: the filter must plan to a single Query index query, got {:?}", + single_plan.expr_filter_plan.index_query + ); + let pushed = single_scan + .index_search_limit(&single_plan.expr_filter_plan, &scanned) + .await + .unwrap(); + assert_eq!( + pushed, + Some(100), + "a lone Query index lookup must still push the limit" + ); + } + + #[tokio::test] + async fn test_limit_pushdown_tolerates_unknown_index_coverage() { + // When an index reports no fragment coverage (a legacy or unsupported index with no + // fragment bitmap, or no segments for the name), `scalar_index_fragment_bitmap` returns + // `None`. `index_search_limit` must treat that as "do not push" and return `Ok(None)`, not + // propagate an error: the pushdown is only an optimization and must never turn an otherwise + // valid scan into a failure. A missing index name reproduces the same `None` coverage. + use lance_index::scalar::SargableQuery; + use lance_index::scalar::expression::ScalarIndexSearch; + + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..1_000))], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + let scanned = dataset.fragments().as_ref().clone(); + + // A single `Query` leaf that names an index with no committed segments, so its coverage + // resolves to `None`. `fragment_bitmap` is `None` for the same reason a legacy index's is. + let plan = ExprFilterPlan { + index_query: Some(ScalarIndexExpr::Query(ScalarIndexSearch { + column: "id".to_string(), + index_name: "does_not_exist_idx".to_string(), + index_type: "BTree".to_string(), + query: Arc::new(SargableQuery::Range( + std::ops::Bound::Unbounded, + std::ops::Bound::Unbounded, + )), + needs_recheck: false, + fragment_bitmap: None, + })), + skip_recheck: true, + refine_expr: None, + full_expr: None, + }; + + let mut scan = dataset.scan(); + scan.scan_in_order(false).limit(Some(100), None).unwrap(); + let pushed = scan.index_search_limit(&plan, &scanned).await.unwrap(); + assert_eq!( + pushed, None, + "unknown index coverage must disable the pushdown, not error" + ); + } + #[test_log::test(tokio::test)] async fn test_limit_cancel() { // If there is a filter and a limit and we can't use the index to satisfy diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index a783ca09905..a828dd6ba47 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -252,6 +252,116 @@ async fn test_overlay_stale_drop_and_new_match(#[values(false, true)] stable_row assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); } +/// A limited, unordered, index-backed scan must return a full page of live rows even when a +/// data overlay is masking index hits. +/// +/// The overlay block removes index-matched rows *after* the index search, which is the same +/// shape as the deletion-file case that `index_search_limit` already refuses to push through: +/// its stated invariant is that every index-covered row must survive into the result, or an +/// early stop can leave fewer than `limit` live rows. The B-tree early stop is page-granular, +/// so a heavily masked page is the worst case. +/// +/// Note on scope: this asserts the observable contract, not the guard. Instrumenting the +/// planner shows that with the guard removed the limit *is* pushed while the block is active +/// (`pushdown_limit=Some(200)`), yet the result is still correct, because the stale-row union +/// and recheck paths currently backfill. No input was found that actually returns short. The +/// guard in `new_filtered_read` / `scalar_indexed_scan` is therefore upholding the documented +/// invariant rather than fixing a demonstrated defect, and this test pins the behavior that +/// must hold either way. +/// +/// Setup: one fragment of 8192 rows (two `DEFAULT_BTREE_BATCH_SIZE` pages) where `age` == `id`, +/// so index order matches storage order. Overlays push the first 4000 rows above the +/// predicate, leaving only 96 live matches in page 0, fewer than the limit of 200. +#[tokio::test] +async fn test_overlay_block_does_not_short_limited_index_scan() { + const PAGE: i32 = 4096; + const NUM_ROWS: i32 = 2 * PAGE; + const STALE: i32 = 4000; + const LIMIT: usize = 200; + const THRESHOLD: i32 = 1_000_000; + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, true), + ArrowField::new("age", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..NUM_ROWS)), + Arc::new(Int32Array::from_iter_values(0..NUM_ROWS)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); + let mut dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: NUM_ROWS as usize, + ..Default::default() + }), + ) + .await + .unwrap(); + build_age_index(&mut dataset).await; + + // Push the first STALE rows, the lowest index-order matches, above the predicate so + // they stop matching. Page 0 keeps only PAGE - STALE = 96 live matches. + let dataset = commit_overlay( + dataset, + "age_limit_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter(0..STALE as u32)), + vec![Arc::new(Int32Array::from_iter_values( + (0..STALE).map(|_| THRESHOLD + 1), + )) as ArrayRef], + ) + .await; + + let filter = format!("age < {THRESHOLD}"); + + // Ground truth: every non-overlaid row still matches. + let mut unlimited = dataset.scan(); + unlimited + .filter(&filter) + .unwrap() + .project(&["id"]) + .unwrap() + .scan_in_order(false); + let total = unlimited.try_into_batch().await.unwrap().num_rows(); + assert_eq!( + total, + (NUM_ROWS - STALE) as usize, + "overlaid rows must drop out of the unlimited result" + ); + assert!( + (PAGE - STALE) < LIMIT as i32, + "the test only bites when page 0's live matches fall below the limit" + ); + + // The limited scan must still return a full page of live rows. With the limit pushed + // into the index while the overlay block is active, the index stops after page 0, 4000 + // of whose hits are masked, and this returns 96 instead of 200. + let mut scanner = dataset.scan(); + scanner + .filter(&filter) + .unwrap() + .project(&["id"]) + .unwrap() + // The limit is only pushed into the index for an unordered scan, which is exactly + // the configuration this guard has to hold for. + .scan_in_order(false) + .limit(Some(LIMIT as i64), None) + .unwrap(); + let batch = scanner.try_into_batch().await.unwrap(); + assert_eq!( + batch.num_rows(), + LIMIT, + "an overlay block must not shorten a limited index scan" + ); +} + /// Row-level BTree precision: when one row in a covered fragment is stale, only that row is /// blocked from the index result and re-evaluated on the stale-Take path. Non-stale rows in /// the same fragment (including one that matches the predicate) remain on the indexed path. diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 4a22aac62e7..3c3ac6c1add 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -154,6 +154,16 @@ impl ScalarIndex for LogicalScalarIndex { options: SearchOptions, metrics: &dyn MetricsCollector, ) -> Result { + // Any limit in `options` is forwarded to every segment unchanged. That is safe + // because each segment independently returns at least `limit` matches, so their + // combination does too, and `combine_search_results` keeps the result `AtLeast`. + // + // A shared budget, decremented by matches already confirmed so later segments read + // less, would cut total reads. It is not used here because it forces the segments + // to be searched one after another, and the parallel fan out below is what keeps + // latency flat as segment count grows. Trading that for fewer reads is a real + // tradeoff rather than a clear win, so it is left to a follow up that can measure + // both sides on a many segment index. let results = try_join_all( self.segments .iter() @@ -585,6 +595,76 @@ mod tests { ); } + #[tokio::test] + async fn test_btree_segment_search_limited_across_segments() { + // A limited search forwards the limit to every segment and combines the results, so the + // combined result must still hold at least `limit` matches across the segments. + let test_dir = TempStrDir::default(); + let dataset = lance_datagen::gen_batch() + .col("value", array::step::()) + .into_dataset( + test_dir.as_str(), + FragmentCount::from(4), + FragmentRowCount::from(16), + ) + .await + .unwrap(); + let mut dataset = dataset; + let fragments = dataset.get_fragments(); + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::BTree); + let mut segments = Vec::new(); + for fragment in &fragments { + segments.push( + CreateIndexBuilder::new(&mut dataset, &["value"], IndexType::BTree, ¶ms) + .name("value_btree_limited".to_string()) + .fragments(vec![fragment.id() as u32]) + .execute_uncommitted() + .await + .unwrap(), + ); + } + dataset + .commit_existing_index_segments("value_btree_limited", "value", segments) + .await + .unwrap(); + + let logical = open_named_scalar_index( + &dataset, + "value", + "value_btree_limited", + &NoOpMetricsCollector, + ) + .await + .unwrap(); + + // All 64 rows match the unbounded range. With a limit the combined result across the + // four segments must still satisfy at least `limit` matches. + let query = SargableQuery::Range(Bound::Unbounded, Bound::Unbounded); + let limit = 10usize; + let result = logical + .search_with_options( + &query, + SearchOptions::default() + .with_track_nulls(false) + .with_limit(Some(limit)), + &NoOpMetricsCollector, + ) + .await + .unwrap(); + // Each B-tree segment short-circuits on the limit and reports `AtLeast`, so the combined + // result must also be `AtLeast` (a lower bound), never `Exact`. Otherwise a caller could + // treat this partial match set as the complete answer. + let row_addrs = match result { + SearchResult::AtLeast(row_addrs) => row_addrs, + other => panic!("limited search must return AtLeast, got {:?}", other), + }; + let count = row_addrs.true_rows().row_addrs().unwrap().count(); + assert!( + count >= limit, + "limited search must return at least {limit} matches, got {count}" + ); + } + #[tokio::test] async fn test_bitmap_segments_commit_and_query_as_logical_index() { let test_dir = TempStrDir::default(); diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index b1d8ed2855e..96285f3b2aa 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -120,16 +120,32 @@ pub struct ScalarIndexExec { properties: Arc, metrics: ExecutionPlanMetricsSet, result_format: IndexExprResultWireFormat, + /// Hint passed to the index search so it can stop once it has found this many + /// matches. `None` means search all matches. + /// + /// This is only an optimization. A downstream `GlobalLimitExec` still applies the + /// exact limit, so the index only needs to return at least this many rows. + limit: Option, } impl DisplayAs for ScalarIndexExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "ScalarIndexQuery: query={}", self.expr) + write!(f, "ScalarIndexQuery: query={}", self.expr)?; + // Surfacing the pushed limit keeps the early-stop visible in EXPLAIN, which + // is otherwise invisible to both users and tests. + if let Some(limit) = self.limit { + write!(f, ", limit={}", limit)?; + } + Ok(()) } DisplayFormatType::TreeRender => { - write!(f, "ScalarIndexQuery\nquery={}", self.expr) + write!(f, "ScalarIndexQuery\nquery={}", self.expr)?; + if let Some(limit) = self.limit { + write!(f, "\nlimit={}", limit)?; + } + Ok(()) } } } @@ -153,9 +169,20 @@ impl ScalarIndexExec { properties, metrics: ExecutionPlanMetricsSet::new(), result_format, + limit: None, } } + /// Push a `limit` hint into the index search so it can stop early. + /// + /// Only set this when returning any `limit` matching rows is safe, such as an + /// unordered scan whose results are not filtered further. Correctness still relies on + /// a downstream limit operator. + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + pub fn dataset(&self) -> &Arc { &self.dataset } @@ -206,12 +233,14 @@ impl ScalarIndexExec { dataset: Arc, plan_metrics: ExecutionPlanMetricsSet, result_format: IndexExprResultWireFormat, + limit: Option, ) -> Result { let metrics = IndexMetrics::new(&plan_metrics, 0); let query_result = { let search_time = plan_metrics.new_time(SCALAR_INDEX_SEARCH_TIME_METRIC, 0); let _timer = search_time.timer(); - expr.evaluate(dataset.as_ref(), &metrics).await? + expr.evaluate_limited(dataset.as_ref(), &metrics, limit) + .await? }; let fragments_covered_by_result = Self::fragments_covered_by_index_query(&expr, dataset.as_ref()).await?; @@ -259,6 +288,7 @@ impl ExecutionPlan for ScalarIndexExec { self.dataset.clone(), self.metrics.clone(), self.result_format, + self.limit, ); let stream = futures::stream::iter(vec![batch_fut]) .then(|batch_fut| batch_fut.map_err(|err| err.into())) @@ -701,16 +731,30 @@ pub struct MaterializeIndexExec { overlay_block: Option, properties: Arc, metrics: ExecutionPlanMetricsSet, + /// Hint passed to the index search so it can stop once it has found this many + /// matches. `None` means materialize all matches. + /// + /// This is only an optimization. A downstream `GlobalLimitExec` still applies the + /// exact limit, so the index only needs to return at least this many rows. + limit: Option, } impl DisplayAs for MaterializeIndexExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "MaterializeIndex: query={}", self.expr) + write!(f, "MaterializeIndex: query={}", self.expr)?; + if let Some(limit) = self.limit { + write!(f, ", limit={}", limit)?; + } + Ok(()) } DisplayFormatType::TreeRender => { - write!(f, "MaterializeIndex\nquery={}", self.expr) + write!(f, "MaterializeIndex\nquery={}", self.expr)?; + if let Some(limit) = self.limit { + write!(f, "\nlimit={}", limit)?; + } + Ok(()) } } } @@ -755,6 +799,33 @@ impl Iterator for FragIdIter<'_> { } } +/// Pick the row-address mask `MaterializeIndexExec` should materialize. +/// +/// The `upper` mask is the candidate set: for `Exact` it is the answer, and for `AtMost` +/// or refined results it is a superset that `LanceFilterExec` prunes downstream, because +/// the full filter reruns on the materialized batches. +/// +/// `AtLeast` carries an unbounded `upper`, so it cannot serve as a candidate set. +/// Materializing its `lower` mask instead is sound ONLY when this node asked for a partial +/// answer, which is exactly the pushed-limit case: the B-tree stops early, every row in +/// `lower` is a confirmed match, and a downstream `GlobalLimitExec` still enforces the +/// exact limit. +/// +/// A lower bound can also arrive for reasons unrelated to any limit. `bloomfilter.rs` and +/// the multi-segment combination in `scalar_logical.rs` both produce `AtLeast` on their +/// own, and there the unconfirmed rows between `lower` and `upper` still need the full +/// recheck. Treating `lower` as the answer would silently drop them, so with no pushed +/// limit this keeps the preexisting behavior rather than guessing. +fn candidate_mask_for(result: IndexExprResult, limit: Option) -> Result { + if result.is_at_least() && !result.is_exact() { + if limit.is_none() { + todo!("Support AtLeast in MaterializeIndexExec") + } + return Ok(result.lower); + } + Ok(result.upper) +} + impl MaterializeIndexExec { pub fn new( dataset: Arc, @@ -774,9 +845,21 @@ impl MaterializeIndexExec { overlay_block: None, properties, metrics: ExecutionPlanMetricsSet::new(), + limit: None, } } + /// Push a `limit` hint into the index search so it can stop early. + /// + /// Only set this when returning any `limit` matching rows is safe, such as an + /// unordered scan whose results are not filtered further and whose rows are not + /// dropped after the search by deletions or an overlay block. Correctness still + /// relies on a downstream limit operator. + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + /// Block specific row addresses (see the `overlay_block` field) from the index result. pub fn with_overlay_block(mut self, block: RowAddrMask) -> Self { self.overlay_block = Some(block); @@ -790,8 +873,9 @@ impl MaterializeIndexExec { fragments: Arc>, overlay_block: Option, metrics: Arc, + limit: Option, ) -> Result { - let expr_result = expr.evaluate(dataset.as_ref(), metrics.as_ref()); + let expr_result = expr.evaluate_limited(dataset.as_ref(), metrics.as_ref(), limit); let span = debug_span!("create_prefilter"); let prefilter = span.in_scope(|| { let fragment_bitmap = @@ -808,19 +892,14 @@ impl MaterializeIndexExec { // gets pruned downstream by `LanceFilterExec` (the full filter // runs on the materialized batches via the scan plan, so any // non-matching candidates in `upper` are dropped before they - // reach the user). `AtLeast` carries an unbounded upper, so the - // candidate set is the whole row space — not actionable here. - let take_upper = |result: IndexExprResult| -> Result { - if result.is_at_least() && !result.is_exact() { - todo!("Support AtLeast in MaterializeIndexExec") - } - Ok(result.upper) - }; + // reach the user). + let candidate_mask = + |result: IndexExprResult| -> Result { candidate_mask_for(result, limit) }; let mut mask = if let Some(prefilter) = prefilter { let (expr_result, prefilter) = futures::try_join!(expr_result, prefilter)?; - take_upper(expr_result)? & (*prefilter).clone() + candidate_mask(expr_result)? & (*prefilter).clone() } else { - take_upper(expr_result.await?)? + candidate_mask(expr_result.await?)? }; if let Some(block) = overlay_block { mask = mask & block; @@ -958,6 +1037,7 @@ impl ExecutionPlan for MaterializeIndexExec { self.fragments.clone(), self.overlay_block.clone(), metrics, + self.limit, ); let stream = futures::stream::iter(vec![batch_fut]) .then(|batch_fut| batch_fut.map_err(|err| err.into())) @@ -993,6 +1073,47 @@ impl ExecutionPlan for MaterializeIndexExec { mod tests { use std::{ops::Bound, sync::Arc}; + use super::{IndexExprResult, RowAddrMask, candidate_mask_for}; + + fn mask_of(addrs: [u64; 2]) -> RowAddrMask { + RowAddrMask::AllowList(RowAddrTreeMap::from_iter(addrs)) + } + + /// A pushed limit is the only thing that makes a lower bound a sufficient answer, so + /// `AtLeast` may collapse to its `lower` mask only when this node asked for one. + #[test] + fn test_candidate_mask_uses_lower_only_for_a_pushed_limit() { + let result = IndexExprResult::at_least(mask_of([1, 2])); + let expected = result.lower.clone(); + let picked = candidate_mask_for(result, Some(10)).unwrap(); + assert_eq!( + picked, expected, + "a limited AtLeast must materialize its confirmed lower bound" + ); + } + + /// Anything that is not a bare lower bound keeps using `upper` as the candidate set, + /// which `LanceFilterExec` prunes downstream. + #[test] + fn test_candidate_mask_uses_upper_for_exact_results() { + let result = IndexExprResult::exact(mask_of([3, 4])); + let expected = result.upper.clone(); + assert_eq!(candidate_mask_for(result, None).unwrap(), expected); + let result = IndexExprResult::exact(mask_of([3, 4])); + let expected = result.upper.clone(); + assert_eq!(candidate_mask_for(result, Some(10)).unwrap(), expected); + } + + /// Regression: `AtLeast` also arrives from sources that have nothing to do with a + /// limit, such as `bloomfilter.rs` and the multi-segment combination in + /// `scalar_logical.rs`. Collapsing those to `lower` would silently drop the + /// unconfirmed rows that still require a recheck, so this path must not do it. + #[test] + #[should_panic(expected = "Support AtLeast in MaterializeIndexExec")] + fn test_candidate_mask_refuses_lower_without_a_pushed_limit() { + let _ = candidate_mask_for(IndexExprResult::at_least(mask_of([1, 2])), None); + } + use crate::index::DatasetIndexExt; use arrow::datatypes::UInt64Type; use arrow::record_batch::RecordBatchIterator;