fix: reject zero vector search probes across SDKs

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.
 

nodejs/__test__/remote.test.ts

@@ -85,6 +85,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("rejects the external blob opt-in without blocking regular adds", async () => {
     let insertRequests = 0;
     const describeRequests: string[] = [];

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",
     );

nodejs/lancedb/query.ts

@@ -532,6 +532,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.
    *

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]

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.
 

python/python/tests/test_query.py

@@ -974,6 +974,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()
@@ -1005,6 +1018,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"):

python/python/tests/test_remote_db.py

@@ -1581,6 +1581,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 == {

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) {

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()

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<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);

rust/lancedb/src/remote/table.rs

@@ -5937,6 +5937,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| {
@@ -6010,6 +6026,7 @@ mod tests {
             .postfilter()
             .distance_type(crate::DistanceType::Cosine)
             .nprobes(12)
+            .unwrap()
             .refine_factor(2)
             .bypass_vector_index()
             .execute()

rust/lancedb/src/table.rs

@@ -1674,6 +1674,7 @@ impl Table {
     ///     .unwrap()
     ///     .refine_factor(5)
     ///     .nprobes(10)
+    ///     .unwrap()
     ///     .execute()
     ///     .await
     ///     .unwrap();