diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index f49963d60..a0a77fd7c 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -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. diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 5422945b6..ec259b79b 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -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) => { diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index eaf59211c..c17b0315a 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -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", ); diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 7c0dc9221..f370e1b31 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -519,6 +519,8 @@ export class VectorQuery extends StandardQueryBase { /** * 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. * diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index f542fe856..712569b2f 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -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] diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index db213ac71..d135dd1d9 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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. diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 6c59a9010..20c0f3a1e 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -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"): diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 47f1f866e..0dcfbe9a9 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -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 == { diff --git a/python/src/query.rs b/python/src/query.rs index 3a4737ee0..1079fb314 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -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) { diff --git a/rust/lancedb/examples/ivf_pq.rs b/rust/lancedb/examples/ivf_pq.rs index 79f674615..b457558ba 100644 --- a/rust/lancedb/examples/ivf_pq.rs +++ b/rust/lancedb/examples/ivf_pq.rs @@ -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() diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index c8b19f013..fffb703fb 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -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 { + 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); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 74f4b70ef..5756d7788 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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() diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index a95d8ddf3..ea6e512fd 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1674,6 +1674,7 @@ impl Table { /// .unwrap() /// .refine_factor(5) /// .nprobes(10) + /// .unwrap() /// .execute() /// .await /// .unwrap();