From a8305e072246d21a97e8857c480e3bbdb9a4178f Mon Sep 17 00:00:00 2001 From: Romin Parekh <3831468+rominparekh@users.noreply.github.com> Date: Fri, 23 Oct 2026 20:08:40 -0700 Subject: [PATCH] fix(rust): accept usize::MAX as a local query limit (#4567) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #2211. Local queries converted `limit`/`offset` with `as i64`, so `limit(usize::MAX)` wrapped to `-1` and Lance failed with `Limit must be non-negative`. The remote client already worked around this (#2354); this fixes the local path. - Saturate scanner offsets to `i64::MAX` and cap limits at `i64::MAX - offset`, keeping both conversion and full-text pagination addition safe. - Use `saturating_add` for vector `top_k = limit + offset`, then cap candidates at the dataset row count to avoid indexed-search capacity and refinement overflows. Keep a positive candidate bound for empty tables because Lance rejects `k = 0`. - Leave public APIs and remote query handling unchanged. ### Testing - Regression coverage includes plain scans, indexed IVF search, refinement, index bypass, empty vector tables, and indexed full-text pagination with finite and extreme limits/offsets. - Reproduced the indexed-search and FTS overflows before applying the follow-up fix. - `cargo test --quiet --features remote -p lancedb --lib query::` — 81 passed. - `cargo clippy --quiet --features remote -p lancedb --tests --examples` — clean. - `cargo fmt --all` and `git diff --check` — clean. Vector planning now reads the dataset row count (metadata fast path) to bound candidate allocation; this was not performance-benchmarked. --------- Co-authored-by: Romin Parekh Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 29 +++--- rust/lancedb/src/query.rs | 153 ++++++++++++++++++++++++++++++++ rust/lancedb/src/table/query.rs | 21 ++++- 3 files changed, 188 insertions(+), 15 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 9d7610f72..52f38d077 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -1561,9 +1561,10 @@ class LanceQueryBuilder(ABC): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -1619,8 +1620,8 @@ class LanceQueryBuilder(ABC): GlobalLimitExec: elapsed=..., skip=0, fetch=10, metrics=[output_rows=..., elapsed_compute=..., output_bytes=...] FilterExec: elapsed=..., _distance@2 IS NOT NULL, metrics=[...] - SortExec: elapsed=..., TopK(fetch=10), expr=[...], - preserve_partitioning=[...], + SortExec: elapsed=..., TopK(fetch=1), expr=[...], + preserve_partitioning=[...], filter=[...], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., row_replacements=...] KNNVectorDistance: elapsed=..., metric=l2, @@ -1633,6 +1634,7 @@ class LanceQueryBuilder(ABC): metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., fragments_scanned=..., ranges_scanned=1, rows_scanned=1, bytes_read=..., iops=..., requests=..., task_wait_time=...] + Parameters ---------- @@ -2959,9 +2961,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters ---------- @@ -3601,9 +3604,10 @@ class AsyncQueryBase(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters @@ -4521,9 +4525,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): LanceRead: uri=..., projection=[text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score, _rowid@0 as _rowid] LanceRead: uri=..., projection=[vector, text], source=stream(_rowid) GlobalLimitExec: skip=0, fetch=10 @@ -4877,9 +4881,10 @@ class BaseQueryBuilder(object): ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance] GlobalLimitExec: skip=0, fetch=10 FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=1), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false] KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vector], ... + LanceRead: uri=..., projection=[vector], num_fragments=..., ... + Parameters diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 00a2b88f3..9bac51209 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -2714,6 +2714,159 @@ mod tests { } } + #[tokio::test] + async fn test_limit_usize_max() { + let conn = connect("memory://foo").execute().await.unwrap(); + let table = conn + .create_table("my_table", make_non_empty_batches()) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected); + + let batches: Vec = table + .query() + .limit(usize::MAX) + .offset(1) + .nearest_to(&[0.1; 4]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, expected - 1); + } + + #[tokio::test] + async fn test_indexed_vector_limit_usize_max() { + use crate::index::vector::IvfFlatIndexBuilder; + + let conn = connect("memory://indexed_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_non_empty_batches()) + .execute() + .await + .unwrap(); + table + .create_index( + &["vector"], + Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(1)), + ) + .execute() + .await + .unwrap(); + let expected = table.count_rows(None).await.unwrap(); + + for limit in [expected, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, usize::MAX] { + for (bypass, refine_factor) in [(false, None), (false, Some(2)), (true, None)] { + let query = table + .query() + .limit(limit) + .offset(offset) + .nearest_to(&[0.1; 4]) + .unwrap(); + let query = if bypass { + query.bypass_vector_index() + } else { + query + }; + let query = if let Some(factor) = refine_factor { + query.refine_factor(factor) + } else { + query + }; + let batches: Vec = + query.execute().await.unwrap().try_collect().await.unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + expected.saturating_sub(offset), + "limit={limit}, offset={offset}, bypass={bypass}, refine={refine_factor:?}" + ); + } + } + } + } + + #[tokio::test] + async fn test_empty_vector_limit_usize_max() { + let conn = connect("memory://empty_limit").execute().await.unwrap(); + let table = conn + .create_table("vectors", make_test_batches()) + .execute() + .await + .unwrap(); + let batches: Vec = table + .query() + .limit(usize::MAX) + .nearest_to(&[0.1; 128]) + .unwrap() + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 0); + } + + #[tokio::test] + async fn test_fts_limit_usize_max_with_offset() { + let conn = connect("memory://fts_limit").execute().await.unwrap(); + let data = RecordBatch::try_from_iter([( + "text", + Arc::new(StringArray::from_iter_values((0..32).map(|_| "hello"))) + as arrow_array::ArrayRef, + )]) + .unwrap(); + let table = conn + .create_table("documents", data) + .execute() + .await + .unwrap(); + table + .create_index(&["text"], Index::FTS(Default::default())) + .execute() + .await + .unwrap(); + + for limit in [32, i64::MAX as usize, usize::MAX] { + for offset in [0, 1, i64::MAX as usize, usize::MAX] { + let batches: Vec = table + .query() + .full_text_search(FullTextSearchQuery::new("hello".into())) + .limit(limit) + .offset(offset) + .execute() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 32usize.saturating_sub(offset), + "limit={limit}, offset={offset}" + ); + } + } + } + #[tokio::test] async fn test_repeated_only_if_combines_with_and() { use crate::expr::{col, lit}; diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index f625a4aff..4ad840350 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -280,7 +280,13 @@ pub async fn create_plan( let (_, element_type) = lance::index::vector::utils::get_vector_type(schema, &column)?; let is_binary = matches!(element_type, DataType::UInt8); - let top_k = query.base.limit.unwrap_or(DEFAULT_TOP_K) + query.base.offset.unwrap_or(0); + let top_k = query + .base + .limit + .unwrap_or(DEFAULT_TOP_K) + .saturating_add(query.base.offset.unwrap_or(0)) + // Lance allocates candidates from k; no search can return more rows than the table. + .min(ds_ref.count_rows(None).await?.max(1)); if is_binary { let query_vector = arrow::compute::cast(&query_vector, &DataType::UInt8)?; @@ -317,9 +323,18 @@ pub async fn create_plan( // For a batch query, `nearest` already applies k to each query vector. // Adding Scanner's global limit would truncate the combined result to k rows. if !is_batch_query { + let offset = query + .base + .offset + .map(|offset| i64::try_from(offset).unwrap_or(i64::MAX)); scanner.limit( - query.base.limit.map(|limit| limit as i64), - query.base.offset.map(|offset| offset as i64), + // Lance's FTS planner adds limit + offset as i64. + query.base.limit.map(|limit| { + i64::try_from(limit) + .unwrap_or(i64::MAX) + .min(i64::MAX - offset.unwrap_or(0)) + }), + offset, )?; }