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 f9e50fe51..842bc52fd 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -146,6 +146,38 @@ 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/db.py b/python/python/lancedb/db.py index 330514484..1af3c1b4a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -409,9 +409,10 @@ class DBConnection(EnforceOverrides): - [LanceModel][lancedb.pydantic.LanceModel] mode: str; default "create" The mode to use when creating the table. - Can be either "create" or "overwrite". + Can be "create", "overwrite", or "exist_ok". By default, if the table already exists, an exception is raised. If you want to overwrite the table, use mode="overwrite". + To open an existing table without adding data, use mode="exist_ok". exist_ok: bool, default False If a table by the same name already exists, then raise an exception if exist_ok=False. If exist_ok=True, then open the existing table; @@ -1546,8 +1547,6 @@ class LanceDBConnection(DBConnection): """ if namespace_path is None: namespace_path = [] - if mode.lower() not in ["create", "overwrite"]: - raise ValueError("mode must be either 'create' or 'overwrite'") validate_table_name(name) tbl = LanceTable.create( @@ -2350,11 +2349,12 @@ class AsyncConnection(object): - pyarrow.Schema - [LanceModel][lancedb.pydantic.LanceModel] - mode: Literal["create", "overwrite"]; default "create" + mode: Literal["create", "overwrite", "exist_ok"]; default "create" The mode to use when creating the table. - Can be either "create" or "overwrite". + Can be "create", "overwrite", or "exist_ok". By default, if the table already exists, an exception is raised. If you want to overwrite the table, use mode="overwrite". + To open an existing table without adding data, use mode="exist_ok". exist_ok: bool, default False If a table by the same name already exists, then raise an exception if exist_ok=False. If exist_ok=True, then open the existing table; diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index db213ac71..b210c6497 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -40,7 +40,7 @@ from .rerankers.base import Reranker from .rerankers.rrf import RRFReranker from .rerankers.util import check_reranker_result from .schema import is_blob_like_field, schema_has_blob_field -from .util import flatten_columns +from .util import _validate_query_vector, flatten_columns from . import _wal_hybrid # WAL-PK-FUSION: delete. from ._blob import ( BLOB_MODE_TO_HANDLING, @@ -907,6 +907,9 @@ class LanceQueryBuilder(ABC): fast_search: bool Skip flat search of unindexed data. """ + if query_type != "fts": + _validate_query_vector(query) + if ordering_field_name is not None: import warnings @@ -1619,6 +1622,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder): str_query: Optional[str] = None, fast_search: bool = None, ): + _validate_query_vector(query) super().__init__(table) self._query = query self._distance_type = None @@ -1693,13 +1697,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 +2501,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 +3662,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/lancedb/util.py b/python/python/lancedb/util.py index 3cf5460a5..3a6a8a102 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -280,6 +280,14 @@ def infer_vector_column_dim(data_type: pa.DataType) -> Optional[int]: return None +def _validate_query_vector(query: Any) -> None: + """Reject empty vector inputs before vector-column inference or execution.""" + if (isinstance(query, list) and not query) or ( + isinstance(query, np.ndarray) and query.size == 0 + ): + raise ValueError("Query vector must not be empty") + + def _query_vector_dim(query: Optional[Any]) -> Optional[int]: if query is None: return None @@ -301,6 +309,9 @@ def infer_vector_column_name( query: Optional[Any], # inferred later in query builder vector_column_name: Optional[str], ): + if query_type != "fts": + _validate_query_vector(query) + if vector_column_name is not None: return vector_column_name diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index df5be5a6d..ebd86bdc3 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -351,6 +351,40 @@ async def test_create_table_from_iterator_async(mem_db_async: lancedb.AsyncConne assert await table.count_rows() == 10 +def test_create_exist_ok_mode(tmp_db: lancedb.DBConnection): + db = tmp_db + data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}] + table = db.create_table("t1", data, mode="exist_ok") + assert table.count_rows() == 2 + original = table.to_arrow() + + reopened = db.create_table("t1", [{"vector": [9.0, 8.0]}], mode="exist_ok") + assert reopened.to_arrow().equals(original) + + +def test_create_empty_exist_ok_mode(tmp_db: lancedb.DBConnection): + schema = pa.schema([pa.field("id", pa.int64())]) + table = tmp_db.create_table("empty", schema=schema, mode="exist_ok") + assert table.schema == schema + assert table.count_rows() == 0 + table.add([{"id": 7}]) + + reopened = tmp_db.create_table("empty", schema=schema, mode="exist_ok") + assert reopened.to_arrow().to_pylist() == [{"id": 7}] + + bad_schema = pa.schema([pa.field("other", pa.int64())]) + with pytest.raises(ValueError): + tmp_db.create_table("empty", schema=bad_schema, mode="exist_ok") + + +def test_create_table_invalid_mode(tmp_db: lancedb.DBConnection): + db = tmp_db + data = [{"vector": [1.1, 1.2]}, {"vector": [0.2, 1.8]}] + with pytest.raises(ValueError, match="Invalid mode bogus"): + db.create_table("tb", data, mode="bogus") + assert "tb" not in db + + def test_create_exist_ok(tmp_db: lancedb.DBConnection): from conftest import pandas_string_type diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index 6c59a9010..d3fedbcc8 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"): @@ -2113,6 +2138,12 @@ def test_ensure_vector_query_nested_empty_list(): ensure_vector_query([[]]) +@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)]) +def test_vector_query_builder_empty_vector(table, query): + with pytest.raises(ValueError, match="^Query vector must not be empty$"): + LanceVectorQueryBuilder(table, query, "vector") + + def test_fast_search(tmp_path): db = lancedb.connect(tmp_path) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 03595285f..a75982991 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -16,6 +16,7 @@ import uuid from packaging.version import Version import lancedb +import numpy as np from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry from lancedb.query import AsyncQuery, ColumnOrdering @@ -1750,6 +1751,20 @@ def test_query_sync_empty_query(): assert data == expected +@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)]) +@pytest.mark.parametrize("vector_column_name", [None, "vector"]) +@pytest.mark.parametrize("query_type", ["auto", "vector"]) +def test_query_sync_empty_vector(query, vector_column_name, query_type): + def handler(body): + pytest.fail("An empty query vector must be rejected before sending a query") + + with query_test_table(handler) as table: + with pytest.raises(ValueError, match="^Query vector must not be empty$"): + table.search( + query, vector_column_name=vector_column_name, query_type=query_type + ).limit(3).to_arrow() + + @pytest.mark.asyncio @pytest.mark.parametrize("search_kwargs", [{}, {"query": None}]) async def test_async_search_without_query(search_kwargs): @@ -1850,6 +1865,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/python/tests/test_table.py b/python/python/tests/test_table.py index 0c12c5c20..e5ee3c3a2 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3777,6 +3777,24 @@ def test_empty_query(mem_db: DBConnection): assert df.num_rows == 42 +@pytest.mark.parametrize("query", [[], np.array([], dtype=np.float32)]) +@pytest.mark.parametrize("vector_column_name", [None, "vector"]) +@pytest.mark.parametrize("query_type", ["auto", "vector"]) +@pytest.mark.parametrize("multiple_vector_columns", [False, True]) +def test_search_empty_vector( + mem_db, query, vector_column_name, query_type, multiple_vector_columns +): + fields = [pa.field("vector", pa.list_(pa.float32(), 8))] + if multiple_vector_columns: + fields.append(pa.field("vec2", pa.list_(pa.float32(), 4))) + table = mem_db.create_table("empty_vector_query", schema=pa.schema(fields)) + + with pytest.raises(ValueError, match="^Query vector must not be empty$"): + table.search( + query, vector_column_name=vector_column_name, query_type=query_type + ).limit(3).to_arrow() + + def test_search_with_schema_inf_single_vector(mem_db: DBConnection): class MyTable(LanceModel): text: str 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 8f28ea3bd..fc5ccff7f 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -6136,6 +6136,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| { @@ -6209,6 +6225,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();