mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-12 23:02:28 +00:00
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:
@@ -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.
|
||||
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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",
|
||||
);
|
||||
|
||||
@@ -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
@@ -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]
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -1674,6 +1674,7 @@ impl Table {
|
||||
/// .unwrap()
|
||||
/// .refine_factor(5)
|
||||
/// .nprobes(10)
|
||||
/// .unwrap()
|
||||
/// .execute()
|
||||
/// .await
|
||||
/// .unwrap();
|
||||
|
||||
Reference in New Issue
Block a user