From 91cbb82afabbd0bc671d270d2da303d41378066d Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Tue, 2 Jun 2026 21:44:57 +0200 Subject: [PATCH 01/14] feat(scalar-index): push scan limit into index search for early termination --- rust/lance-index/src/scalar.rs | 21 +++ rust/lance-index/src/scalar/btree.rs | 144 ++++++++++++++++++-- rust/lance-index/src/scalar/expression.rs | 34 ++++- rust/lance/src/dataset/scanner.rs | 154 ++++++++++++++++++++-- rust/lance/src/index/scalar_logical.rs | 18 +++ rust/lance/src/io/exec/scalar_index.rs | 43 +++++- 6 files changed, 382 insertions(+), 32 deletions(-) diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 5ab138ff481..3100a6892d8 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -983,6 +983,27 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { metrics: &dyn MetricsCollector, ) -> Result; + /// Like [`Self::search`] but with a best-effort `limit` hint: when `limit` is `Some(n)` + /// an index may stop after finding `n` matches (it may still return more). Only push a + /// limit for a single positive lookup. The default ignores it and calls [`Self::search`]. + /// + /// ``` + /// # use lance_core::Result; + /// # use lance_index::{metrics::NoOpMetricsCollector, scalar::{AnyQuery, ScalarIndex}}; + /// # async fn example(index: &dyn ScalarIndex, query: &dyn AnyQuery) -> Result<()> { + /// let _result = index.search_limited(query, &NoOpMetricsCollector, Some(10)).await?; + /// # Ok(()) + /// # } + /// ``` + async fn search_limited( + &self, + query: &dyn AnyQuery, + metrics: &dyn MetricsCollector, + _limit: Option, + ) -> Result { + self.search(query, metrics).await + } + /// Returns true if the remap operation is supported fn can_remap(&self) -> bool; diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 6de490b9572..024cd515800 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -1696,14 +1696,19 @@ impl Index for BTreeIndex { } } -#[async_trait] -impl ScalarIndex for BTreeIndex { - async fn search( +impl BTreeIndex { + /// Shared implementation for [`ScalarIndex::search`] and + /// [`ScalarIndex::search_limited`]. + /// + /// When `limit` is `Some(n)` the pages are searched in order and the search stops once + /// it has at least `n` matching row ids. The result may hold more than `n` rows but + /// never fewer unless the query matches fewer. + async fn do_search( &self, - query: &dyn AnyQuery, + query: &SargableQuery, metrics: &dyn MetricsCollector, + limit: Option, ) -> Result { - let query = query.as_any().downcast_ref::().unwrap(); let mut pages = match query { SargableQuery::Equals(val) => self .page_lookup @@ -1764,7 +1769,10 @@ impl ScalarIndex for BTreeIndex { // We add them as Matches::Some (not Matches::All) so that // FlatIndex::search() evaluates the predicate and correctly marks // the rows as NULL rather than TRUE. - if !matches!(query, SargableQuery::IsNull()) { + // + // When a `limit` is set the query is a single positive lookup, so null tracking + // is not needed and skipping null pages helps us stop early. + if limit.is_none() && !matches!(query, SargableQuery::IsNull()) { let existing: HashSet = pages.iter().map(|m| m.page_id()).collect(); for &page_id in self .page_lookup @@ -1789,19 +1797,60 @@ 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?; + // Collect row IDs from the pages. `buffered` keeps page order. When a `limit` is + // set we read one page at a time and stop once we have enough matches, so we do + // not issue I/O for pages we never need. Without a limit we fan out across CPUs + // (I/O and compute are mixed, but the important case is the index being cached). + 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 { + // Count only TRUE matches. NULL rows never match, so they must not count + // toward the limit. `len()` already excludes nulls. + matches_found += page_result.len().unwrap_or(0); + results.push(page_result); + if matches_found >= limit as u64 { + break; + } + } else { + results.push(page_result); + } + } // Merge matching row IDs let selection = NullableRowAddrSet::union_all(&results); Ok(SearchResult::Exact(selection)) } +} + +#[async_trait] +impl ScalarIndex for BTreeIndex { + async fn search( + &self, + query: &dyn AnyQuery, + metrics: &dyn MetricsCollector, + ) -> Result { + let query = query.as_any().downcast_ref::().unwrap(); + self.do_search(query, metrics, None).await + } + + async fn search_limited( + &self, + query: &dyn AnyQuery, + metrics: &dyn MetricsCollector, + limit: Option, + ) -> Result { + let query = query.as_any().downcast_ref::().unwrap(); + self.do_search(query, metrics, limit).await + } fn can_remap(&self) -> bool { true @@ -4990,6 +5039,75 @@ mod tests { } } + /// `search_limited` returns at least `limit` matches but stops reading pages early, + /// so for a multi-page range it returns fewer rows than an unlimited search. + #[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()), + )); + + // Enough rows to span several btree pages, with no nulls so every row matches an + // unbounded range. `train_btree_index` makes pages of `DEFAULT_BTREE_BATCH_SIZE` + // rows, so this gives five pages. + 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 schema = data.schema(); + let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + 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. + let full = index.search(&everything, &metrics).await.unwrap(); + let full_len = full.row_addrs().len().unwrap(); + assert_eq!(full_len, num_rows); + + // A limit that reaches into the second page. The search must satisfy it but stop + // well before reading all five pages. + let limit = (DEFAULT_BTREE_BATCH_SIZE + 100) as usize; + let limited = index + .search_limited(&everything, &metrics, Some(limit)) + .await + .unwrap(); + 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" + ); + } + fn sample_lookup_batch() -> RecordBatch { record_batch!( ("min", Int32, [Some(0), Some(10), Some(20)]), diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index 187d5be999f..51e826787ca 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1309,21 +1309,24 @@ impl ScalarIndexExpr { &self, index_loader: &dyn ScalarIndexLoader, metrics: &dyn MetricsCollector, + limit: Option, ) -> Result { match self { + // A limit only applies to a single positive lookup. NOT, AND, and OR need the + // full result of each side, so the limit is dropped when recursing into them. Self::Not(inner) => { - let result = inner.evaluate_nullable(index_loader, metrics).await?; + let result = inner.evaluate_nullable(index_loader, metrics, None).await?; Ok(!result) } Self::And(lhs, rhs) => { - let lhs_result = lhs.evaluate_nullable(index_loader, metrics); - let rhs_result = rhs.evaluate_nullable(index_loader, metrics); + let lhs_result = lhs.evaluate_nullable(index_loader, metrics, None); + let rhs_result = rhs.evaluate_nullable(index_loader, metrics, None); 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_nullable(index_loader, metrics); - let rhs_result = rhs.evaluate_nullable(index_loader, metrics); + let lhs_result = lhs.evaluate_nullable(index_loader, metrics, None); + let rhs_result = rhs.evaluate_nullable(index_loader, metrics, None); let (lhs_result, rhs_result) = try_join!(lhs_result, rhs_result)?; Ok(lhs_result | rhs_result) } @@ -1331,7 +1334,9 @@ impl ScalarIndexExpr { let index = index_loader .load_index(&search.column, &search.index_name, metrics) .await?; - let search_result = index.search(search.query.as_ref(), metrics).await?; + let search_result = index + .search_limited(search.query.as_ref(), metrics, limit) + .await?; Ok(search_result.into()) } } @@ -1342,9 +1347,24 @@ impl ScalarIndexExpr { &self, index_loader: &dyn ScalarIndexLoader, metrics: &dyn MetricsCollector, + ) -> Result { + self.evaluate_limited(index_loader, metrics, None).await + } + + /// Like [`Self::evaluate`] but pushes a `limit` hint into the index search so it can + /// stop once it has found at least `limit` matches. + /// + /// See [`crate::scalar::ScalarIndex::search_limited`] for the rules on when a limit + /// may be pushed down. + #[instrument(level = "debug", skip_all)] + pub async fn evaluate_limited( + &self, + index_loader: &dyn ScalarIndexLoader, + metrics: &dyn MetricsCollector, + limit: Option, ) -> Result { Ok(self - .evaluate_nullable(index_loader, metrics) + .evaluate_nullable(index_loader, metrics, limit) .await? .drop_nulls()) } diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 9a5cd94dd09..f635e420c20 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2821,6 +2821,16 @@ impl Scanner { fragments: Option>>, scan_range: Option>, ) -> Result> { + // Decide whether a limit can be pushed into the index search. The fragments the + // read covers (used for the deletion check) are the requested subset, or the whole + // dataset when none was given. + 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 pushdown_limit = self.index_search_limit(filter_plan, scanned_fragments); + let mut read_options = FilteredReadOptions::basic_full_read(&self.dataset) .with_filter_plan(filter_plan.clone()) .with_projection(projection); @@ -2859,11 +2869,10 @@ 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 + Arc::new( + ScalarIndexExec::new(self.dataset.clone(), index_query, result_format) + .with_limit(pushdown_limit), + ) as Arc }); Ok(Arc::new(FilteredReadExec::try_new( @@ -4042,6 +4051,55 @@ impl Scanner { Ok((relevant_frags, missing_frags)) } + /// 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 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 relevant fragments have no deletions. Deleted rows are pruned after the index + /// search, so stopping early could leave fewer than `limit` live rows. + /// + /// Returns `None` when no limit can be pushed. + fn index_search_limit( + &self, + filter_plan: &ExprFilterPlan, + relevant_fragments: &[Fragment], + ) -> Option { + let limit = self.limit?; + if limit <= 0 { + return None; + } + if self.ordering.is_some() + || self.nearest.is_some() + || self.full_text_query.is_some() + || self.aggregate.is_some() + || filter_plan.has_refine() + { + return None; + } + if filter_plan + .index_query + .as_ref() + .is_some_and(|query| query.needs_recheck()) + { + return None; + } + if relevant_fragments + .iter() + .any(|fragment| fragment.deletion_file.is_some()) + { + return None; + } + let offset = self.offset.unwrap_or(0).max(0) as usize; + 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( @@ -4066,11 +4124,18 @@ impl Scanner { .partition_frags_by_coverage(index_expr, fragments) .await?; - let mut plan: Arc = Arc::new(MaterializeIndexExec::new( - self.dataset.clone(), - index_expr.clone(), - Arc::new(relevant_frags), - )); + // A limit can be pushed into the index search, but only when its rows are used as + // is and the relevant fragments have no deletions. + let pushdown_limit = self.index_search_limit(filter_plan, &relevant_frags); + + let mut plan: Arc = Arc::new( + MaterializeIndexExec::new( + self.dataset.clone(), + index_expr.clone(), + Arc::new(relevant_frags), + ) + .with_limit(pushdown_limit), + ); let refine_expr = filter_plan.refine_expr.as_ref(); @@ -5768,6 +5833,75 @@ mod test { assert_eq!(ids, &(10..20).collect::>()); } + #[tokio::test] + async fn test_limit_pushed_into_scalar_index() { + // When a scan filter is fully served by a scalar index (no refine, no recheck, no + // ordering) the limit can be pushed into the index search. The result must still + // be exactly `limit` rows that all match the filter. Early stop must not drop or + // duplicate matches. + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "id", + DataType::Int32, + false, + )])); + // Span several btree pages so a small limit short-circuits before the end. + let num_rows = 20_000; + 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 mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + dataset + .create_index( + &["id"], + IndexType::BTree, + None, + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + let limit = 100; + let scan_ids = |dataset: Arc| async move { + let batch = dataset + .scan() + .filter("id >= 5") + .unwrap() + .limit(Some(limit), None) + .unwrap() + .try_into_batch() + .await + .unwrap(); + batch + .column_by_name("id") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }; + + let ids = scan_ids(Arc::new(dataset.clone())).await; + assert_eq!(ids.len(), limit as usize); + assert!( + ids.iter().all(|&id| id >= 5), + "every returned row must satisfy the filter" + ); + + // With deletions present the limit must not be pushed, since deleted rows are + // pruned after the index search. The scan must still return exactly `limit` live + // matches. + dataset.delete("id >= 5 AND id < 10000").await.unwrap(); + let ids = scan_ids(Arc::new(dataset)).await; + assert_eq!(ids.len(), limit as usize); + assert!( + ids.iter().all(|&id| id >= 10000), + "deleted rows must not be returned" + ); + } + #[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/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 162a36a0c97..8e49dd03af6 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -132,6 +132,24 @@ impl ScalarIndex for LogicalScalarIndex { combine_search_results(results) } + async fn search_limited( + &self, + query: &dyn AnyQuery, + metrics: &dyn MetricsCollector, + limit: Option, + ) -> Result { + // Forwarding the limit to every segment is safe. Each segment returns at least + // `limit` matches when it has them, so the combined result still has at least + // `limit` matches overall. + let results = try_join_all( + self.segments + .iter() + .map(|segment| segment.search_limited(query, metrics, limit)), + ) + .await?; + combine_search_results(results) + } + fn can_remap(&self) -> bool { false } diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index ade4995fb4b..0eb35cc9eea 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -76,6 +76,12 @@ 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 { @@ -109,9 +115,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 } @@ -161,12 +178,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?; @@ -218,6 +237,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())) @@ -485,6 +505,12 @@ pub struct MaterializeIndexExec { fragments: Arc>, 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 { @@ -557,17 +583,29 @@ impl MaterializeIndexExec { fragments, 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. Correctness still relies on + /// a downstream limit operator. + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + #[instrument(name = "materialize_scalar_index", skip_all, level = "debug")] async fn do_execute( expr: ScalarIndexExpr, dataset: Arc, fragments: Arc>, 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 = @@ -734,6 +772,7 @@ impl ExecutionPlan for MaterializeIndexExec { self.dataset.clone(), self.fragments.clone(), metrics, + self.limit, ); let stream = futures::stream::iter(vec![batch_fut]) .then(|batch_fut| batch_fut.map_err(|err| err.into())) From 8de2c7b020d880fc021c940d2f4aa234720afbc7 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 20 Jun 2026 14:14:35 +0200 Subject: [PATCH 02/14] fix(scanner): only push scalar-index limit for unordered scans --- rust/lance-index/src/scalar/btree.rs | 19 +++----- rust/lance-index/src/scalar/expression.rs | 3 +- rust/lance/src/dataset/scanner.rs | 55 +++++++++++++---------- rust/lance/src/index/scalar_logical.rs | 4 +- 4 files changed, 39 insertions(+), 42 deletions(-) diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 45e691da3fb..19141ad4033 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -1802,9 +1802,7 @@ impl BTreeIndex { // We add them as Matches::Some (not Matches::All) so that // FlatIndex::search() evaluates the predicate and correctly marks // the rows as NULL rather than TRUE. - // - // When a `limit` is set the query is a single positive lookup, so null tracking - // is not needed and skipping null pages helps us stop early. + // A `limit` implies a single positive lookup, so skip null tracking to stop early. if limit.is_none() && !matches!(query, SargableQuery::IsNull()) { let existing: HashSet = pages.iter().map(|m| m.page_id()).collect(); for &page_id in self @@ -1830,10 +1828,7 @@ impl BTreeIndex { .collect::>(); debug!("Searching {} btree pages", page_tasks.len()); - // Collect row IDs from the pages. `buffered` keeps page order. When a `limit` is - // set we read one page at a time and stop once we have enough matches, so we do - // not issue I/O for pages we never need. Without a limit we fan out across CPUs - // (I/O and compute are mixed, but the important case is the index being cached). + // With a `limit`, read one page at a time and stop once we have enough; otherwise fan out across CPUs. let parallelism = if limit.is_some() { 1 } else { @@ -1845,8 +1840,7 @@ impl BTreeIndex { let mut matches_found: u64 = 0; while let Some(page_result) = page_stream.try_next().await? { if let Some(limit) = limit { - // Count only TRUE matches. NULL rows never match, so they must not count - // toward the limit. `len()` already excludes nulls. + // Count only TRUE matches toward the limit; `len()` already excludes nulls. matches_found += page_result.len().unwrap_or(0); results.push(page_result); if matches_found >= limit as u64 { @@ -5025,9 +5019,7 @@ mod tests { Arc::new(LanceCache::no_cache()), )); - // Enough rows to span several btree pages, with no nulls so every row matches an - // unbounded range. `train_btree_index` makes pages of `DEFAULT_BTREE_BATCH_SIZE` - // rows, so this gives five pages. + // Five btree pages of `DEFAULT_BTREE_BATCH_SIZE` rows, with 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); @@ -5063,8 +5055,7 @@ mod tests { let full_len = full.row_addrs().len().unwrap(); assert_eq!(full_len, num_rows); - // A limit that reaches into the second page. The search must satisfy it but stop - // well before reading all five pages. + // A limit reaching into the second page: satisfied but stops before reading all five. let limit = (DEFAULT_BTREE_BATCH_SIZE + 100) as usize; let limited = index .search_limited(&everything, &metrics, Some(limit)) diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index d5844c948bd..d4e05ffa058 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1327,8 +1327,7 @@ impl ScalarIndexExpr { limit: Option, ) -> Result { match self { - // A limit only applies to a single positive lookup. NOT, AND, and OR need the - // full result of each side, so the limit is dropped when recursing into them. + // A limit applies only to a single positive lookup, so drop it for NOT/AND/OR. Self::Not(inner) => { let result = inner.evaluate_nullable(index_loader, metrics, None).await?; Ok(!result) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index c8c6832ae38..192ff4b8d17 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4084,6 +4084,11 @@ impl Scanner { /// 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)`). In the default ordered mode the + /// scan returns the first matches in storage (row address) order, but a B-tree + /// stops after collecting matches in index-value page order. Those are different + /// subsets whenever storage order and index order disagree, so pushing the limit + /// would silently 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 @@ -4101,7 +4106,9 @@ impl Scanner { if limit <= 0 { return None; } - if self.ordering.is_some() + // Ordered scans return storage-order matches, while a B-tree stops in index-value order. + if self.ordered + || self.ordering.is_some() || self.nearest.is_some() || self.full_text_query.is_some() || self.aggregate.is_some() @@ -4150,8 +4157,7 @@ impl Scanner { .partition_frags_by_coverage(index_expr, fragments) .await?; - // A limit can be pushed into the index search, but only when its rows are used as - // is and the relevant fragments have no deletions. + // A limit can be pushed into the index search only when safe; see index_search_limit. let pushdown_limit = self.index_search_limit(filter_plan, &relevant_frags); let mut plan: Arc = Arc::new( @@ -5882,20 +5888,17 @@ mod test { #[tokio::test] async fn test_limit_pushed_into_scalar_index() { - // When a scan filter is fully served by a scalar index (no refine, no recheck, no - // ordering) the limit can be pushed into the index search. The result must still - // be exactly `limit` rows that all match the filter. Early stop must not drop or - // duplicate matches. + // 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 so a small limit short-circuits before the end. - let num_rows = 20_000; + // 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))], + vec![Arc::new(Int32Array::from_iter_values((0..num_rows).rev()))], ) .unwrap(); let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); @@ -5911,17 +5914,15 @@ mod test { .await .unwrap(); - let limit = 100; - let scan_ids = |dataset: Arc| async move { - let batch = dataset - .scan() - .filter("id >= 5") + let limit = 100i64; + let scan_ids = |dataset: Arc, ordered: bool| async move { + let mut scan = dataset.scan(); + scan.filter("id >= 5") .unwrap() + .scan_in_order(ordered) .limit(Some(limit), None) - .unwrap() - .try_into_batch() - .await .unwrap(); + let batch = scan.try_into_batch().await.unwrap(); batch .column_by_name("id") .unwrap() @@ -5930,18 +5931,26 @@ mod test { .to_vec() }; - let ids = scan_ids(Arc::new(dataset.clone())).await; + // 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).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)] + ); + + // Unordered scan: limit pushed into the index, but still exactly `limit` matching rows. + let ids = scan_ids(Arc::new(dataset.clone()), false).await; assert_eq!(ids.len(), limit as usize); assert!( ids.iter().all(|&id| id >= 5), "every returned row must satisfy the filter" ); - // With deletions present the limit must not be pushed, since deleted rows are - // pruned after the index search. The scan must still return exactly `limit` live - // matches. + // 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)).await; + let ids = scan_ids(Arc::new(dataset), false).await; assert_eq!(ids.len(), limit as usize); assert!( ids.iter().all(|&id| id >= 10000), diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 9d63e6b4368..dfbca7a0e5b 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -138,9 +138,7 @@ impl ScalarIndex for LogicalScalarIndex { metrics: &dyn MetricsCollector, limit: Option, ) -> Result { - // Forwarding the limit to every segment is safe. Each segment returns at least - // `limit` matches when it has them, so the combined result still has at least - // `limit` matches overall. + // Forwarding the limit to every segment is safe: the combined result still has at least `limit` matches. let results = try_join_all( self.segments .iter() From d6d9377a23cbf16345b5f821c246da502af12338 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 20 Jun 2026 15:46:31 +0200 Subject: [PATCH 03/14] fix(scanner): don't push scalar-index limit for fragment-subset scans --- rust/lance-index/src/scalar.rs | 7 +- rust/lance-index/src/scalar/btree.rs | 3 + rust/lance/src/dataset/scanner.rs | 118 +++++++++++++++++++++++-- rust/lance/src/index/scalar_logical.rs | 61 +++++++++++++ 4 files changed, 180 insertions(+), 9 deletions(-) diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 021a93e9ea8..aceffdce17f 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -1047,8 +1047,11 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { ) -> Result; /// Like [`Self::search`] but with a best-effort `limit` hint: when `limit` is `Some(n)` - /// an index may stop after finding `n` matches (it may still return more). Only push a - /// limit for a single positive lookup. The default ignores it and calls [`Self::search`]. + /// an index may stop after finding `n` matching rows (it may still return more). The hint + /// applies to positive lookups that keep matches as-is (equality, range, `IsIn`); negating + /// or combining operators ignore it. The caller must also discard null rows, since an index + /// may skip null tracking when a limit is set. The default ignores the hint and calls + /// [`Self::search`]. /// /// ``` /// # use lance_core::Result; diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index fc8470b9623..bd24a5e55e3 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2168,6 +2168,9 @@ impl BTreeIndex { // could refine that classification (see #6802). // // A `limit` implies a single positive lookup, so skip null tracking to stop early. + // Correctness then relies on the caller discarding nulls: every `search_limited` + // path goes through `evaluate_limited` -> `drop_nulls`, so the untracked null rows + // are dropped anyway. A future caller that keeps nulls must not pass a limit here. if limit.is_none() && !matches!(query, SargableQuery::IsNull()) { let existing: HashSet = pages.iter().map(|m| m.page_id()).collect(); for &page_id in self diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 678e194eec4..111dba9b467 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4107,6 +4107,10 @@ impl Scanner { /// those re-filters rows later and could drop matches. /// - The relevant fragments have no deletions. Deleted rows are pruned after the index /// search, so stopping early could leave fewer than `limit` live rows. + /// - The scan is not restricted to a fragment subset (`with_fragments`). The index search + /// runs over the whole dataset, and the fragment restriction is applied afterwards, so + /// an early stop could return `limit` matches that all fall outside the subset and leave + /// fewer than `limit` rows once it is applied. /// /// Returns `None` when no limit can be pushed. fn index_search_limit( @@ -4119,7 +4123,10 @@ impl Scanner { return None; } // Ordered scans return storage-order matches, while a B-tree stops in index-value order. + // A fragment subset is restricted only after the (global) index search, so an early stop + // could leave fewer than `limit` rows once the restriction is applied. if self.ordered + || self.fragments.is_some() || self.ordering.is_some() || self.nearest.is_some() || self.full_text_query.is_some() @@ -5898,8 +5905,14 @@ mod test { assert_eq!(ids, &(10..20).collect::>()); } + #[rstest] #[tokio::test] - async fn test_limit_pushed_into_scalar_index() { + 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", @@ -5914,7 +5927,13 @@ mod test { ) .unwrap(); let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone()); - let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + 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"], @@ -5927,12 +5946,12 @@ mod test { .unwrap(); let limit = 100i64; - let scan_ids = |dataset: Arc, ordered: bool| async move { + 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), None) + .limit(Some(limit), offset) .unwrap(); let batch = scan.try_into_batch().await.unwrap(); batch @@ -5944,7 +5963,7 @@ mod test { }; // 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).await; + 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), @@ -5953,16 +5972,26 @@ mod test { ); // Unordered scan: limit pushed into the index, but still exactly `limit` matching rows. - let ids = scan_ids(Arc::new(dataset.clone()), false).await; + 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" ); + // 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)); + // 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), false).await; + let ids = scan_ids(Arc::new(dataset), false, None).await; assert_eq!(ids.len(), limit as usize); assert!( ids.iter().all(|&id| id >= 10000), @@ -5970,6 +5999,81 @@ mod test { ); } + #[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" + ); + } + #[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/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index eda4b8ca117..4e1c8edea88 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -538,6 +538,67 @@ mod tests { ); } + #[tokio::test] + async fn test_btree_segment_search_limited_across_segments() { + // `search_limited` 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_limited(&query, &NoOpMetricsCollector, Some(limit)) + .await + .unwrap(); + let row_addrs = match result { + SearchResult::Exact(row_addrs) | SearchResult::AtLeast(row_addrs) => row_addrs, + other => panic!("unexpected result variant from limited search: {:?}", 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(); From c5dbe6d5c7ebb4011779766f450c2a63e3bfca71 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 27 Jun 2026 14:55:00 +0200 Subject: [PATCH 04/14] fix(scanner): gate scalar-index limit pushdown on retired fragments and return AtLeast --- rust/lance-index/src/scalar/btree.rs | 17 ++- rust/lance/src/dataset/scanner.rs | 195 ++++++++++++++++++++----- rust/lance/src/index/scalar_logical.rs | 7 +- rust/lance/src/io/exec/scalar_index.rs | 21 ++- 4 files changed, 197 insertions(+), 43 deletions(-) diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index bd24a5e55e3..f67965108af 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2236,7 +2236,16 @@ impl BTreeIndex { // Merge matching row IDs let selection = NullableRowAddrSet::union_all(&results); - Ok(SearchResult::Exact(selection)) + // A limited search may stop before reading every matching page, so the returned set is + // not the complete answer. Every row in it does satisfy the query (there may be more), + // which is exactly `AtLeast`. Reporting `Exact` here would let callers treat a partial + // match set as the full one. We conservatively report `AtLeast` for any limited search, + // even one that happened to read all pages, since the answer is still a valid lower bound. + Ok(if limit.is_some() { + SearchResult::AtLeast(selection) + } else { + SearchResult::Exact(selection) + }) } } @@ -5578,6 +5587,12 @@ mod tests { .search_limited(&everything, &metrics, Some(limit)) .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, diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 111dba9b467..a10637c02a6 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2861,15 +2861,23 @@ impl Scanner { fragments: Option>>, scan_range: Option>, ) -> Result> { - // Decide whether a limit can be pushed into the index search. The fragments the - // read covers (used for the deletion check) are the requested subset, or the whole - // dataset when none was given. + // 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 + // therefore 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 pushdown_limit = self.index_search_limit(filter_plan, scanned_fragments); + let pushdown_limit = self + .index_search_limit(filter_plan, scanned_fragments) + .await?; let mut read_options = FilteredReadOptions::basic_full_read(&self.dataset) .with_filter_plan(filter_plan.clone()) @@ -4105,51 +4113,62 @@ impl Scanner { /// - 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 relevant fragments have no deletions. Deleted rows are pruned after the index - /// search, so stopping early could leave fewer than `limit` live rows. - /// - The scan is not restricted to a fragment subset (`with_fragments`). The index search - /// runs over the whole dataset, and the fragment restriction is applied afterwards, so - /// an early stop could return `limit` matches that all fall outside the subset and leave - /// fewer than `limit` rows once it is applied. + /// - The index cannot yield row addresses that are filtered out after the search. The + /// index search returns row addresses from every fragment its segments cover, and those + /// that do not survive into the final result are pruned *after* the search. An early stop + /// would then spend its budget on rows that get dropped and could leave fewer than `limit` + /// live rows. A row address is dropped after the search when it belongs to a fragment that + /// has deletions (the deleted rows are masked out) or one that is not in the scanned set (a + /// retired/compacted-away fragment the index still has stale entries for, or a fragment + /// excluded by `with_fragments`). The single safe condition is therefore that every fragment + /// the index covers is in the scanned set *and* has no deletion file. /// /// Returns `None` when no limit can be pushed. - fn index_search_limit( + async fn index_search_limit( &self, filter_plan: &ExprFilterPlan, - relevant_fragments: &[Fragment], - ) -> Option { - let limit = self.limit?; + scanned_fragments: &[Fragment], + ) -> Result> { + let Some(limit) = self.limit else { + return Ok(None); + }; if limit <= 0 { - return None; + return Ok(None); } - // Ordered scans return storage-order matches, while a B-tree stops in index-value order. - // A fragment subset is restricted only after the (global) index search, so an early stop - // could leave fewer than `limit` rows once the restriction is applied. + // 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.fragments.is_some() || self.ordering.is_some() || self.nearest.is_some() || self.full_text_query.is_some() || self.aggregate.is_some() || filter_plan.has_refine() { - return None; - } - if filter_plan - .index_query - .as_ref() - .is_some_and(|query| query.needs_recheck()) - { - return None; + return Ok(None); } - if relevant_fragments + let Some(index_query) = filter_plan.index_query.as_ref() else { + return Ok(None); + }; + if index_query.needs_recheck() { + 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), they never remove index hits. + let live_undeleted: RoaringBitmap = scanned_fragments .iter() - .any(|fragment| fragment.deletion_file.is_some()) - { - return None; + .filter(|fragment| fragment.deletion_file.is_none()) + .map(|fragment| fragment.id as u32) + .collect(); + let covered_frags = self.fragments_covered_by_index_query(index_query).await?; + if !covered_frags.is_subset(&live_undeleted) { + return Ok(None); } let offset = self.offset.unwrap_or(0).max(0) as usize; - Some((limit as usize).saturating_add(offset)) + 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 @@ -4177,7 +4196,12 @@ impl Scanner { .await?; // A limit can be pushed into the index search only when safe; see index_search_limit. - let pushdown_limit = self.index_search_limit(filter_plan, &relevant_frags); + // `relevant_frags` is `covered ∩ scanned`, so requiring the index's covered fragments to + // be a subset of it rejects both retired/uncovered-scanned fragments and `with_fragments` + // subsets that drop covered fragments. + let pushdown_limit = self + .index_search_limit(filter_plan, &relevant_frags) + .await?; let mut plan: Arc = Arc::new( MaterializeIndexExec::new( @@ -6074,6 +6098,111 @@ mod test { ); } + #[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" + ); + } + #[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/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 4e1c8edea88..6575f1bb2da 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -588,9 +588,12 @@ mod tests { .search_limited(&query, &NoOpMetricsCollector, Some(limit)) .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::Exact(row_addrs) | SearchResult::AtLeast(row_addrs) => row_addrs, - other => panic!("unexpected result variant from limited search: {:?}", other), + 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!( diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index 864a65229fe..b4a51c50955 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -700,19 +700,26 @@ 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 { + // reach the user). + // + // `AtLeast` carries an unbounded `upper`, so we cannot use it as the + // candidate set. Instead we materialize its `lower` mask: every row + // in `lower` is a guaranteed match, so it is a sound (possibly + // partial) answer. This is exactly what the limit pushdown produces — + // the B-tree stops early and returns a confirmed lower bound, and a + // downstream `GlobalLimitExec` still enforces the exact limit. + let candidate_mask = |result: IndexExprResult| -> Result { if result.is_at_least() && !result.is_exact() { - todo!("Support AtLeast in MaterializeIndexExec") + Ok(result.lower) + } else { + Ok(result.upper) } - Ok(result.upper) }; let 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?)? }; let ids = row_ids_for_mask(mask, &dataset, &fragments).await?; let ids = UInt64Array::from(ids); From d2132de18dd4093bb6c4d24418a2f1dd98d9db26 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:13:40 +0300 Subject: [PATCH 05/14] simplified documentation --- rust/lance-index/src/scalar.rs | 12 ++--- rust/lance-index/src/scalar/btree.rs | 16 +++---- rust/lance/src/dataset/scanner.rs | 63 +++++++++++++------------- rust/lance/src/index/scalar_logical.rs | 7 +-- rust/lance/src/io/exec/scalar_index.rs | 6 +-- 5 files changed, 52 insertions(+), 52 deletions(-) diff --git a/rust/lance-index/src/scalar.rs b/rust/lance-index/src/scalar.rs index 85b44444c69..25e36e571d3 100644 --- a/rust/lance-index/src/scalar.rs +++ b/rust/lance-index/src/scalar.rs @@ -1103,12 +1103,12 @@ pub trait ScalarIndex: Send + Sync + std::fmt::Debug + Index + DeepSizeOf { metrics: &dyn MetricsCollector, ) -> Result; - /// Like [`Self::search`] but with a best-effort `limit` hint: when `limit` is `Some(n)` - /// an index may stop after finding `n` matching rows (it may still return more). The hint - /// applies to positive lookups that keep matches as-is (equality, range, `IsIn`); negating - /// or combining operators ignore it. The caller must also discard null rows, since an index - /// may skip null tracking when a limit is set. The default ignores the hint and calls - /// [`Self::search`]. + /// Like [`Self::search`] but with a best-effort `limit` hint. When `limit` is `Some(n)` + /// an index may stop after finding `n` matching rows and may still return more. The hint + /// applies to positive lookups that keep matches as-is, such as equality, range, and + /// `IsIn`. Negating or combining operators ignore it. The caller must also discard null + /// rows, since an index may skip null tracking when a limit is set. The default ignores + /// the hint and calls [`Self::search`]. /// /// ``` /// # use lance_core::Result; diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 8292b90e538..5d00810a336 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2176,8 +2176,8 @@ impl BTreeIndex { // could refine that classification (see #6802). // // A `limit` implies a single positive lookup, so skip null tracking to stop early. - // Correctness then relies on the caller discarding nulls: every `search_limited` - // path goes through `evaluate_limited` -> `drop_nulls`, so the untracked null rows + // Correctness then relies on the caller discarding nulls. Every `search_limited` + // path goes through `evaluate_limited` and `drop_nulls`, so the untracked null rows // are dropped anyway. A future caller that keeps nulls must not pass a limit here. if limit.is_none() && !matches!(query, SargableQuery::IsNull()) { let existing: HashSet = pages.iter().map(|m| m.page_id()).collect(); @@ -2218,7 +2218,7 @@ impl BTreeIndex { .collect::>(); debug!("Searching {} btree pages", page_tasks.len()); - // With a `limit`, read one page at a time and stop once we have enough; otherwise fan out across CPUs. + // With a `limit`, read one page at a time and stop once we have enough. Otherwise fan out across CPUs. let parallelism = if limit.is_some() { 1 } else { @@ -2230,7 +2230,7 @@ impl BTreeIndex { let mut matches_found: u64 = 0; while let Some(page_result) = page_stream.try_next().await? { if let Some(limit) = limit { - // Count only TRUE matches toward the limit; `len()` already excludes nulls. + // Count only TRUE matches toward the limit. `len()` already excludes nulls. matches_found += page_result.len().unwrap_or(0); results.push(page_result); if matches_found >= limit as u64 { @@ -2245,10 +2245,10 @@ impl BTreeIndex { let selection = NullableRowAddrSet::union_all(&results); // A limited search may stop before reading every matching page, so the returned set is - // not the complete answer. Every row in it does satisfy the query (there may be more), - // which is exactly `AtLeast`. Reporting `Exact` here would let callers treat a partial - // match set as the full one. We conservatively report `AtLeast` for any limited search, - // even one that happened to read all pages, since the answer is still a valid lower bound. + // not the complete answer. Every row in it still satisfies the query, so the result is a + // valid lower bound, which is exactly `AtLeast`. Reporting `Exact` would let callers treat + // a partial match set as complete, so report `AtLeast` for any limited search even when it + // happened to read all pages. Ok(if limit.is_some() { SearchResult::AtLeast(selection) } else { diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 53c96df24a2..e1e22b33434 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -2892,14 +2892,14 @@ impl Scanner { scan_range: Option>, ) -> 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 + // 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 - // therefore mutually exclusive and need no extra coordination here. + // 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() @@ -4135,24 +4135,19 @@ impl Scanner { /// 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)`). In the default ordered mode the - /// scan returns the first matches in storage (row address) order, but a B-tree - /// stops after collecting matches in index-value page order. Those are different - /// subsets whenever storage order and index order disagree, so pushing the limit - /// would silently change which rows `LIMIT`/`OFFSET` returns. + /// - 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 cannot yield row addresses that are filtered out after the search. The - /// index search returns row addresses from every fragment its segments cover, and those - /// that do not survive into the final result are pruned *after* the search. An early stop - /// would then spend its budget on rows that get dropped and could leave fewer than `limit` - /// live rows. A row address is dropped after the search when it belongs to a fragment that - /// has deletions (the deleted rows are masked out) or one that is not in the scanned set (a - /// retired/compacted-away fragment the index still has stale entries for, or a fragment - /// excluded by `with_fragments`). The single safe condition is therefore that every fragment - /// the index covers is in the scanned set *and* has no deletion file. + /// - 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( @@ -4186,9 +4181,9 @@ impl Scanner { } // 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), they never remove index hits. + // 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()) @@ -4226,10 +4221,10 @@ 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 `covered ∩ scanned`, so requiring the index's covered fragments to - // be a subset of it rejects both retired/uncovered-scanned fragments and `with_fragments` - // subsets that drop covered fragments. + // 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?; @@ -6125,17 +6120,19 @@ mod test { #[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. + // 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. + // 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. + // 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(), @@ -6178,7 +6175,8 @@ mod test { .to_vec() }; - // Ordered scan (the default): limit not pushed, so the first matches are the largest ids (descending storage). + // 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!( @@ -6205,7 +6203,8 @@ mod test { ); assert!(ids.iter().all(|&id| id >= 5)); - // With deletions the limit must not be pushed even when unordered, since deleted rows are pruned after the index search. + // 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), false, None).await; assert_eq!(ids.len(), limit as usize); @@ -6221,9 +6220,9 @@ mod test { #[values(LanceFileVersion::Legacy, LanceFileVersion::Stable)] data_storage_version: LanceFileVersion, ) { - // The scalar-index search runs over the whole dataset; a `with_fragments` subset is + // 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 + // 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", @@ -6301,7 +6300,7 @@ mod test { // 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. + // fragment, even though no live fragment has a deletion file. let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( "id", DataType::Int32, diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 82951ba627a..596569ef395 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -142,7 +142,8 @@ impl ScalarIndex for LogicalScalarIndex { metrics: &dyn MetricsCollector, limit: Option, ) -> Result { - // Forwarding the limit to every segment is safe: the combined result still has at least `limit` matches. + // Forwarding the limit to every segment is safe. The combined result still has at least + // `limit` matches. let results = try_join_all( self.segments .iter() @@ -580,7 +581,7 @@ mod tests { .await .unwrap(); - // All 64 rows match the unbounded range; with a limit the combined result across the + // 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; @@ -589,7 +590,7 @@ mod tests { .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 + // 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, diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index 03859c8d4a0..bf615dff93a 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -810,10 +810,10 @@ impl MaterializeIndexExec { // reach the user). // // `AtLeast` carries an unbounded `upper`, so we cannot use it as the - // candidate set. Instead we materialize its `lower` mask: every row + // candidate set. Instead we materialize its `lower` mask. Every row // in `lower` is a guaranteed match, so it is a sound (possibly - // partial) answer. This is exactly what the limit pushdown produces — - // the B-tree stops early and returns a confirmed lower bound, and a + // partial) answer. This is exactly what the limit pushdown produces. + // The B-tree stops early and returns a confirmed lower bound, and a // downstream `GlobalLimitExec` still enforces the exact limit. let candidate_mask = |result: IndexExprResult| -> Result { if result.is_at_least() && !result.is_exact() { From 2abbbcda4a36ecdc7cb5fb68dd9dc1a91191a110 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:49:04 +0300 Subject: [PATCH 06/14] fix(scanner): restrict scalar-index limit pushdown to single-lookup filters --- rust/lance/src/dataset/scanner.rs | 102 ++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index e1e22b33434..31473b6e6b1 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4142,6 +4142,9 @@ impl Scanner { /// - 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 @@ -4179,6 +4182,14 @@ impl Scanner { 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. + if !matches!(index_query, ScalarIndexExpr::Query(_)) { + 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 @@ -6394,6 +6405,97 @@ mod test { ); } + #[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).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).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" + ); + } + #[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 From 2574c15bc3eb7cbb8aa729bf0c5b7389be3cdf13 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Fri, 10 Jul 2026 23:18:53 +0300 Subject: [PATCH 07/14] fix(scanner): don't error the limit pushdown on unknown index coverage --- rust/lance/src/dataset/scanner.rs | 67 +++++++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 3 deletions(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 31473b6e6b1..514aec10f81 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4187,9 +4187,9 @@ impl Scanner { // 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. - if !matches!(index_query, ScalarIndexExpr::Query(_)) { + 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 @@ -4200,7 +4200,16 @@ impl Scanner { .filter(|fragment| fragment.deletion_file.is_none()) .map(|fragment| fragment.id as u32) .collect(); - let covered_frags = self.fragments_covered_by_index_query(index_query).await?; + // 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); } @@ -6496,6 +6505,58 @@ mod test { ); } + #[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 From e809b39d8527633835f4fe5a046b0d3e0e399a26 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:33:52 +0300 Subject: [PATCH 08/14] refactor(scalar-index): remove redundant offset clamp --- rust/lance/src/dataset/scanner.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 514aec10f81..fe70311cb4f 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4213,7 +4213,7 @@ impl Scanner { if !covered_frags.is_subset(&live_undeleted) { return Ok(None); } - let offset = self.offset.unwrap_or(0).max(0) as usize; + let offset = self.offset.unwrap_or(0) as usize; Ok(Some((limit as usize).saturating_add(offset))) } From 160a1b091d1a6ffba9bf9ec01c828c8921b284d3 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:14:46 +0200 Subject: [PATCH 09/14] fix(scalar-index): retain limited expression evaluation --- rust/lance-index/src/scalar/expression.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index c301641070d..970c52ab91e 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1688,21 +1688,22 @@ impl ScalarIndexExpr { &self, index_loader: &dyn ScalarIndexLoader, metrics: &dyn MetricsCollector, + limit: Option, ) -> Result { match self { Self::Not(inner) => { - let result = inner.evaluate_nullable(index_loader, metrics).await?; + let result = inner.evaluate_nullable(index_loader, metrics, None).await?; Ok(!result) } Self::And(lhs, rhs) => { - let lhs_result = lhs.evaluate_nullable(index_loader, metrics); - let rhs_result = rhs.evaluate_nullable(index_loader, metrics); + let lhs_result = lhs.evaluate_nullable(index_loader, metrics, None); + let rhs_result = rhs.evaluate_nullable(index_loader, metrics, None); 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_nullable(index_loader, metrics); - let rhs_result = rhs.evaluate_nullable(index_loader, metrics); + let lhs_result = lhs.evaluate_nullable(index_loader, metrics, None); + let rhs_result = rhs.evaluate_nullable(index_loader, metrics, None); let (lhs_result, rhs_result) = try_join!(lhs_result, rhs_result)?; Ok(lhs_result | rhs_result) } From 384699e1f591e109d2fa8b3c8e7e31d1c019b434 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:15:13 +0200 Subject: [PATCH 10/14] fix(scalar-index): restore limited evaluation --- rust/lance-index/src/scalar/expression.rs | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index 970c52ab91e..5d38f390f82 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -1734,7 +1734,25 @@ impl ScalarIndexExpr { metrics: &dyn MetricsCollector, ) -> Result { Ok(self - .evaluate_nullable(index_loader, metrics) + .evaluate_nullable(index_loader, metrics, None) + .await? + .drop_nulls()) + } + + /// Like [`Self::evaluate`] but pushes a `limit` hint into the index search so it can + /// stop once it has found at least `limit` matches. + /// + /// See [`crate::scalar::ScalarIndex::search_limited`] for the rules on when a limit + /// may be pushed down. + #[instrument(level = "debug", skip_all)] + pub async fn evaluate_limited( + &self, + index_loader: &dyn ScalarIndexLoader, + metrics: &dyn MetricsCollector, + limit: Option, + ) -> Result { + Ok(self + .evaluate_nullable(index_loader, metrics, limit) .await? .drop_nulls()) } From ae9aef37d49bcd617bcd3ddfcd1127e547ae45b4 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:50:41 +0300 Subject: [PATCH 11/14] style: fix overlayed/overlaid typos flagged by spell check --- rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 5c84a39242e..e0aa2cd7f85 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -321,7 +321,7 @@ async fn test_overlay_block_does_not_short_limited_index_scan() { let filter = format!("age < {THRESHOLD}"); - // Ground truth: every non-overlayed row still matches. + // Ground truth: every non-overlaid row still matches. let mut unlimited = dataset.scan(); unlimited .filter(&filter) @@ -333,7 +333,7 @@ async fn test_overlay_block_does_not_short_limited_index_scan() { assert_eq!( total, (NUM_ROWS - STALE) as usize, - "overlayed rows must drop out of the unlimited result" + "overlaid rows must drop out of the unlimited result" ); assert!( (PAGE - STALE) < LIMIT as i32, From 22a3029bd54ec54996c72ea7fb7d2fc320a2ffdb Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 22 Aug 2026 18:44:55 +0300 Subject: [PATCH 12/14] fix(scanner): only collapse AtLeast to its lower bound for a pushed limit --- rust/lance/src/io/exec/scalar_index.rs | 83 ++++++++++++++++++++++---- 1 file changed, 70 insertions(+), 13 deletions(-) diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index 82099778d28..73a77a6ef00 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -781,6 +781,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, @@ -849,19 +876,8 @@ impl MaterializeIndexExec { // non-matching candidates in `upper` are dropped before they // reach the user). // - // `AtLeast` carries an unbounded `upper`, so we cannot use it as the - // candidate set. Instead we materialize its `lower` mask. Every row - // in `lower` is a guaranteed match, so it is a sound (possibly - // partial) answer. This is exactly what the limit pushdown produces. - // The B-tree stops early and returns a confirmed lower bound, and a - // downstream `GlobalLimitExec` still enforces the exact limit. - let candidate_mask = |result: IndexExprResult| -> Result { - if result.is_at_least() && !result.is_exact() { - Ok(result.lower) - } else { - Ok(result.upper) - } - }; + 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)?; candidate_mask(expr_result)? & (*prefilter).clone() @@ -1040,6 +1056,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; From 436f3d96a434e31e29a1ec2e4b0206bd24d7908e Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sat, 22 Aug 2026 19:48:53 +0300 Subject: [PATCH 13/14] test(scanner): assert limit pushdown reaches the index exec nodes --- rust/lance-index-core/src/scalar.rs | 2 +- rust/lance-index/src/scalar/btree.rs | 6 ++-- rust/lance-index/src/scalar/expression.rs | 2 +- rust/lance/src/dataset/scanner.rs | 35 ++++++++++++++++++- .../tests/dataset_overlay_index_masking.rs | 4 +-- rust/lance/src/index/scalar_logical.rs | 9 ++++- rust/lance/src/io/exec/scalar_index.rs | 26 +++++++++++--- 7 files changed, 71 insertions(+), 13 deletions(-) diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index bd7887ca381..eee7d06242a 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -512,7 +512,7 @@ pub struct SearchOptions { /// /// 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 -- equality, range, `IsIn`. + /// 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. /// diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index e9b48eea0ce..43d0721940a 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2248,7 +2248,7 @@ impl ScalarIndex for BTreeIndex { // 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. + // 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 options.limit().is_some() { 1 @@ -2260,7 +2260,7 @@ impl ScalarIndex for BTreeIndex { let mut matches_found: u64 = 0; while let Some(page_result) = page_stream.try_next().await? { if let Some(limit) = options.limit() { - // Count only TRUE matches toward the limit; `len()` already excludes nulls. + // Count only TRUE matches toward the limit. `len()` already excludes nulls. matches_found += page_result.len().unwrap_or(0); results.push(page_result); if matches_found >= limit as u64 { @@ -2287,7 +2287,7 @@ impl ScalarIndex for BTreeIndex { // 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, + // 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. Ok(if options.limit().is_some() { SearchResult::AtLeast(selection) diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index 4d89107abd6..39ea8fc4f68 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -2012,7 +2012,7 @@ 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`, + /// 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 diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index b8aa8fed972..267254d62db 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -7936,6 +7936,19 @@ mod test { .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; @@ -7945,6 +7958,11 @@ mod test { "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; @@ -7953,6 +7971,11 @@ mod test { 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)`. @@ -7963,16 +7986,26 @@ mod test { "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), false, None).await; + 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] 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 e0aa2cd7f85..a828dd6ba47 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -271,7 +271,7 @@ async fn test_overlay_stale_drop_and_new_match(#[values(false, true)] stable_row /// /// 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. +/// 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; @@ -305,7 +305,7 @@ async fn test_overlay_block_does_not_short_limited_index_scan() { .unwrap(); build_age_index(&mut dataset).await; - // Push the first STALE rows -- the lowest index-order matches -- above the predicate so + // 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, diff --git a/rust/lance/src/index/scalar_logical.rs b/rust/lance/src/index/scalar_logical.rs index 706ae7d5935..3c3ac6c1add 100644 --- a/rust/lance/src/index/scalar_logical.rs +++ b/rust/lance/src/index/scalar_logical.rs @@ -156,7 +156,14 @@ impl ScalarIndex for LogicalScalarIndex { ) -> 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; `combine_search_results` keeps the result `AtLeast`. + // 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() diff --git a/rust/lance/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index 73a77a6ef00..d05e3bec2d3 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -132,10 +132,20 @@ 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(()) } } } @@ -733,10 +743,18 @@ 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(()) } } } From d9a736216a05d41b1f4e2e5014d36ed6eec5abe3 Mon Sep 17 00:00:00 2001 From: gstamatakis95 <126914070+gstamatakis95@users.noreply.github.com> Date: Sun, 23 Aug 2026 11:48:48 +0300 Subject: [PATCH 14/14] fix(btree): decline the pushed search limit when nulls are tracked --- rust/lance-index-core/src/scalar.rs | 7 +- rust/lance-index/src/scalar/btree.rs | 105 +++++++++++++++++++++++-- rust/lance/src/io/exec/scalar_index.rs | 1 - 3 files changed, 104 insertions(+), 9 deletions(-) diff --git a/rust/lance-index-core/src/scalar.rs b/rust/lance-index-core/src/scalar.rs index eee7d06242a..2bd497da61b 100644 --- a/rust/lance-index-core/src/scalar.rs +++ b/rust/lance-index-core/src/scalar.rs @@ -516,9 +516,12 @@ pub struct SearchOptions { /// Negating and combining operators ignore it, because a partial match set /// cannot be complemented or intersected soundly. /// - /// A limited search reports [`SearchResult::AtLeast`], never + /// 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. + /// 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, } diff --git a/rust/lance-index/src/scalar/btree.rs b/rust/lance-index/src/scalar/btree.rs index 43d0721940a..5a875ce835a 100644 --- a/rust/lance-index/src/scalar/btree.rs +++ b/rust/lance-index/src/scalar/btree.rs @@ -2245,12 +2245,23 @@ impl ScalarIndex for BTreeIndex { .collect::>(); debug!("Searching {} btree pages", page_tasks.len()); + // 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 options.limit().is_some() { + let parallelism = if limit.is_some() { 1 } else { get_num_compute_intensive_cpus() @@ -2259,9 +2270,10 @@ impl ScalarIndex for BTreeIndex { 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) = options.limit() { - // Count only TRUE matches toward the limit. `len()` already excludes nulls. - matches_found += page_result.len().unwrap_or(0); + 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; @@ -2288,8 +2300,9 @@ impl ScalarIndex for BTreeIndex { // 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. - Ok(if options.limit().is_some() { + // 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) @@ -5714,6 +5727,86 @@ mod tests { ); } + /// 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/src/io/exec/scalar_index.rs b/rust/lance/src/io/exec/scalar_index.rs index d05e3bec2d3..96285f3b2aa 100644 --- a/rust/lance/src/io/exec/scalar_index.rs +++ b/rust/lance/src/io/exec/scalar_index.rs @@ -893,7 +893,6 @@ impl MaterializeIndexExec { // runs on the materialized batches via the scan plan, so any // non-matching candidates in `upper` are dropped before they // reach the user). - // let candidate_mask = |result: IndexExprResult| -> Result { candidate_mask_for(result, limit) }; let mut mask = if let Some(prefilter) = prefilter {