fix: reject zero vector search probes across SDKs (#4418)

`nprobes(0)` bypassed the validation used by the minimum and maximum
probe setters, allowing zero into query requests and producing
inconsistent local and remote results. Sync Python could additionally
interpret the zero maximum as an unbounded search.

Reject zero in the Rust vector-query builder and propagate its input
error through the Python vector/hybrid and Node bindings. Validate
nonpositive counts in the sync Python vector and hybrid builders so
every SDK and backend reports `nprobes must be greater than 0` before
execution. Update the Rust examples and document the positive-count
requirement.

Regression tests cover local and remote vector/hybrid builders,
rejection before remote query requests, promised TypeScript query
vectors, and positive counts replacing both probe bounds.

Validation:

- Rust probe regression tests: 2 passed.
- Focused Python query, remote, hybrid, and serialization tests: 22
passed. The four sync regression cases fail with the original setters.
- Focused TypeScript indexed local and mocked remote tests: 2 passed.
- Workspace `cargo check --features remote --tests --examples` and
Clippy passed.
- Rust formatting, Python formatting/lint, and TypeScript
build/lint/docs passed. Repository-wide `ruff check .` reports 29
existing errors in unchanged CI scripts; `ruff check python` passes.

Fixes #4388

BREAKING CHANGE: Rust `VectorQuery::nprobes` now returns `Result<Self>`
instead of `Self`; callers must handle the result, as with the existing
minimum/maximum probe setters.

<!-- lance-gatekeeper-fix:v1 agent=e92fda5a212c7f9e80fa55e7b4e4634d
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
This commit is contained in:
lancedb-gatefixer[bot]
2026-10-03 21:29:47 -07:00
committed by GitHub
parent 0be3ae960e
commit 3e0a1a8780
13 changed files with 188 additions and 11 deletions
+2
View File
@@ -447,6 +447,8 @@ nprobes(nprobes): VectorQuery
Set the number of partitions to search (probe)
The number of probes must be greater than 0.
This argument is only used when the vector column has an IVF PQ index.
If there is no index then this value is ignored.
+32
View File
@@ -93,6 +93,38 @@ async function withMockDatabase(
}
describe("remote connection", () => {
it("rejects nprobes(0) before sending a query", async () => {
const requests: string[] = [];
await withMockDatabase(
(req, res) => {
requests.push(req.url ?? "");
if (req.url === "/v1/table/test/describe/") {
res
.writeHead(200, { "Content-Type": "application/json" })
.end(JSON.stringify({ version: 1, schema: { fields: [] } }));
} else {
res.writeHead(404).end();
}
},
async (db) => {
const table = await db.openTable("test");
expect(() => table.vectorSearch([0, 0]).nprobes(0)).toThrow(
"Invalid input, nprobes must be greater than 0",
);
expect(() =>
table.query().nearestTo([0, 0]).fullTextSearch("dog").nprobes(0),
).toThrow("Invalid input, nprobes must be greater than 0");
await expect(
table
.vectorSearch(Promise.resolve([0, 0]))
.nprobes(0)
.toArrow(),
).rejects.toThrow("Invalid input, nprobes must be greater than 0");
},
);
expect(requests).toEqual(["/v1/table/test/describe/"]);
});
it.each([false, true])(
"preserves an empty query's schema with an empty batch: %s",
async (withEmptyBatch) => {
+6
View File
@@ -1211,6 +1211,12 @@ describe("When creating an index", () => {
.toArrow();
expect(rst.numRows).toBe(2);
expect(() => tbl.vectorSearch(queryVec).nprobes(0)).toThrow(
"Invalid input, nprobes must be greater than 0",
);
expect(() =>
tbl.query().nearestTo(queryVec).fullTextSearch("dog").nprobes(0),
).toThrow("Invalid input, nprobes must be greater than 0");
expect(() => tbl.query().nearestTo(queryVec).minimumNprobes(0)).toThrow(
"Invalid input, minimum_nprobes must be greater than 0",
);
+2
View File
@@ -519,6 +519,8 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
/**
* Set the number of partitions to search (probe)
*
* The number of probes must be greater than 0.
*
* This argument is only used when the vector column has an IVF PQ index.
* If there is no index then this value is ignored.
*
+7 -2
View File
@@ -289,8 +289,13 @@ impl VectorQuery {
}
#[napi]
pub fn nprobes(&mut self, nprobe: u32) {
self.inner = self.inner.clone().nprobes(nprobe as usize);
pub fn nprobes(&mut self, nprobe: u32) -> napi::Result<()> {
self.inner = self
.inner
.clone()
.nprobes(nprobe as usize)
.default_error()?;
Ok(())
}
#[napi]
+8 -2
View File
@@ -1693,13 +1693,15 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceVectorQueryBuilder
The LanceQueryBuilder object.
"""
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
return self
@@ -2495,13 +2497,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
Parameters
----------
nprobes: int
The number of probes to use.
The number of probes to use. Must be greater than 0.
Returns
-------
LanceHybridQueryBuilder
The LanceHybridQueryBuilder object.
"""
if nprobes <= 0:
raise ValueError("Invalid input, nprobes must be greater than 0")
self._minimum_nprobes = nprobes
self._maximum_nprobes = nprobes
return self
@@ -3654,6 +3658,8 @@ class AsyncVectorQueryBase:
"""
Set the number of partitions to search (probe)
The number of probes must be greater than 0.
This argument is only used when the vector column has an IVF-based index.
If there is no index then this value is ignored.
+25
View File
@@ -975,6 +975,19 @@ def test_query_builder_with_filter(table):
assert all(np.array(rs[0]["vector"]) == [3, 4])
@pytest.mark.parametrize("query_type", ["vector", "hybrid"])
@pytest.mark.parametrize("nprobes", [0, -1])
def test_nprobes_nonpositive_sync(table, query_type, nprobes):
if query_type == "hybrid":
query = table.search(query_type="hybrid").vector([0, 0]).text("a")
else:
query = table.search([0, 0])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(nprobes)
def test_invalid_nprobes_sync(table):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
LanceVectorQueryBuilder(table, [0, 0], "vector").minimum_nprobes(0).to_list()
@@ -1006,6 +1019,18 @@ def test_multiple_nprobes_calls_works_sync(table):
).minimum_nprobes(20).to_list()
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_nprobes_zero_async(table_async: AsyncTable, hybrid):
query = table_async.query().nearest_to([0, 0])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
@pytest.mark.asyncio
async def test_invalid_nprobes_async(table_async: AsyncTable):
with pytest.raises(ValueError, match="minimum_nprobes must be greater than 0"):
+47
View File
@@ -1692,6 +1692,53 @@ def test_query_sync_maximal():
)
@pytest.mark.parametrize("hybrid", [False, True])
def test_query_sync_nprobes_zero(hybrid):
query_requests = []
def handler(body):
query_requests.append(body)
return pa.table({"id": []})
with query_test_table(handler) as table:
if hybrid:
query = table.search(query_type="hybrid").vector([1, 2, 3]).text("dog")
else:
query = table.search([1, 2, 3])
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
assert query_requests == []
@pytest.mark.asyncio
@pytest.mark.parametrize("hybrid", [False, True])
async def test_query_async_nprobes_zero(hybrid):
requests = []
def handler(request):
requests.append(request.path)
if request.path == "/v1/table/test/describe/":
send_json(request, {"version": 1, "schema": {"fields": []}})
else:
request.send_response(404)
request.end_headers()
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
query = table.query().nearest_to([1, 2, 3])
if hybrid:
query = query.nearest_to_text("dog")
with pytest.raises(
ValueError, match="^Invalid input, nprobes must be greater than 0$"
):
query.nprobes(0)
assert requests == ["/v1/table/test/describe/"]
def test_query_sync_nprobes():
def handler(body):
assert body == {
+5 -4
View File
@@ -995,8 +995,9 @@ impl VectorQuery {
self.inner = self.inner.clone().refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner = self.inner.clone().nprobes(nprobe as usize);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner = self.inner.clone().nprobes(nprobe as usize).infer_error()?;
Ok(())
}
pub fn minimum_nprobes(&mut self, minimum_nprobes: u32) -> PyResult<()> {
@@ -1188,8 +1189,8 @@ impl HybridQuery {
self.inner_vec.refine_factor(refine_factor);
}
pub fn nprobes(&mut self, nprobe: u32) {
self.inner_vec.nprobes(nprobe);
pub fn nprobes(&mut self, nprobe: u32) -> PyResult<()> {
self.inner_vec.nprobes(nprobe)
}
pub fn ef(&mut self, ef: u32) {
+1 -1
View File
@@ -142,7 +142,7 @@ async fn search_index(table: &Table) -> Result<()> {
// Override the default of 10 to get more rows
.limit(15)
// Override the default of 20 to search more partitions
.nprobes(30)
.nprobes(30)?
// Override the default of None to apply a refine step
.refine_factor(1)
.execute()
+35 -2
View File
@@ -1252,10 +1252,17 @@ impl VectorQuery {
/// This method sets both the minimum and maximum number of partitions to search.
/// For more fine-grained control see [`VectorQuery::minimum_nprobes`] and
/// [`VectorQuery::maximum_nprobes`].
pub fn nprobes(mut self, nprobes: usize) -> Self {
///
/// Returns an error if `nprobes` is not greater than 0.
pub fn nprobes(mut self, nprobes: usize) -> Result<Self> {
if nprobes == 0 {
return Err(Error::InvalidInput {
message: "nprobes must be greater than 0".to_string(),
});
}
self.request.minimum_nprobes = nprobes;
self.request.maximum_nprobes = Some(nprobes);
self
Ok(self)
}
/// Set the minimum number of partitions to search
@@ -2361,6 +2368,7 @@ mod tests {
.nearest_to(&[9.8, 8.7])
.unwrap()
.nprobes(1000)
.unwrap()
.postfilter()
.distance_type(DistanceType::Cosine)
.approx_mode(ApproxMode::Accurate)
@@ -2386,6 +2394,31 @@ mod tests {
assert_eq!(query.request.refine_factor, Some(999));
}
#[tokio::test]
async fn test_nprobes_validation() {
let conn = connect("memory://").execute().await.unwrap();
let table = conn
.create_table("my_table", make_test_batches())
.execute()
.await
.unwrap();
let query = table.query().nearest_to(&[0.1, 0.2]).unwrap();
assert!(matches!(
query.clone().nprobes(0),
Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0"
));
// Fixed probe counts replace both bounds, even when moving below or
// above the previous range.
let mut query = query;
for nprobes in [30, 1, 50] {
query = query.nprobes(nprobes).unwrap();
assert_eq!(query.request.minimum_nprobes, nprobes);
assert_eq!(query.request.maximum_nprobes, Some(nprobes));
}
}
#[test]
fn test_approx_mode_serde_parse_default_and_display() {
assert_eq!(ApproxMode::default(), ApproxMode::Normal);
+17
View File
@@ -6060,6 +6060,22 @@ mod tests {
assert_eq!(data[0].as_ref().unwrap(), &expected_data);
}
#[tokio::test]
async fn test_query_vector_nprobes_zero() {
let table = Table::new_with_handler::<&str>("my_table", |_| {
panic!("invalid nprobes must be rejected before sending a request")
});
let result = table
.query()
.nearest_to(vec![0.1, 0.2, 0.3])
.unwrap()
.nprobes(0);
assert!(matches!(
result,
Err(Error::InvalidInput { message }) if message == "nprobes must be greater than 0"
));
}
#[tokio::test]
async fn test_query_vector_all_params() {
let table = Table::new_with_handler("my_table", |request| {
@@ -6133,6 +6149,7 @@ mod tests {
.postfilter()
.distance_type(crate::DistanceType::Cosine)
.nprobes(12)
.unwrap()
.refine_factor(2)
.bypass_vector_index()
.execute()
+1
View File
@@ -1674,6 +1674,7 @@ impl Table {
/// .unwrap()
/// .refine_factor(5)
/// .nprobes(10)
/// .unwrap()
/// .execute()
/// .await
/// .unwrap();