rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-25 21:22:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rust): accept usize::MAX as a local query limit (#4567)

Browse Source 
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 <roparekh@adobe.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Xuanwo <github@xuanwo.io>
This commit is contained in:
Romin Parekh
2026-10-24 11:08:40 +08:00
committed by GitHub
co-authored by Romin Parekh Copilot Xuanwo
parent 6777a40dc2
commit a8305e0722
3 changed files with 188 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/query.py
+17 -12
View File
@@ -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=..., ...
<BLANKLINE>
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=...]
<BLANKLINE>
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=..., ...
<BLANKLINE>
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=..., ...
<BLANKLINE>
<BLANKLINE>
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=..., ...
<BLANKLINE>
<BLANKLINE>
Parameters
rust/lancedb/src/query.rs
+153
View File
@@ -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<RecordBatch> = 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<RecordBatch> = 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<RecordBatch> =
query.execute().await.unwrap().try_collect().await.unwrap();
assert_eq!(
batches.iter().map(RecordBatch::num_rows).sum::<usize>(),
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<RecordBatch> = 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::<usize>(), 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<RecordBatch> = 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::<usize>(),
32usize.saturating_sub(offset),
"limit={limit}, offset={offset}"
);
}
}
}
#[tokio::test]
async fn test_repeated_only_if_combines_with_and() {
use crate::expr::{col, lit};
rust/lancedb/src/table/query.rs
+18 -3
View File
@@ -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,
)?;
}