From 0be3ae960eb39b43faad5685cd381c5126a305c5 Mon Sep 17 00:00:00 2001 From: Chi-Wei Lee Date: Sat, 3 Oct 2026 23:12:29 +0800 Subject: [PATCH 1/3] fix(python): support exist_ok in local create_table (#4421) Synchronous local Python `create_table` callers using `mode="exist_ok"` received `ValueError` instead of creating or opening a table. Related to [#2900](https://github.com/lancedb/lancedb/issues/2900). The local wrapper rejected this mode before the shared Rust parser could validate it. Removing that guard delegates mode validation to the parser, and the sync and async docs now describe `exist_ok`. Remote option gaps and the TypeScript table type gap remain outside this change. The added tests retain both calls from [@wjones127's comment](https://github.com/lancedb/lancedb/issues/2900#issuecomment-5964167206). This builds on the code @wjones127 posted there. Three regression tests fail before the change, pass after it, fail with the original `db.py` restored, then pass again on HEAD. Tested on Linux x86_64 with Python 3.14.7. From the Python directory, `uv run --no-sync --extra tests --extra dev python -m pytest python/tests/test_db.py -q` reports 53 passed. `uv run --no-sync --extra tests --extra dev python -m pytest python/tests -q -rs -m "not slow and not s3_test"` reports 1428 passed, 34 skipped, 68 deselected, and one failure: the existing `test_get_lsm_write_spec`, also failing on clean base (1425 passed, one failed). Ruff 0.9.9 `ruff format --check .` and `ruff check .` pass. Tests used worktree Python code with the released `lancedb==0.40.0b12` native wheel. No native source build was run. Co-authored-by: Will Jones --- python/python/lancedb/db.py | 10 +++++----- python/python/tests/test_db.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 5 deletions(-) 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/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 From 3e0a1a87806b049e64c20173a49a95f16b86aec7 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 21:29:47 -0700 Subject: [PATCH 2/3] 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` instead of `Self`; callers must handle the result, as with the existing minimum/maximum probe setters. Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- docs/src/js/classes/VectorQuery.md | 2 ++ nodejs/__test__/remote.test.ts | 32 ++++++++++++++++++ nodejs/__test__/table.test.ts | 6 ++++ nodejs/lancedb/query.ts | 2 ++ nodejs/src/query.rs | 9 +++-- python/python/lancedb/query.py | 10 ++++-- python/python/tests/test_query.py | 25 ++++++++++++++ python/python/tests/test_remote_db.py | 47 +++++++++++++++++++++++++++ python/src/query.rs | 9 ++--- rust/lancedb/examples/ivf_pq.rs | 2 +- rust/lancedb/src/query.rs | 37 +++++++++++++++++++-- rust/lancedb/src/remote/table.rs | 17 ++++++++++ rust/lancedb/src/table.rs | 1 + 13 files changed, 188 insertions(+), 11 deletions(-) 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(); From 7f5933594f42eeeec1ce73a7ce8e0cd0113cbd82 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Sat, 3 Oct 2026 21:30:13 -0700 Subject: [PATCH 3/3] fix(python): reject empty query vectors before column inference (#4415) Empty list and NumPy vector queries were validated too late: local search inferred a vector column first, producing misleading schema errors, while remote search indexed `vector[0]` and raised `IndexError`. Add a shared empty-vector check before column inference and query creation, and validate direct vector-builder inputs. Local and remote sync search now raise `ValueError: Query vector must not be empty` for these inputs. Validation: - All 26 new regression cases failed on the original code and pass with the fix, covering list/NumPy inputs, automatic/explicit vector search, inferred/explicit column names, one/two local vector columns, and remote search before query submission. - After merging current `main`, all 32 focused cases passed, covering these regressions and upstream async searches without a query. - The query, table, remote, utility, hybrid-query, and database suites passed in the rebuilt uv environment: 541 passed, 1 skipped, 1 slow test deselected. - `ruff check .`, `ruff format .`, and `cargo fmt --all` passed from the repository root. Fixes #4389 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/query.py | 6 +++++- python/python/lancedb/util.py | 11 +++++++++++ python/python/tests/test_query.py | 6 ++++++ python/python/tests/test_remote_db.py | 15 +++++++++++++++ python/python/tests/test_table.py | 18 ++++++++++++++++++ 5 files changed, 55 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index d135dd1d9..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 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_query.py b/python/python/tests/test_query.py index 20c0f3a1e..d3fedbcc8 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -2138,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 0dcfbe9a9..a9e411a3b 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 MockTextEmbeddingFunction from lancedb.query import AsyncQuery, ColumnOrdering from lancedb.remote import ClientConfig @@ -1592,6 +1593,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): 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