From ba7b7f15c2a576062a480e6863abfa092b6e0e58 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:07:18 +0800 Subject: [PATCH 01/68] fix(node): reject musl binaries with unresolved AVX-512 symbol (#3791) ## Summary - inspect the assembled x64 musl addon for the unresolved AVX-512 distance-table symbol before npm publication - fail the publish lifecycle with an actionable error if the broken symbol is present - add regression coverage using the dynamic-symbol output from the affected 0.27.x artifact ## Root cause LanceDB 0.27.x could enable the Rust AVX-512 distance-table callsite when the separate C kernel build did not emit its object, leaving `sum_4bit_dist_table_32bytes_batch_avx512` undefined in the musl addon. The current Lance dependency includes the upstream split-kernel cfg fix, and the published 0.33.0 musl artifact no longer has the undefined symbol, but the LanceDB publish path did not verify that property before releasing an addon. This adds a release guard around the corrected implementation so the same invalid binary cannot be published again. ## Validation - `pnpm lint` - `pnpm run build` - `pnpm test --runInBand` (680 passed, 5 skipped) - `pnpm run docs` - checked published artifacts: 0.27.1 is rejected for the unresolved symbol; 0.33.0 is accepted Fixes #3200 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- nodejs/.npmignore | 1 + nodejs/__test__/native_symbols.test.ts | 47 ++++++++++++++++ nodejs/package.json | 3 +- nodejs/scripts/check-native-symbols.js | 77 ++++++++++++++++++++++++++ 4 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 nodejs/__test__/native_symbols.test.ts create mode 100644 nodejs/scripts/check-native-symbols.js diff --git a/nodejs/.npmignore b/nodejs/.npmignore index 41f6643a6..128c0d36d 100644 --- a/nodejs/.npmignore +++ b/nodejs/.npmignore @@ -15,6 +15,7 @@ renovate.json src lancedb examples +scripts nodejs-artifacts Cargo.toml biome.json diff --git a/nodejs/__test__/native_symbols.test.ts b/nodejs/__test__/native_symbols.test.ts new file mode 100644 index 000000000..5445867f6 --- /dev/null +++ b/nodejs/__test__/native_symbols.test.ts @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const { + checkNativeBinary, + findForbiddenUndefinedSymbols, +} = require("../scripts/check-native-symbols.js"); + +test("detects the unresolved AVX-512 symbol from broken musl binaries", () => { + const output = [ + " U napi_create_function", + " U sum_4bit_dist_table_32bytes_batch_avx512", + " U strlen", + ].join("\n"); + + expect(findForbiddenUndefinedSymbols(output)).toEqual([ + "sum_4bit_dist_table_32bytes_batch_avx512", + ]); +}); + +test("accepts native binaries without unresolved internal Lance symbols", () => { + const runNm = jest.fn(() => ({ + status: 0, + stdout: + " U napi_create_function\n U strlen\n", + stderr: "", + })); + + expect(() => checkNativeBinary("lancedb.node", runNm)).not.toThrow(); + expect(runNm).toHaveBeenCalledWith( + "nm", + ["-D", "--undefined-only", "lancedb.node"], + { encoding: "utf8" }, + ); +}); + +test("rejects native binaries with the unresolved AVX-512 symbol", () => { + const runNm = jest.fn(() => ({ + status: 0, + stdout: " U sum_4bit_dist_table_32bytes_batch_avx512\n", + stderr: "", + })); + + expect(() => checkNativeBinary("lancedb.node", runNm)).toThrow( + "lancedb.node contains unresolved internal Lance symbols: sum_4bit_dist_table_32bytes_batch_avx512", + ); +}); diff --git a/nodejs/package.json b/nodejs/package.json index 9bade7200..21f903235 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -90,7 +90,8 @@ "postdocs": "node typedoc_post_process.js", "lint": "biome check . && biome format .", "lint-fix": "biome check --write . && biome format --write .", - "prepublishOnly": "napi prepublish -t npm", + "check:native-symbols": "node scripts/check-native-symbols.js", + "prepublishOnly": "pnpm check:native-symbols && napi prepublish -t npm", "test": "jest --verbose", "integration": "S3_TEST=1 pnpm test", "universal": "napi universalize", diff --git a/nodejs/scripts/check-native-symbols.js b/nodejs/scripts/check-native-symbols.js new file mode 100644 index 000000000..2f6dbf9b8 --- /dev/null +++ b/nodejs/scripts/check-native-symbols.js @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const { spawnSync } = require("node:child_process"); +const { existsSync } = require("node:fs"); +const path = require("node:path"); + +const FORBIDDEN_UNDEFINED_SYMBOLS = new Set([ + "sum_4bit_dist_table_32bytes_batch_avx512", +]); + +function findForbiddenUndefinedSymbols(output) { + const found = new Set(); + + for (const line of output.split(/\r?\n/)) { + const columns = line.trim().split(/\s+/); + const symbol = columns.at(-1)?.split("@")[0]; + if (symbol && FORBIDDEN_UNDEFINED_SYMBOLS.has(symbol)) { + found.add(symbol); + } + } + + return [...found].sort(); +} + +function checkNativeBinary(binaryPath, runNm = spawnSync) { + const result = runNm("nm", ["-D", "--undefined-only", binaryPath], { + encoding: "utf8", + }); + + if (result.error) { + throw new Error(`Unable to inspect ${binaryPath}: ${result.error.message}`); + } + if (result.status !== 0) { + throw new Error( + `Unable to inspect ${binaryPath}: nm exited with status ${result.status}\n${result.stderr}`, + ); + } + + const forbidden = findForbiddenUndefinedSymbols(result.stdout); + if (forbidden.length > 0) { + throw new Error( + `${binaryPath} contains unresolved internal Lance symbols: ${forbidden.join( + ", ", + )}`, + ); + } +} + +function main() { + const binaryPath = path.resolve( + __dirname, + "..", + "npm", + "linux-x64-musl", + "lancedb.linux-x64-musl.node", + ); + + if (!existsSync(binaryPath)) { + throw new Error( + `Missing ${binaryPath}; assemble the native artifacts before publishing`, + ); + } + + checkNativeBinary(binaryPath); +} + +if (require.main === module) { + try { + main(); + } catch (error) { + console.error(error instanceof Error ? error.message : error); + process.exitCode = 1; + } +} + +module.exports = { checkNativeBinary, findForbiddenUndefinedSymbols }; From a98667749d923512c1f26b65167da5aea0c6669c Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:09:01 +0800 Subject: [PATCH 02/68] fix(python): verify Lance-created IVF_RQ compatibility (#3793) ## Summary - add regression coverage for an IVF_RQ index created through the Lance Python API and consumed through LanceDB - verify a fresh LanceDB handle can list and query the externally-created index ## Root cause Lance 4.0.0-beta.5 could serialize `rotate_mat_position` as null, while the Lance 2.0.0 reader used by LanceDB 0.29.2 required a `u32`. Upstream Lance fixed the incompatible IVF_RQ metadata handling and index versioning, and the current LanceDB dependency includes that repair. LanceDB was missing coverage for the reported mixed-API workflow, allowing the compatibility regression to go undetected. ## Validation - `cd python && uv run --extra tests pytest python/tests/test_index.py -q` (24 passed) - `uv run --project python --extra dev ruff format .` - `uv run --project python --extra dev ruff check .` - `uv run --project python --extra dev ruff format --check python/python/tests/test_index.py` - `uv run --project python --extra dev ruff check python/python/tests/test_index.py` Fixes #3092 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/tests/test_index.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index fe6ebe87a..0252a09ad 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -430,6 +430,28 @@ async def test_create_ivfrq_index(some_table: AsyncTable): assert indices[0].name == "vector_idx" +@pytest.mark.asyncio +async def test_read_ivfrq_index_created_with_lance(db_async, some_table: AsyncTable): + pytest.importorskip("lance") + dataset = await some_table.to_lance() + dataset.create_index( + "vector", + index_type="IVF_RQ", + metric="cosine", + num_partitions=1, + ) + + # Open a fresh LanceDB handle so the index is discovered from disk. + table = await db_async.open_table("some_table") + indices = await table.list_indices() + assert len(indices) == 1 + assert indices[0].index_type == "IvfRq" + assert indices[0].columns == ["vector"] + + results = await table.query().nearest_to(list(range(DIM))).limit(8).to_arrow() + assert len(results) == 8 + + @pytest.mark.asyncio async def test_create_hnswpq_index(some_table: AsyncTable): await some_table.create_index("vector", config=HnswPq(num_partitions=10)) From 102518ebda39513aa0bf337be96706df5ad558a0 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:09:16 +0800 Subject: [PATCH 03/68] fix(python): preserve nested fields in filtered scans (#3816) ## Summary - add a Python regression for filtered scans over deeply nested list/struct values - verify the public query path returns every nested sibling field without an Arrow schema mismatch ## Root cause Older Lance scanner code split filter projections into partial nested arrays and did not recursively merge List children, so a sibling field could be omitted from the returned array while remaining in the declared schema. That merge was fixed upstream in lance-format/lance#4838, and the Lance revision currently pinned by LanceDB contains that fix. This test makes the repair durable at the LanceDB Python API boundary. ## Validation - cd python && uv run --extra tests pytest python/tests/test_query.py::test_filter_preserves_deeply_nested_struct_fields -q - cd python && uv run --extra tests pytest python/tests/test_query.py -q - uv run --project python --extra tests --extra dev ruff check . - uv run --project python --extra tests --extra dev ruff format --check python/python/tests/test_query.py Fixes #2217 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/tests/test_query.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/python/python/tests/test_query.py b/python/python/tests/test_query.py index c8d0ad159..58ea7deb5 100644 --- a/python/python/tests/test_query.py +++ b/python/python/tests/test_query.py @@ -833,6 +833,28 @@ def test_with_row_id(table: lancedb.table.Table): assert rs["_rowid"].to_pylist() == [0, 1] +def test_filter_preserves_deeply_nested_struct_fields(tmp_db): + # Regression test for https://github.com/lancedb/lancedb/issues/2217. + data = [ + { + "item": "test", + "sub_obj": { + "second_sub_obj": [ + { + "third_sub_obj": {"int_val": 1}, + "str_val": "test", + } + ] + }, + } + ] + table = tmp_db.create_table("deeply_nested_struct", data=data) + + result = table.search().where("item = 'test'").limit(1).to_arrow() + + assert result.to_pylist() == data + + def test_blob_v2_query_omits_auto_row_id(tmp_db): table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid") From 40af9f350e785d762cc887e029718fc9904bbe0b Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:09:32 +0800 Subject: [PATCH 04/68] fix(python): avoid deadlock for query-backed iterators (#3820) ## Summary - move synchronous iterator preparation off LanceDB's singleton background event loop - keep the loop available when an input iterator executes a synchronous table query - cover query-backed iterators for both `Table.add` and `create_table` ## Root cause The async implementations prepared input data on the singleton background loop and peeked at the first item of iterators there. If that iterator called a synchronous LanceDB query, `LOOP.run()` scheduled the query on the same loop and blocked waiting for it, so neither operation could progress. ## Validation - `ruff format` on the changed Python files - `ruff check .` - `cd python && uv run --frozen --extra tests pytest python/tests/test_table.py -q` (138 passed) - the broader non-slow suite was attempted but hit a native segmentation fault in the external DataFusion `register_table_provider` path; an isolated control using the pre-change inline preparation behavior reproduced the same unrelated crash Fixes #2107 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/db.py | 13 +++++++++++-- python/python/lancedb/table.py | 16 ++++++++++++---- python/python/tests/test_table.py | 24 ++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 6 deletions(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 5def89b01..e1e5d270a 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -5,6 +5,7 @@ from __future__ import annotations from abc import abstractmethod +import asyncio from datetime import timedelta from pathlib import Path import sys @@ -2652,8 +2653,16 @@ class AsyncConnection(object): if fill_value is None: fill_value = 0.0 - data, schema = sanitize_create_table( - data, schema, metadata, on_bad_vectors, fill_value + # Input preparation may advance a user-provided iterator. Keep that work + # off the background event loop so an iterator can use the synchronous + # LanceDB API without blocking the loop that API needs to make progress. + data, schema = await asyncio.to_thread( + sanitize_create_table, + data, + schema, + metadata, + on_bad_vectors, + fill_value, ) validate_schema(schema) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index d5872e6f0..006ef8cb9 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -6213,8 +6213,12 @@ class AsyncTable: if mode == "overwrite": # For overwrite, apply the same preprocessing as create_table # so vector columns are inferred as FixedSizeList. - data, _ = sanitize_create_table( - data, None, on_bad_vectors=on_bad_vectors, fill_value=fill_value + data, _ = await asyncio.to_thread( + sanitize_create_table, + data, + None, + on_bad_vectors=on_bad_vectors, + fill_value=fill_value, ) if blob_v2_column_paths(schema): # Plain bytes cannot be inferred as blob v2. Keep matching blob @@ -6224,7 +6228,8 @@ class AsyncTable: elif on_bad_vectors != "error" or ( schema.metadata is not None and b"embedding_functions" in schema.metadata ): - data = _sanitize_data( + data = await asyncio.to_thread( + _sanitize_data, data, schema, metadata=schema.metadata, @@ -6233,7 +6238,10 @@ class AsyncTable: allow_subschema=True, ) _register_optional_converters() - data = to_scannable(data) + # Converting an iterator peeks at its first item. A synchronous query in + # that iterator schedules work on LOOP, so peeking on LOOP's own thread + # would deadlock waiting for itself. + data = await asyncio.to_thread(to_scannable, data) progress, owns = _normalize_progress(progress) try: return await self._inner.add( diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 755e94d7f..095f0b911 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -518,6 +518,30 @@ def test_add(mem_db: DBConnection): _add(table, schema) +def test_add_from_iterator_that_queries_table(mem_db: DBConnection): + source = mem_db.create_table("source", data=pa.table({"id": range(16)})) + target = mem_db.create_table("target", schema=source.schema) + + def batches(): + for _ in range(5): + yield source.search().limit(10).to_arrow() + + target.add(batches()) + + assert target.count_rows() == 50 + + +def test_create_table_from_iterator_that_queries_table(mem_db: DBConnection): + source = mem_db.create_table("source", data=pa.table({"id": range(16)})) + + def batches(): + yield source.search().limit(10).to_arrow() + + target = mem_db.create_table("target", data=batches()) + + assert target.count_rows() == 10 + + def test_add_releases_arrow_buffers_without_gc(mem_db: DBConnection): """Regression test for https://github.com/lancedb/lancedb/issues/2512.""" schema = pa.schema([pa.field("x", pa.int64())]) From a1230cf9a490c1768cbec2709edf772543d27331 Mon Sep 17 00:00:00 2001 From: Li Zinuo Date: Fri, 23 Oct 2026 01:09:46 +0800 Subject: [PATCH 05/68] fix: preserve unbounded plain queries in namespace pushdown (#3927) ## What Preserve unbounded plain-query semantics in the Rust Namespace `QueryTable` adapter. Plain queries without an explicit limit are now encoded with `i32::MAX`. Explicit limits remain unchanged, while vector and FTS queries continue to use the default top-k of 10. ## Root cause The Rust Namespace adapter converted a missing limit using: ```rust q.limit.unwrap_or(10) ``` This turned an unbounded plain query into an explicit `k=10`, silently truncating the results returned by Namespace pushdown. ## Tests Added regression coverage for unbounded and explicitly limited plain queries, while verifying that vector and FTS queries retain their default limit. ```shell cargo fmt --all -- --check cargo test --quiet --features remote -p lancedb \ test_convert_to_namespace_query --lib ``` close #3926. --- rust/lancedb/src/table/query.rs | 56 ++++++++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index cf70c6626..f625a4aff 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -671,6 +671,12 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { }); } + let k = match q.limit { + Some(limit) => limit as i32, + None if q.full_text_search.is_some() => DEFAULT_TOP_K as i32, + None => i32::MAX, + }; + let filter = q.filter.as_ref().map(filter_to_sql).transpose()?; let columns = match &q.select { @@ -723,7 +729,7 @@ fn convert_to_namespace_query(query: &AnyQuery) -> Result { Ok(NsQueryTableRequest { id: None, // Will be set by caller vector, - k: q.limit.unwrap_or(10) as i32, + k, filter, columns, prefilter: Some(q.prefilter), @@ -1073,6 +1079,54 @@ mod tests { assert!(ns_request.vector.single_vector.as_ref().unwrap().is_empty()); } + #[test] + fn test_convert_to_namespace_query_plain_query_without_limit() { + let any_query = AnyQuery::Query(QueryRequest::default()); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, i32::MAX); + } + + #[test] + fn test_convert_to_namespace_query_plain_query_with_explicit_limit() { + let any_query = AnyQuery::Query(QueryRequest { + limit: Some(10), + ..Default::default() + }); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, 10); + } + + #[test] + fn test_convert_to_namespace_query_vector_without_limit() { + let query_vector = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])); + let any_query = AnyQuery::VectorQuery(VectorQueryRequest { + query_vector: vec![query_vector as ArrayRef], + ..Default::default() + }); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, 10); + } + + #[test] + fn test_convert_to_namespace_query_fts_without_limit() { + let any_query = AnyQuery::Query(QueryRequest { + full_text_search: Some(lance_index::scalar::FullTextSearchQuery::new( + "hello world".to_string(), + )), + ..Default::default() + }); + + let ns_request = convert_to_namespace_query(&any_query).unwrap(); + + assert_eq!(ns_request.k, 10); + } + #[tokio::test] async fn test_execute_query_local_routing() { use crate::connect; From 548a1903216075756ae7e6e6eecb671a78f38dd8 Mon Sep 17 00:00:00 2001 From: primorLee Date: Fri, 23 Oct 2026 01:10:00 +0800 Subject: [PATCH 06/68] fix(python): refresh zero-lifetime OAuth tokens (#3971) ## Summary - treat an OAuth `expires_in` value of zero as an explicit immediate expiration - refresh zero-lifetime tokens instead of caching them indefinitely - add a regression test covering two consecutive header requests ## Tests - `pytest python/tests/test_header_provider.py -q` (16 passed) - `ruff format --check python/lancedb/remote/header.py python/tests/test_header_provider.py` - `ruff check python/lancedb/remote/header.py python/tests/test_header_provider.py` --- python/python/lancedb/remote/header.py | 2 +- python/python/tests/test_header_provider.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/remote/header.py b/python/python/lancedb/remote/header.py index 206b7a7cb..5f5bd12d8 100644 --- a/python/python/lancedb/remote/header.py +++ b/python/python/lancedb/remote/header.py @@ -141,7 +141,7 @@ class OAuthProvider(HeaderProvider): # Set expiration if provided expires_in = token_data.get("expires_in") - if expires_in: + if expires_in is not None: self._token_expires_at = time.time() + expires_in else: # Token doesn't expire or expiration unknown diff --git a/python/python/tests/test_header_provider.py b/python/python/tests/test_header_provider.py index 187b0a50a..e50ff3270 100644 --- a/python/python/tests/test_header_provider.py +++ b/python/python/tests/test_header_provider.py @@ -114,6 +114,24 @@ class TestOAuthProvider: "x-lancedb-credential-type": "oidc", } + def test_zero_lifetime_token_is_refreshed(self): + """A token that expires immediately must not be cached indefinitely.""" + call_count = 0 + + def fetcher(): + nonlocal call_count + call_count += 1 + return { + "access_token": f"token{call_count}", + "expires_in": 0, + } + + provider = OAuthProvider(fetcher, refresh_buffer_seconds=0) + + assert provider.get_headers() == {"Authorization": "Bearer token1"} + assert provider.get_headers() == {"Authorization": "Bearer token2"} + assert call_count == 2 + def test_missing_access_token(self): """Test error handling when access_token is missing.""" From 6d92307fc34b81adc0c040da367c6f9b64f0aeba Mon Sep 17 00:00:00 2001 From: Zhewen Tan <1090179959@qq.com> Date: Fri, 23 Oct 2026 01:10:15 +0800 Subject: [PATCH 07/68] fix(python): accept tuple and Arrow vector queries (#3972) ## Summary Sync `Table.search()` declares tuple, `pyarrow.Array`, and `pyarrow.ChunkedArray` query vectors as supported, and the async query path already accepts them. The sync query builder only classified lists and NumPy arrays as vectors, however, so these inputs were routed through embedding/FTS resolution and then rejected. Arrow and tuple inputs also did not contribute a query dimension when selecting among multiple vector columns. This change: - recognizes the supported tuple and Arrow containers as sync vector queries for both `auto` and explicit `vector` query types; - converts them to the existing internal `float32` NumPy representation at the sync builder boundary; - infers dimensions from tuple, primitive Arrow, and list-valued Arrow queries so schemas with differently sized vector columns select the matching column. The regression test uses two vector columns with dimensions 2 and 3. It verifies tuple, Arrow array, and chunked Arrow queries through both sync query modes and confirms the dimension-2 column is selected. ## Test plan - `pytest python/tests/test_table.py::test_sync_search_accepts_supported_vector_containers -q` (6 passed) - `pytest python/tests/test_util.py python/tests/test_query.py python/tests/test_table.py -q -m "not slow and not s3_test"` (308 passed, 1 skipped, 1 deselected) - `pytest python/tests -q -m "not slow and not s3_test" --ignore=python/tests/docs` (981 passed, 25 skipped, 65 deselected) - `ruff format --check python` - `ruff check python` - `cargo fmt --all --check` Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 9 +++++--- python/python/lancedb/util.py | 11 +++++++++- python/python/tests/test_table.py | 36 +++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index c427820e8..53f56a2c0 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -82,6 +82,7 @@ if TYPE_CHECKING: T = TypeVar("T", bound="LanceModel") AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"] +_SYNC_VECTOR_QUERY_TYPES = (list, tuple, np.ndarray, pa.Array, pa.ChunkedArray) # Number of rows a hybrid query returns when no limit was set on it. This # mirrors the default the Rust query builder applies to its sub-queries. @@ -1024,10 +1025,12 @@ class LanceQueryBuilder(ABC): fast_search=fast_search, ) - if isinstance(query, list): + if isinstance(query, (list, tuple)): query = np.array(query, dtype=np.float32) elif isinstance(query, np.ndarray): query = query.astype(np.float32) + elif isinstance(query, (pa.Array, pa.ChunkedArray)): + query = np.array(query.to_pylist(), dtype=np.float32) else: raise TypeError(f"Unsupported query type: {type(query)}") @@ -1049,7 +1052,7 @@ class LanceQueryBuilder(ABC): query = cls._query_to_vector(table, query, vector_column_name) return query, query_type elif query_type == "auto": - if isinstance(query, (list, np.ndarray)): + if isinstance(query, _SYNC_VECTOR_QUERY_TYPES): return query, "vector" else: conf = table.embedding_functions.get(vector_column_name) @@ -1065,7 +1068,7 @@ class LanceQueryBuilder(ABC): @classmethod def _query_to_vector(cls, table, query, vector_column_name): - if isinstance(query, (list, np.ndarray)): + if isinstance(query, _SYNC_VECTOR_QUERY_TYPES): return query conf = table.embedding_functions.get(vector_column_name) if conf is not None: diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index d4533887b..3babc3872 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -309,7 +309,16 @@ def _query_vector_dim(query: Optional[Any]) -> Optional[int]: if query.ndim == 0: return None return query.shape[-1] - if isinstance(query, list) and query: + if isinstance(query, (pa.Array, pa.ChunkedArray)): + if len(query) == 0: + return None + if pa.types.is_fixed_size_list(query.type): + return query.type.list_size + if pa.types.is_list(query.type) or pa.types.is_large_list(query.type): + first = query[0].as_py() + return len(first) if first is not None else None + return len(query) + if isinstance(query, (list, tuple)) and query: first = query[0] if isinstance(first, (list, tuple, np.ndarray)): return len(first) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 095f0b911..7a4f5981c 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -3697,6 +3697,42 @@ def test_multiple_vector_columns(mem_db: DBConnection): assert result1["text"][0] != result2["text"][0] +@pytest.mark.parametrize( + "query", + [ + (0.0, 0.0), + pa.array([0.0, 0.0]), + pa.chunked_array([[0.0], [0.0]]), + ], +) +@pytest.mark.parametrize("query_type", ["auto", "vector"]) +def test_sync_search_accepts_supported_vector_containers( + mem_db: DBConnection, query, query_type +): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("small", pa.list_(pa.float32(), 2)), + pa.field("large", pa.list_(pa.float32(), 3)), + ] + ) + table = mem_db.create_table( + "supported_vector_containers", + data=pa.table( + { + "id": [1, 2], + "small": [[0.0, 0.0], [1.0, 1.0]], + "large": [[1.0, 1.0, 1.0], [0.0, 0.0, 0.0]], + }, + schema=schema, + ), + ) + + result = table.search(query, query_type=query_type).limit(1).to_arrow() + + assert result["id"].to_pylist() == [1] + + def test_create_scalar_index(mem_db: DBConnection): vec_array = pa.array( [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]], pa.list_(pa.float32(), 2) From ae97d43f4bc3c1c01950c3a0d8369f21ceba014d Mon Sep 17 00:00:00 2001 From: Vinayak Vathare <137673662+VathareVinayak@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:40:30 +0530 Subject: [PATCH 08/68] fix(python): preserve valid embeddings when an OpenAI batch is rejected (#4005) Retries individual inputs one at a time when the OpenAI API rejects a batch with BadRequestError, instead of discarding the whole batch. Valid inputs now still get embeddings; only the actually-bad input resolves to None, and the log names the exact failing input. Fixes #1677 --- python/python/lancedb/embeddings/openai.py | 25 ++++++++---- python/python/tests/test_embeddings.py | 45 ++++++++++++++++++++++ 2 files changed, 63 insertions(+), 7 deletions(-) diff --git a/python/python/lancedb/embeddings/openai.py b/python/python/lancedb/embeddings/openai.py index 1edfaa2ca..d4ecbdf3f 100644 --- a/python/python/lancedb/embeddings/openai.py +++ b/python/python/lancedb/embeddings/openai.py @@ -101,24 +101,35 @@ class OpenAIEmbeddings(TextEmbeddingFunction): valid_indices.append(idx) # TODO retry, rate limit, token limit - try: + def embed_batch(texts_batch, indices_batch): kwargs = { - "input": valid_texts, + "input": texts_batch, "model": self.name, } if self.name != "text-embedding-ada-002": kwargs["dimensions"] = self.dim rs = self._openai_client.embeddings.create(**kwargs) - valid_embeddings = { - idx: v.embedding for v, idx in zip(rs.data, valid_indices) - } + return {idx: v.embedding for v, idx in zip(rs.data, indices_batch)} + + def embed_one_or_log(idx, text): + # Retry a single rejected input + try: + return embed_batch([text], [idx]) + except openai.BadRequestError: + logging.exception("Bad request for input at index %d: %r", idx, text) + return {} + + try: + valid_embeddings = embed_batch(valid_texts, valid_indices) except openai.AuthenticationError: logging.error("Authentication failed: Invalid API key provided") raise except openai.BadRequestError: - logging.exception("Bad request: %s", texts) - return [None] * len(texts) + # Retry each input individually + valid_embeddings = {} + for idx, text in zip(valid_indices, valid_texts): + valid_embeddings.update(embed_one_or_log(idx, text)) except Exception: logging.exception("OpenAI embeddings error") raise diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 9434033e1..c42172f77 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -596,6 +596,51 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +def test_openai_bad_request_partial_batch_recovery(): + # Regression test for issue #1677 + + # Fake exception types, avoids requiring the optional openai/httpx deps + class FakeBadRequestError(Exception): + pass + + class FakeAuthenticationError(Exception): + pass + + fake_openai_module = MagicMock() + fake_openai_module.BadRequestError = FakeBadRequestError + fake_openai_module.AuthenticationError = FakeAuthenticationError + + registry = get_registry() + model = registry.get("openai").create(name="text-embedding-ada-002") + + def fake_create(input, model=None, **kwargs): + # Full batch always rejected + if len(input) > 1: + raise FakeBadRequestError("bad request") + if input[0] == "bad text": + raise FakeBadRequestError("bad request") + embedding = MagicMock() + embedding.embedding = [0.1, 0.2, 0.3] + result = MagicMock() + result.data = [embedding] + return result + + mock_client = MagicMock() + mock_client.embeddings.create.side_effect = fake_create + + with ( + patch( + "lancedb.embeddings.openai.attempt_import_or_raise", + return_value=fake_openai_module, + ), + patch.object(type(model), "_openai_client", new_callable=lambda: mock_client), + ): + embeddings = model.generate_embeddings(["good text", "bad text"]) + + assert embeddings[0] == [0.1, 0.2, 0.3] + assert embeddings[1] is None + + def test_ollama_embeddings_pickle(): """OllamaEmbeddings must pickle even after the cached client is created.""" registry = get_registry() From 0c3a1c92c862d1f4d251d55eacafabe225277722 Mon Sep 17 00:00:00 2001 From: Madan Kumar Date: Thu, 22 Oct 2026 22:40:45 +0530 Subject: [PATCH 09/68] fix(python): don't crash inferring schema for an all-null vector column (#4034) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #4033 `create_table(data=...)` without an explicit schema raised `IndexError: index out of bounds` when a `vector`/`embedding`-named list column was all-null. `_modal_list_size` infers the dimension with `pc.mode(pc.list_value_length(arr))[0]`, and pyarrow's `mode` skips nulls, so an all-null column produces a zero-length result and the `[0]` subscript blows up. The sibling `_infer_vector_dim` already guards the identical `pc.mode(...)` with a `len == 0` check and returns `None`; `_modal_list_size` just didn't. This makes `_modal_list_size` return `Optional[int]` and return `None` when every value is null, and the call site then leaves the column as a variable-length list instead of inventing a fixed size — the dimension is genuinely unknowable when there are no non-null rows. Added a regression test that creates a table from an all-null `list` vector column and asserts it succeeds, the column stays a `list`, and the rows are queryable. It fails on `main` with the `IndexError` and passes with this change. --- python/python/lancedb/table.py | 17 ++++++++++++++--- python/python/tests/test_table.py | 24 ++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 006ef8cb9..ff866d4db 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5335,6 +5335,12 @@ def _infer_target_schema( # Use the most common length of the list as the dimensions dim = _modal_list_size(peeked.column(i)) + if dim is None: + # Every list entry is null (e.g. embeddings not computed yet), + # so the dimension is unknowable. Leave the column as a + # variable-length list instead of guessing a fixed size. + continue + # Determine target type based on value type if pa.types.is_floating(field.type.value_type): target_type = pa.list_(pa.float32(), dim) @@ -5374,9 +5380,14 @@ def _infer_target_schema( return schema, reader -def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> int: - # Use the most common length of the list as the dimensions - return pc.mode(pc.list_value_length(arr))[0].as_py()["mode"] +def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> Optional[int]: + # Use the most common length of the list as the dimensions. pyarrow's + # `mode` skips nulls, so an all-null column yields an empty result; return + # None in that case so the caller can keep a variable-length list. + modes = pc.mode(pc.list_value_length(arr)) + if len(modes) == 0: + return None + return modes[0].as_py()["mode"] def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]: diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 7a4f5981c..827f2ef58 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -280,6 +280,30 @@ def test_create_table_infers_large_int_vectors(mem_db: DBConnection): assert vector_column.to_pylist() == [[0.0, 300.0]] +def test_create_table_all_null_vector_column(mem_db: DBConnection): + # A vector column can legitimately arrive with every value null, e.g. when + # the embeddings have not been computed yet. The dimension is unknowable in + # that case, so schema inference must not crash and should leave the column + # as a variable-length list rather than a fixed-size one. + data = pa.table( + { + "id": pa.array([1, 2]), + "vector": pa.array([None, None], type=pa.list_(pa.float32())), + } + ) + + table = mem_db.create_table("all_null_vector", data=data) + + vector_field = table.schema.field("vector") + assert pa.types.is_list(vector_field.type) + assert vector_field.type.value_type == pa.float32() + + result = table.to_arrow() + assert result.num_rows == 2 + assert result.column("vector").to_pylist() == [None, None] + assert table.count_rows() == 2 + + @pytest.mark.asyncio async def test_create_table_async_infers_large_int_vectors( mem_db_async: AsyncConnection, From 79e256959fa44f6f335e4b0e6409a6c21b6dcde3 Mon Sep 17 00:00:00 2001 From: Mohd Quamar Tyagi <104281681+Tyagiquamar@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:54:57 +0530 Subject: [PATCH 10/68] fix(python): accept pathlib.Path in get_uri_scheme and get_uri_location (#4132) This PR converts the uri parameter to string before calling urlparse in get_uri_scheme and get_uri_location, allowing pathlib.Path objects to be passed seamlessly without raising a TypeError. --------- Co-authored-by: Xuanwo --- python/python/lancedb/util.py | 15 ++++++++------- python/python/tests/test_util.py | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index 3babc3872..f1a652924 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -34,20 +34,20 @@ def safe_import_adlfs(): adlfs = safe_import_adlfs() -def get_uri_scheme(uri: str) -> str: +def get_uri_scheme(uri: Union[str, pathlib.Path]) -> str: """ Get the scheme of a URI. If the URI does not have a scheme, assume it is a file URI. Parameters ---------- - uri : str + uri : str or pathlib.Path The URI to parse. Returns ------- str: The scheme of the URI. """ - parsed = urlparse(uri) + parsed = urlparse(str(uri)) scheme = parsed.scheme if not scheme: scheme = "file" @@ -61,25 +61,26 @@ def get_uri_scheme(uri: str) -> str: return scheme -def get_uri_location(uri: str) -> str: +def get_uri_location(uri: Union[str, pathlib.Path]) -> str: """ Get the location of a URI. If the parameter is not a url, assumes it is just a path Parameters ---------- - uri : str + uri : str or pathlib.Path The URI to parse. Returns ------- str: Location part of the URL, without scheme """ - parsed = urlparse(uri) + uri_str = str(uri) + parsed = urlparse(uri_str) if len(parsed.scheme) == 1: # Windows drive names are parsed as the scheme # e.g. "c:\path" -> ParseResult(scheme="c", netloc="", path="/path", ...) # So we add special handling here for schemes that are a single character - return uri + return uri_str if not parsed.netloc: return parsed.path diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 741214794..4671109b8 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -30,6 +30,7 @@ import pytest import lancedb from lancedb.util import ( flatten_columns, + get_uri_location, get_uri_scheme, inf_vector_column_query, join_uri, @@ -86,6 +87,22 @@ def test_normalize_uri(): assert parsed_scheme == expected_scheme +def test_get_uri_scheme_and_location_pathlib(): + # Test pathlib.Path instances + abs_path = pathlib.Path("/tmp/data") + rel_path = pathlib.Path("relative/data") + + assert get_uri_scheme(abs_path) == "file" + assert get_uri_scheme(rel_path) == "file" + + assert get_uri_location(abs_path) == str(abs_path) + assert get_uri_location(rel_path) == str(rel_path) + + # String URIs still behave identically + assert get_uri_scheme("s3://bucket/data") == "s3" + assert get_uri_location("s3://bucket/data") == "bucket/data" + + def test_join_uri_remote(): schemes = ["s3", "az", "gs"] for scheme in schemes: From d4b2290babbc8ab83a4287551b0145fd7ed05f44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E5=BF=97=E8=B0=A6?= <89645338+simpleqt@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:25:13 +0800 Subject: [PATCH 11/68] fix(python): correct the voyageai reranker env var docstring (#4201) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `VoyageAIReranker`'s `api_key` docstring said it falls back to `OPENAI_API_KEY`, but the class actually reads `VOYAGE_API_KEY` (and raises a "VOYAGE_API_KEY not set" error without it) — the docstring was copy-pasted from the OpenAI reranker. Docstring-only fix. --- python/python/lancedb/rerankers/voyageai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/python/lancedb/rerankers/voyageai.py b/python/python/lancedb/rerankers/voyageai.py index 0f41ba927..d3dbd59ac 100644 --- a/python/python/lancedb/rerankers/voyageai.py +++ b/python/python/lancedb/rerankers/voyageai.py @@ -32,7 +32,7 @@ class VoyageAIReranker(Reranker): return_score : str, default "relevance" options are "relevance" or "all". Only "relevance" is supported for now. api_key : str, default None - The API key to use. If None, will use the OPENAI_API_KEY environment variable. + The API key to use. If None, will use the VOYAGE_API_KEY environment variable. truncation : Optional[bool], default None """ From 9a590913767d31a0ba5ff524c7884052cec2451f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=99=88=E5=BF=97=E8=B0=A6?= <89645338+simpleqt@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:25:29 +0800 Subject: [PATCH 12/68] fix(python): add missing space in the hybrid search error message (#4202) `LanceHybridQueryBuilder._validate_query`'s ValueError concatenated two adjacent string literals without a separator, rendering "...explicitly for hybrid search.**But** not both.". Added the trailing space. String-only change. --- python/python/lancedb/query.py | 2 +- python/python/tests/test_hybrid_query.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 53f56a2c0..fe8c33861 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2441,7 +2441,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): if query is not None and (vector is not None or text is not None): raise ValueError( "You can either provide a string query in search() method " - "or set `vector()` and `text()` explicitly for hybrid search." + "or set `vector()` and `text()` explicitly for hybrid search. " "But not both." ) diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 8d670a136..19050093c 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -536,3 +536,15 @@ def test_normalize_scores(): assert pc.equal(result, expected), ( f"Expected {expected} but got {result} for invert={invert}" ) + + +def test_hybrid_query_rejects_query_with_vector(sync_table: Table): + # Pins the full message: the space between "hybrid search." and "But" is + # load-bearing (see the _validate_query literal concatenation). + with pytest.raises(ValueError) as excinfo: + sync_table.search("dog", query_type="hybrid").vector([0.5, -0.5]).to_arrow() + assert str(excinfo.value) == ( + "You can either provide a string query in search() method " + "or set `vector()` and `text()` explicitly for hybrid search. " + "But not both." + ) From 82bfe7d38e38285424e843f2491b9680c59c30d2 Mon Sep 17 00:00:00 2001 From: Joaquin Hui Gomez <132194176+joaquinhuigomez@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:25:43 +0100 Subject: [PATCH 13/68] fix: table_names should not list loose .lance files as tables (#4203) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `table_names()` enumerates with `read_dir`, which returns files as well as directories, filtered only on the `.lance` extension — so a loose `foo.lance` *file* at the database prefix is reported as a table. `list_tables()` already applies the correct rule ("Only child directories can be tables"), and `test_listing_ignores_non_table_children` documents it, so the two listing APIs currently disagree. The practical damage is in `drop_all_tables()`, which is built on `table_names`: with a loose `.lance` file present it deletes some real tables, fails on the file with `os error 20 (NotADirectory)`, and then fails the same way on every retry — the database can never be emptied through the API, and the partial drop is not recoverable. `table_names` now lists via `list_with_delimiter` and filters `common_prefixes` through the same `table_name` helper `list_tables` uses; the sort/`start_after`/`limit` handling is unchanged. `drop_tables` additionally attempts every name before returning, still reporting the first error unchanged (kept rather than an aggregate error so the mapped exception types in the bindings stay exactly as they are). Extended the non-table-children test to cover `table_names`, plus tests for the drop behavior — all three fail without the fix. Full `cargo test -p lancedb --lib --features remote` 1320 passed on the repo-pinned 1.97 toolchain; clippy/fmt clean. --- rust/lancedb/src/database/listing.rs | 120 +++++++++++++++++++++------ 1 file changed, 94 insertions(+), 26 deletions(-) diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 71e9f0ef1..700c855dd 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -750,26 +750,39 @@ impl ListingDatabase { uri.push_str(&format!("?{}", query_string)); } let commit_handler = commit_handler_from_url(&uri, &Some(object_store_params)).await?; + // One name that cannot be dropped must not strand the names after it: a caller + // dropping every table would otherwise be left with a partly dropped database that + // the same call cannot finish. Every name is attempted and the first failure is + // still what the call reports, so a caller matching on the error sees what it saw + // before. + let mut first_error = None; for name in names { let dir_name = format!("{}.{}", name, LANCE_EXTENSION); let full_path = self.base_path.clone().join(dir_name.clone()); - commit_handler.delete(&full_path).await?; + let dropped: Result<()> = async { + commit_handler.delete(&full_path).await?; - self.object_store - .remove_dir_all(full_path.clone()) - .await - .map_err(|err| match err { - // this error is not lance::Error::DatasetNotFound, as the method - // `remove_dir_all` may be used to remove something not be a dataset - lance::Error::NotFound { .. } => Error::TableNotFound { - name: name.clone(), - source: Box::new(err), - }, - _ => Error::from(err), - })?; + self.object_store + .remove_dir_all(full_path.clone()) + .await + .map_err(|err| match err { + // this error is not lance::Error::DatasetNotFound, as the method + // `remove_dir_all` may be used to remove something not be a dataset + lance::Error::NotFound { .. } => Error::TableNotFound { + name: name.clone(), + source: Box::new(err), + }, + _ => Error::from(err), + }) + } + .await; + + if let Err(err) = dropped { + first_error.get_or_insert(err); + } } - Ok(()) + first_error.map_or(Ok(()), Err) } /// Inherit storage options from the connection into the target map @@ -930,20 +943,18 @@ impl Database for ListingDatabase { if !request.namespace_path.is_empty() { return self.namespace_database().table_names(request).await; } - let mut f = self + // Only child directories can be tables, the same rule `list_tables` lists by: a + // loose object named like a table is not one, and a name handed back for it is a + // name the caller can neither open nor drop. + let dir_suffix = format!(".{LANCE_EXTENSION}"); + let listing = self .object_store - .read_dir(self.base_path.clone()) - .await? + .list_with_delimiter(Some(&self.base_path)) + .await?; + let mut f = listing + .common_prefixes .iter() - .map(Path::new) - .filter(|path| { - let is_lance = path - .extension() - .and_then(|e| e.to_str()) - .map(|e| e == LANCE_EXTENSION); - is_lance.unwrap_or(false) - }) - .filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from))) + .filter_map(|location| table_name(location, &dir_suffix)) .collect::>(); f.sort(); if let Some(start_after) = request.start_after { @@ -1697,6 +1708,63 @@ mod tests { .unwrap(); assert_eq!(page.tables, vec!["real"]); + + // The deprecated name listing answers the same question and so follows the same + // rule: a name it hands back has to be one a caller can open and drop. + #[allow(deprecated)] + let names = db.table_names(TableNamesRequest::default()).await.unwrap(); + assert_eq!(names, vec!["real"]); + } + + /// Dropping the database drops its tables and leaves everything else alone. A loose + /// `*.lance` object under the database prefix is not a table, so it is neither dropped + /// nor allowed to strand the tables that are. + #[tokio::test] + async fn test_drop_all_tables_leaves_loose_lance_files() { + let (tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b"]).await; + // Sorts between the two tables, so a drop that stops at the first failure would + // leave `b` behind. + let loose = tempdir.path().join("aaa-loose.lance"); + std::fs::write(&loose, b"not a table").unwrap(); + + db.drop_all_tables(&[]).await.unwrap(); + + #[allow(deprecated)] + let names = db.table_names(TableNamesRequest::default()).await.unwrap(); + assert!(names.is_empty(), "a table survived the drop: {names:?}"); + assert!(!tempdir.path().join("a.lance").exists()); + assert!(!tempdir.path().join("b.lance").exists()); + assert!( + loose.exists(), + "a loose object was dropped as if it were a table" + ); + } + + /// A name that cannot be dropped does not strand the names after it: every name is + /// attempted, and the failure is still what the call reports. + #[tokio::test] + async fn test_drop_tables_attempts_every_name() { + let (tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b"]).await; + + let result = db + .drop_tables(vec![ + "a".to_string(), + "missing".to_string(), + "b".to_string(), + ]) + .await; + + assert!( + matches!(result, Err(Error::TableNotFound { .. })), + "dropping a table that is not there did not report it: {result:?}" + ); + assert!(!tempdir.path().join("a.lance").exists()); + assert!( + !tempdir.path().join("b.lance").exists(), + "the table after the failing name was left behind" + ); } #[tokio::test] From d37a6b76a465a1fc62372ffe51d4cf4bb5233633 Mon Sep 17 00:00:00 2001 From: SIDDARTHA REDDY <75976672+SIDDARTHAREDDY8@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:25:58 -0400 Subject: [PATCH 14/68] fix: gate remote job error construction behind the remote feature (#4233) ## Summary Fixes #4142: `lancedb` failed to compile without the `remote` feature. `TerminalResult::decode` in `rust/lancedb/src/job.rs` constructed the `Error::Http` variant unconditionally, but that variant is declared behind `#[cfg(feature = "remote")]` in `rust/lancedb/src/error.rs`. Since the crate's default feature set is empty, a plain `lancedb = "0.38"` dependency failed with: ``` error[E0599]: no variant named `Http` found for enum `error::Error` --> rust/lancedb/src/job.rs:102:40 ``` The module itself is not remote-only (it carries local job support), so instead of gating `pub mod job`, both `Error::Http` constructions now go through a small `remote_job_error` helper that resolves to `Error::Http` with the `remote` feature and falls back to `Error::Runtime` without it. A `request_id` is only ever `Some` for remote jobs, so the fallback is unreachable in practice without the feature. ## Testing - Reproduced the reported failure on current `main`: `cargo check -p lancedb --no-default-features` failed with the exact E0599 errors from the issue before the change, and passes after it. - `cargo check -p lancedb --features remote` passes with no errors and no new warnings. - Added two regression tests in `rust/lancedb/src/job.rs` (`decode_missing_result_uses_feature_appropriate_error`, `decode_unparsable_result_uses_feature_appropriate_error`) asserting the error variant matches the feature configuration. The test code compiles under the fixed configuration (`cargo check --tests` passes), and the branching logic was executed standalone against both feature cfgs (3/3 green with and without `remote`). The full `cargo test` binary link could not complete on the small build VM available here (the linker stalled), so CI coverage of the new tests is appreciated. - `cargo fmt --all -- --check` is clean. ## Disclosure AI assistance was used while preparing this change (code search, editing, and verification commands were agent-driven); the diagnosis, fix design, and all test results above were verified by actually running the builds. --- rust/lancedb/src/job.rs | 80 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 71 insertions(+), 9 deletions(-) diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 711a73033..6865e59ee 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -101,21 +101,19 @@ impl TerminalResult { fn decode(self) -> Result { let value = self.value.ok_or_else(|| match &self.request_id { - Some(request_id) => Error::Http { - source: "successful typed job response did not contain a result".into(), - request_id: request_id.clone(), - status_code: None, - }, + Some(request_id) => remote_job_error( + request_id.clone(), + "successful typed job response did not contain a result".to_string(), + ), None => Error::Runtime { message: "successful typed job did not contain a result".to_string(), }, })?; serde_json::from_value(value).map_err(|error| match self.request_id { - Some(request_id) => Error::Http { - source: format!("failed to parse typed job result: {error}").into(), + Some(request_id) => remote_job_error( request_id, - status_code: None, - }, + format!("failed to parse typed job result: {error}"), + ), None => Error::Runtime { message: format!("failed to parse typed job result: {error}"), }, @@ -123,6 +121,28 @@ impl TerminalResult { } } +/// Build the error for a failed remote typed-job result. +/// +/// The `Error::Http` variant is gated behind the `remote` feature, so without +/// it we fall back to `Error::Runtime`. A `request_id` is only ever `Some` +/// for remote jobs, making the fallback unreachable in practice without the +/// feature. +#[cfg(feature = "remote")] +fn remote_job_error(request_id: String, source: String) -> Error { + Error::Http { + source: source.into(), + request_id, + status_code: None, + } +} + +#[cfg(not(feature = "remote"))] +fn remote_job_error(request_id: String, source: String) -> Error { + Error::Runtime { + message: format!("{source} (request_id={request_id})"), + } +} + type ResultDecoder = Arc Result + Send + Sync>; enum JobInner { @@ -634,4 +654,46 @@ mod tests { assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. }))); assert_eq!(job.status().await.unwrap(), "cancelled"); } + + #[test] + fn decode_missing_result_uses_feature_appropriate_error() { + // A `request_id` is only `Some` for remote jobs. The error for a + // missing result must use `Error::Http` with the `remote` feature and + // fall back to `Error::Runtime` without it (where `Error::Http` + // does not exist). See https://github.com/lancedb/lancedb/issues/4142. + let result = TerminalResult { + value: None, + request_id: Some("req-1".to_string()), + }; + let err = result.decode::().unwrap_err(); + #[cfg(feature = "remote")] + assert!( + matches!(err, Error::Http { .. }), + "expected Error::Http, got {err:?}" + ); + #[cfg(not(feature = "remote"))] + assert!( + matches!(err, Error::Runtime { .. }), + "expected Error::Runtime, got {err:?}" + ); + } + + #[test] + fn decode_unparsable_result_uses_feature_appropriate_error() { + let result = TerminalResult { + value: Some(Value::String("not a u64".to_string())), + request_id: Some("req-1".to_string()), + }; + let err = result.decode::().unwrap_err(); + #[cfg(feature = "remote")] + assert!( + matches!(err, Error::Http { .. }), + "expected Error::Http, got {err:?}" + ); + #[cfg(not(feature = "remote"))] + assert!( + matches!(err, Error::Runtime { .. }), + "expected Error::Runtime, got {err:?}" + ); + } } From 66b5b5bbdcb816b547927fb082f3f96925951e5e Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:26:13 +0800 Subject: [PATCH 15/68] fix(python): clarify unavailable CUDA index acceleration (#4249) ## Summary - preflight CUDA availability and initialization before accelerated index training - direct unavailable or misconfigured CUDA environments to the explicit CPU path - cover both legacy and unified Python index APIs ## Root cause LanceDB delegated accelerator requests directly to PyLance. When PyTorch lacked CUDA support or CUDA initialization failed, users received a low-level assertion or runtime error with no indication that omitting `accelerator` or setting `accelerator=None` creates the index on CPU. The repair validates CUDA before training and preserves explicit accelerator semantics without silently changing the requested backend. ## Scope This PR addresses only the separately reproduced CUDA availability and initialization error. It does not close #1766; the reported query failure after successful CUDA index creation remains open pending reproducible evidence. ## Validation - `uv run --extra tests pytest python/tests/test_table.py -k create_index -q` (10 passed) - `ruff format .` - `ruff check .` - CPU-only integration check confirms the actionable fallback message Related to #1766 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- python/python/lancedb/table.py | 35 +++++++++++++++++++++-- python/python/tests/test_table.py | 47 ++++++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index ff866d4db..35ec7a466 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -119,6 +119,33 @@ def _should_push_down_query_table( return namespace_client is not None and "QueryTable" in pushdown_operations +_CUDA_UNAVAILABLE_MESSAGE = ( + "CUDA acceleration is unavailable in this Python environment. Install a " + "CUDA-enabled PyTorch build and verify that the CUDA device is accessible, or " + "omit `accelerator` (or set `accelerator=None`) to create the index on CPU." +) + + +def _validate_index_accelerator(accelerator: Any) -> None: + if not isinstance(accelerator, str) or accelerator.split(":", 1)[0] != "cuda": + return + + try: + import torch + + cuda_available = torch.cuda.is_available() + except (ImportError, RuntimeError, AssertionError) as err: + raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err + + if not cuda_available: + raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) + + try: + torch.cuda.init() + except (RuntimeError, AssertionError) as err: + raise RuntimeError(_CUDA_UNAVAILABLE_MESSAGE) from err + + def _polars_predicate_pushdown_barrier(frame: Any) -> Any: """Return a Polars frame unchanged while blocking predicate pushdown.""" return frame @@ -3600,7 +3627,9 @@ class LanceTable(Table): # Handle accelerator through pylance if accelerator is not None: - self.to_lance().create_index( + dataset = self.to_lance() + _validate_index_accelerator(accelerator) + dataset.create_index( column=column, index_type=index_type, metric=metric, @@ -3625,6 +3654,8 @@ class LanceTable(Table): acc = getattr(config, "accelerator", None) if acc is not None: # Dispatch to pylance for GPU acceleration + dataset = self.to_lance() + _validate_index_accelerator(acc) index_type_map = { "IvfFlat": "IVF_FLAT", "IvfSq": "IVF_SQ", @@ -3636,7 +3667,7 @@ class LanceTable(Table): cfg_type = type(config).__name__ lance_index_type = index_type_map.get(cfg_type, "IVF_PQ") - self.to_lance().create_index( + dataset.create_index( column=column, index_type=lance_index_type, metric=getattr(config, "distance_type", "l2"), diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 827f2ef58..b241adb0a 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -15,7 +15,7 @@ from datetime import date, datetime, timedelta from decimal import Decimal, localcontext from time import sleep from typing import List -from unittest.mock import patch +from unittest.mock import MagicMock, patch import lancedb from lancedb import table as table_module @@ -1773,6 +1773,51 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection): ) +@pytest.mark.parametrize("api", ["legacy", "unified"]) +@pytest.mark.parametrize("cuda_failure", ["unavailable", "init_error"]) +def test_create_index_cuda_failure_suggests_cpu( + api, cuda_failure, mem_db: DBConnection +): + table = mem_db.create_table( + f"test_cuda_unavailable_{api}", + data=[{"vector": [3.1, 4.1]}, {"vector": [5.9, 26.5]}], + ) + torch = MagicMock() + torch.cuda.is_available.return_value = cuda_failure == "init_error" + if cuda_failure == "init_error": + torch.cuda.init.side_effect = RuntimeError("CUDA driver initialization failed") + + with ( + patch.dict(sys.modules, {"torch": torch}), + patch.object(table, "to_lance") as to_lance, + ): + if api == "legacy": + with ( + pytest.warns(DeprecationWarning, match="create_index"), + pytest.raises(RuntimeError, match="accelerator=None.*CPU"), + ): + table.create_index( + metric="dot", + num_partitions=1, + num_sub_vectors=1, + vector_column_name="vector", + accelerator="cuda", + ) + else: + with pytest.raises(RuntimeError, match="accelerator=None.*CPU"): + table.create_index( + "vector", + config=IvfPq( + distance_type="dot", + num_partitions=1, + num_sub_vectors=1, + accelerator="cuda", + ), + ) + + to_lance.return_value.create_index.assert_not_called() + + @patch("lancedb.table.AsyncTable.create_index") def test_create_index_name_and_train_parameters( mock_create_index, mem_db: DBConnection From 624defd109c161505e84b26e66f43a35b71ac16c Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 01:26:28 +0800 Subject: [PATCH 16/68] fix: preserve sliced child validity during struct merge (#4261) ## Root cause The pinned Lance Arrow crate rebuilt a sliced struct child with its original Arrow offset but supplied a validity bitmap containing only the visible rows. Arrow requires the bitmap to cover the offset as well, so filtered vector searches involving nullable nested fields could panic with `null_bit_buffer size too small`. ## Fix - Carry an offset-aligned validity bitmap when propagating parent nulls to a sliced child. - Vendor the pinned `lance-arrow` crate until a Lance release includes the fix; include it in Maturin source distributions. - Add unit regressions for sliced children and parent bitmaps, plus a record-batch merge regression. ## Validation - The issue's 8,600-row Python reproducer completes its filtered flat dot search and returns 8,448 rows without a panic. - The targeted Rust merge regression and two focused Python query tests pass. - Rust formatting, Ruff formatting and lint, the vendored crate tests, and the workspace Cargo checks passed. The Maturin source archive includes the patched crate. Fixes #4260 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- Cargo.lock | 2 +- Cargo.toml | 7 +- python/Cargo.toml | 2 + python/vendor/lance-arrow/Cargo.toml | 33 + python/vendor/lance-arrow/README.lancedb.md | 7 + python/vendor/lance-arrow/README.md | 6 + python/vendor/lance-arrow/src/bfloat16.rs | 406 +++ python/vendor/lance-arrow/src/deepcopy.rs | 356 +++ python/vendor/lance-arrow/src/floats.rs | 292 ++ python/vendor/lance-arrow/src/ipc.rs | 621 ++++ python/vendor/lance-arrow/src/json.rs | 1538 +++++++++ python/vendor/lance-arrow/src/lib.rs | 3102 +++++++++++++++++++ python/vendor/lance-arrow/src/list.rs | 177 ++ python/vendor/lance-arrow/src/memory.rs | 92 + python/vendor/lance-arrow/src/scalar.rs | 302 ++ python/vendor/lance-arrow/src/schema.rs | 208 ++ python/vendor/lance-arrow/src/stream.rs | 554 ++++ python/vendor/lance-arrow/src/struct.rs | 227 ++ rust/lancedb/tests/nullable_struct_merge.rs | 64 + 19 files changed, 7994 insertions(+), 2 deletions(-) create mode 100644 python/vendor/lance-arrow/Cargo.toml create mode 100644 python/vendor/lance-arrow/README.lancedb.md create mode 100644 python/vendor/lance-arrow/README.md create mode 100644 python/vendor/lance-arrow/src/bfloat16.rs create mode 100644 python/vendor/lance-arrow/src/deepcopy.rs create mode 100644 python/vendor/lance-arrow/src/floats.rs create mode 100644 python/vendor/lance-arrow/src/ipc.rs create mode 100644 python/vendor/lance-arrow/src/json.rs create mode 100644 python/vendor/lance-arrow/src/lib.rs create mode 100644 python/vendor/lance-arrow/src/list.rs create mode 100644 python/vendor/lance-arrow/src/memory.rs create mode 100644 python/vendor/lance-arrow/src/scalar.rs create mode 100644 python/vendor/lance-arrow/src/schema.rs create mode 100644 python/vendor/lance-arrow/src/stream.rs create mode 100644 python/vendor/lance-arrow/src/struct.rs create mode 100644 rust/lancedb/tests/nullable_struct_merge.rs diff --git a/Cargo.lock b/Cargo.lock index 7f36dc088..34ee05cb1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5136,7 +5136,6 @@ dependencies = [ [[package]] name = "lance-arrow" version = "15.0.0-beta.1" -source = "git+https://github.com/lance-format/lance.git?tag=v15.0.0-beta.1#46a7cfe1fba5efa1f90deea4c7b14d414b991409" dependencies = [ "arrow-array", "arrow-buffer", @@ -5834,6 +5833,7 @@ dependencies = [ "datafusion-common", "env_logger", "futures", + "lance-arrow", "lance-core", "lance-io", "lance-namespace", diff --git a/Cargo.toml b/Cargo.toml index 3ac246149..870ef5d07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,7 @@ lance-datagen = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" lance-file = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-io = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-index = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-jsonb = { version = "=15.0.0-beta.1", tag = "v15.0.0-beta.1", git = "https://github.com/lance-format/lance.git" } lance-linalg = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-namespace = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-namespace-impls = { "version" = "=15.0.0-beta.1", default-features = false, "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } @@ -26,7 +27,7 @@ lance-table = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = lance-testing = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-datafusion = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } lance-encoding = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } -lance-arrow = { "version" = "=15.0.0-beta.1", "tag" = "v15.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" } +lance-arrow = { version = "=15.0.0-beta.1", path = "python/vendor/lance-arrow" } lancedb = { path = "rust/lancedb", default-features = false } ahash = "0.8" # Note that this one does not include pyarrow @@ -79,6 +80,10 @@ tonic = { version = "0.14", features = ["tls-native-roots", "tls-aws-lc"] } uuid = { version = "1.7.0", features = ["v4", "v7"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } +[patch."https://github.com/lance-format/lance.git"] +# Carry the offset-safe validity fix until it is available in a Lance release. +lance-arrow = { path = "python/vendor/lance-arrow" } + [patch.crates-io] # Lance v14.0.0-beta.18 requires multipart write modes from this upstream revision. object_store = { git = "https://github.com/apache/arrow-rs-object-store.git", rev = "c10758c207c0103355f940f0c736fb8ecb157b04" } diff --git a/python/Cargo.toml b/python/Cargo.toml index ad1790a99..076299a7d 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -22,6 +22,8 @@ bytes.workspace = true lancedb.workspace = true datafusion-common.workspace = true lance-core.workspace = true +# The workspace path dependency includes the patched crate in Maturin sources. +lance-arrow.workspace = true lance-namespace.workspace = true lance-namespace-impls.workspace = true lance-io.workspace = true diff --git a/python/vendor/lance-arrow/Cargo.toml b/python/vendor/lance-arrow/Cargo.toml new file mode 100644 index 000000000..9cb156f2e --- /dev/null +++ b/python/vendor/lance-arrow/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "lance-arrow" +version = "15.0.0-beta.1" +edition = "2024" +authors = ["Lance Devs "] +license = "Apache-2.0" +repository = "https://github.com/lance-format/lance" +readme = "README.md" +description = "Arrow Extension for Lance" +keywords = ["data-format", "data-science", "machine-learning", "apache-arrow", "data-analytics"] +categories = ["database-implementations"] + +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +arrow-array.workspace = true +arrow-buffer.workspace = true +arrow-data.workspace = true +arrow-ipc = { workspace = true, features = ["zstd"] } +arrow-ord.workspace = true +arrow-schema.workspace = true +arrow-select.workspace = true +bytemuck = { version = "1", default-features = false, features = ["extern_crate_alloc"] } +bytes.workspace = true +futures.workspace = true +half = { workspace = true, features = ["std", "bytemuck"] } +lance-jsonb.workspace = true +num-traits.workspace = true +rand = { workspace = true, features = ["small_rng"] } + +[target.'cfg(target_arch = "wasm32")'.dependencies] +getrandom = { version = "0.2", features = ["js"] } diff --git a/python/vendor/lance-arrow/README.lancedb.md b/python/vendor/lance-arrow/README.lancedb.md new file mode 100644 index 000000000..5ed714e08 --- /dev/null +++ b/python/vendor/lance-arrow/README.lancedb.md @@ -0,0 +1,7 @@ +This directory contains `rust/lance-arrow` from Lance tag `v15.0.0-beta.1` +(commit `46a7cfe1fba5efa1f90deea4c7b14d414b991409`), licensed under Apache-2.0. + +LanceDB carries a local fix in `src/lib.rs` for sliced child validity bitmaps +until a Lance release includes the fix. The patch is selected in the root +`Cargo.toml`; the rest of Lance stays on the pinned release. An explicit unsafe +block in `src/bfloat16.rs` lets this copy pass LanceDB's Rust 2024 warning gate. diff --git a/python/vendor/lance-arrow/README.md b/python/vendor/lance-arrow/README.md new file mode 100644 index 000000000..e0f129a3a --- /dev/null +++ b/python/vendor/lance-arrow/README.md @@ -0,0 +1,6 @@ +# lance-arrow + +`lance-arrow` is an internal sub-crate, containing [Apache-Arrow](https://github.com/apache/arrow-rs) +extensions used by [Lance](https://github.com/lance-format/lance). + +**Important Note**: This crate is **not intended for external usage**. diff --git a/python/vendor/lance-arrow/src/bfloat16.rs b/python/vendor/lance-arrow/src/bfloat16.rs new file mode 100644 index 000000000..1a60e5ece --- /dev/null +++ b/python/vendor/lance-arrow/src/bfloat16.rs @@ -0,0 +1,406 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! bfloat16 support for Apache Arrow. + +use std::fmt::Formatter; +use std::slice; + +use arrow_array::{Array, FixedSizeBinaryArray, builder::BooleanBufferBuilder}; +use arrow_buffer::{Buffer, MutableBuffer}; +use arrow_data::ArrayData; +use arrow_schema::{ArrowError, DataType, Field as ArrowField}; +use half::bf16; + +use crate::{ARROW_EXT_NAME_KEY, FloatArray}; + +/// The name of the bfloat16 extension in Arrow metadata +pub const BFLOAT16_EXT_NAME: &str = "lance.bfloat16"; + +/// Check whether the given field is a bfloat16 field +/// +/// A field is a bfloat16 field if it has a data type of `FixedSizeBinary(2)` and the metadata +/// contains the bfloat16 extension name. +pub fn is_bfloat16_field(field: &ArrowField) -> bool { + field.data_type() == &DataType::FixedSizeBinary(2) + && field + .metadata() + .get(ARROW_EXT_NAME_KEY) + .map(|name| name == BFLOAT16_EXT_NAME) + .unwrap_or_default() +} + +/// The bfloat16 data type +/// +/// This implements the [`ArrowFloatType`](crate::floats::ArrowFloatType) trait for bfloat16 values. +#[derive(Debug)] +pub struct BFloat16Type {} + +/// An array of bfloat16 values +/// +/// Note that bfloat16 is not the same thing as fp16 which is supported natively by arrow-rs. +#[derive(Clone)] +pub struct BFloat16Array { + inner: FixedSizeBinaryArray, +} + +impl std::fmt::Debug for BFloat16Array { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "BFloat16Array\n[\n")?; + from_arrow::print_long_array(&self.inner, f, |array, i, f| { + if array.is_null(i) { + write!(f, "null") + } else { + let binary_values = array.value(i); + let value = + bf16::from_bits(u16::from_le_bytes([binary_values[0], binary_values[1]])); + write!(f, "{:?}", value) + } + })?; + write!(f, "]") + } +} + +impl BFloat16Array { + pub fn from_iter_values(iter: impl IntoIterator) -> Self { + let values: Vec = iter.into_iter().collect(); + values.into() + } + + pub fn len(&self) -> usize { + self.inner.len() + } + + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + pub fn is_null(&self, i: usize) -> bool { + self.inner.is_null(i) + } + + pub fn null_count(&self) -> usize { + self.inner.null_count() + } + + pub fn iter(&self) -> BFloat16Iter<'_> { + BFloat16Iter { + array: self, + index: 0, + } + } + + pub fn value(&self, i: usize) -> bf16 { + assert!( + i < self.len(), + "Trying to access an element at index {} from a BFloat16Array of length {}", + i, + self.len() + ); + // Safety: + // `i < self.len() + unsafe { self.value_unchecked(i) } + } + + /// # Safety + /// Caller must ensure that `i < self.len()` + pub unsafe fn value_unchecked(&self, i: usize) -> bf16 { + // SAFETY: the caller guarantees that `i` is within this array's bounds. + let binary_value = unsafe { self.inner.value_unchecked(i) }; + bf16::from_bits(u16::from_le_bytes([binary_value[0], binary_value[1]])) + } + + pub fn into_inner(self) -> FixedSizeBinaryArray { + self.inner + } +} + +impl FromIterator> for BFloat16Array { + fn from_iter>>(iter: I) -> Self { + let mut buffer = MutableBuffer::new(10); + // No null buffer builder :( + let mut nulls = BooleanBufferBuilder::new(10); + let mut len = 0; + + for maybe_value in iter { + if let Some(value) = maybe_value { + let bytes = value.to_le_bytes(); + buffer.extend(bytes); + } else { + buffer.extend([0u8, 0u8]); + } + nulls.append(maybe_value.is_some()); + len += 1; + } + + let null_buffer = nulls.finish(); + let num_valid = null_buffer.count_set_bits(); + let null_buffer = if num_valid == len { + None + } else { + Some(null_buffer.into_inner()) + }; + + let array_data = ArrayData::builder(DataType::FixedSizeBinary(2)) + .len(len) + .add_buffer(buffer.into()) + .null_bit_buffer(null_buffer); + // SAFETY: the value buffer contains exactly `2 * len` bytes (two bytes + // pushed per iteration of the loop above, including the zero-fill for + // null slots), which matches the `FixedSizeBinary(2)` storage layout. + // The null bit buffer, when present, has `len` bits appended above, so + // its length covers the array's logical range. + let array_data = unsafe { array_data.build_unchecked() }; + Self { + inner: FixedSizeBinaryArray::from(array_data), + } + } +} + +impl FromIterator for BFloat16Array { + fn from_iter>(iter: I) -> Self { + Self::from_iter_values(iter) + } +} + +impl From> for BFloat16Array { + fn from(data: Vec) -> Self { + let len = data.len(); + // Zero-copy: `bf16` is `#[repr(transparent)]` over `u16` and derives + // `bytemuck::Pod`, so `cast_vec` reinterprets the allocation in place — + // no per-element copy or heap alloc. The crate-root `compile_error!` + // pins `target_endian = "little"`, so the resulting bytes match the + // `FixedSizeBinary(2)` on-disk order Lance writes elsewhere. + let raw: Vec = bytemuck::cast_vec(data); + let array_data = ArrayData::builder(DataType::FixedSizeBinary(2)) + .len(len) + .add_buffer(Buffer::from_vec(raw)); + // SAFETY: the value buffer contains exactly `2 * len` bytes — one + // `u16` per element after the layout-compatible cast — matching the + // `FixedSizeBinary(2)` storage layout. No null buffer is attached, so + // every element is logically valid. + let array_data = unsafe { array_data.build_unchecked() }; + Self { + inner: FixedSizeBinaryArray::from(array_data), + } + } +} + +impl TryFrom for BFloat16Array { + type Error = ArrowError; + + fn try_from(value: FixedSizeBinaryArray) -> Result { + if value.value_length() == 2 { + Ok(Self { inner: value }) + } else { + Err(ArrowError::InvalidArgumentError( + "FixedSizeBinaryArray must have a value length of 2".to_string(), + )) + } + } +} + +impl PartialEq for BFloat16Array { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +pub struct BFloat16Iter<'a> { + array: &'a BFloat16Array, + index: usize, +} + +impl<'a> Iterator for BFloat16Iter<'a> { + type Item = Option; + + fn next(&mut self) -> Option { + if self.index >= self.array.len() { + return None; + } + let i = self.index; + self.index += 1; + if self.array.is_null(i) { + Some(None) + } else { + Some(Some(self.array.value(i))) + } + } +} + +/// Methods that are lifted from arrow-rs temporarily until they are made public. +mod from_arrow { + use arrow_array::Array; + + /// Helper function for printing potentially long arrays. + pub(super) fn print_long_array( + array: &A, + f: &mut std::fmt::Formatter, + print_item: F, + ) -> std::fmt::Result + where + A: Array, + F: Fn(&A, usize, &mut std::fmt::Formatter) -> std::fmt::Result, + { + let head = std::cmp::min(10, array.len()); + + for i in 0..head { + if array.is_null(i) { + writeln!(f, " null,")?; + } else { + write!(f, " ")?; + print_item(array, i, f)?; + writeln!(f, ",")?; + } + } + if array.len() > 10 { + if array.len() > 20 { + writeln!(f, " ...{} elements...,", array.len() - 20)?; + } + + let tail = std::cmp::max(head, array.len() - 10); + + for i in tail..array.len() { + if array.is_null(i) { + writeln!(f, " null,")?; + } else { + write!(f, " ")?; + print_item(array, i, f)?; + writeln!(f, ",")?; + } + } + } + Ok(()) + } +} + +impl FloatArray for FixedSizeBinaryArray { + type FloatType = BFloat16Type; + + /// Returns the underlying `bf16` values as a borrowed slice. + /// + /// # Preconditions + /// + /// - `value_length()` must be 2 (the `FixedSizeBinary(2)` storage shape + /// used by [`BFloat16Array`]). Asserted at entry. + /// - The value buffer must be at least 2-byte aligned. Lance's in-tree + /// constructors always satisfy this: value buffers are built either via + /// `MutableBuffer` (aligned to arrow-buffer's `ALIGNMENT` constant, ≥32 + /// bytes) or via `Buffer::from_vec::` (aligned to `align_of::()` + /// == 2); both meet `bf16`'s 2-byte requirement. Externally-built + /// `FixedSizeBinaryArray`s arriving via FFI, IPC, or + /// `Buffer::from_custom_allocation` are not required by arrow-rs to be + /// aligned beyond a single byte; passing one to this method violates the + /// precondition. A `debug_assert` below catches such inputs in debug and + /// test builds. + /// + /// # Endianness + /// + /// `lance-arrow` is gated on `target_endian = "little"` at the crate root, + /// so this method always returns values in the same byte order Lance writes + /// (see [`BFloat16Array::value`] and the [`FromIterator`] impls). + fn as_slice(&self) -> &[bf16] { + assert_eq!( + self.value_length(), + 2, + "BFloat16 arrays must use FixedSizeBinary(2) storage" + ); + debug_assert_eq!( + (self.value_data().as_ptr() as usize) % std::mem::align_of::(), + 0, + "BFloat16 value buffer must be at least 2-byte aligned" + ); + // SAFETY: + // - The assert above pins `value_size == 2`, so `value_data().len() / 2` + // equals the array's logical element count. + // `FixedSizeBinaryArray::From` constructs its value buffer + // as `buffers[0].slice_with_length(offset * 2, len * 2)` (arrow-array + // `fixed_size_binary_array.rs`), so `value_data()` already returns + // the offset-adjusted slice. Do not replace `value_data()` with an + // accessor that returns the un-sliced backing buffer. + // - `bf16` is `#[repr(transparent)]` over `u16` (size 2, alignment 2); + // every `u16` bit pattern is a valid `bf16`, so any byte content + // yields a defined value — never UB. + // - Alignment is the caller's responsibility per the precondition + // documented above. The `debug_assert_eq!` immediately preceding this + // block catches violations in debug and test builds only — release + // builds rely on callers honoring the precondition. arrow-rs + // declares `FixedSizeBinary(n)`'s + // `BufferSpec::FixedWidth { alignment: align_of::() == 1 }` + // (arrow-data `data.rs`), so arrow-rs alone does not guarantee + // 2-byte alignment. Lance's in-tree construction paths build value + // buffers via `MutableBuffer` (arrow-buffer `ALIGNMENT` constant, + // ≥32 bytes) or `Buffer::from_vec::` (2-byte aligned), both of + // which satisfy `bf16`'s 2-byte requirement. + // - The returned slice borrows from `self`; the underlying ref-counted, + // immutable Arrow buffer cannot be mutated or freed for the slice's + // lifetime. + unsafe { + slice::from_raw_parts( + self.value_data().as_ptr() as *const bf16, + self.value_data().len() / 2, + ) + } + } + + fn from_values(values: Vec) -> Self { + BFloat16Array::from(values).into_inner() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basics() { + let values: Vec = vec![1.0, 2.0, 3.0]; + let values: Vec = values.iter().map(|v| bf16::from_f32(*v)).collect(); + + let array = BFloat16Array::from_iter_values(values.clone()); + let array2 = BFloat16Array::from(values.clone()); + assert_eq!(array, array2); + assert_eq!(array.len(), 3); + + // Pin the raw little-endian bytes emitted by `From>` (rewritten to + // reinterpret the Vec via `bytemuck::cast_vec`), so a layout/byte-order + // regression is caught directly rather than only through Debug formatting. + // bf16 is the high 16 bits of the f32: 1.0->0x3F80, 2.0->0x4000, 3.0->0x4040. + let inner = array2.clone().into_inner(); + let raw_bytes: Vec = (0..inner.len()) + .flat_map(|i| inner.value(i).to_vec()) + .collect(); + assert_eq!(raw_bytes, vec![0x80, 0x3F, 0x00, 0x40, 0x40, 0x40]); + + let expected_fmt = "BFloat16Array\n[\n 1.0,\n 2.0,\n 3.0,\n]"; + assert_eq!(expected_fmt, format!("{:?}", array)); + + for (expected, value) in values.iter().zip(array.iter()) { + assert_eq!(Some(*expected), value); + } + + for (expected, value) in values.as_slice().iter().zip(array2.iter()) { + assert_eq!(Some(*expected), value); + } + + let arrow_array = array.into_inner(); + assert_eq!(arrow_array.as_slice(), values.as_slice()); + } + + #[test] + fn test_nulls() { + let values: Vec> = + vec![Some(bf16::from_f32(1.0)), None, Some(bf16::from_f32(3.0))]; + let array = BFloat16Array::from_iter(values.clone()); + assert_eq!(array.len(), 3); + assert_eq!(array.null_count(), 1); + + let expected_fmt = "BFloat16Array\n[\n 1.0,\n null,\n 3.0,\n]"; + assert_eq!(expected_fmt, format!("{:?}", array)); + + for (expected, value) in values.iter().zip(array.iter()) { + assert_eq!(*expected, value); + } + } +} diff --git a/python/vendor/lance-arrow/src/deepcopy.rs b/python/vendor/lance-arrow/src/deepcopy.rs new file mode 100644 index 000000000..5df4339b5 --- /dev/null +++ b/python/vendor/lance-arrow/src/deepcopy.rs @@ -0,0 +1,356 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_array::{Array, RecordBatch, make_array}; +use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer}; +use arrow_data::{ArrayData, ArrayDataBuilder, transform::MutableArrayData}; +use arrow_schema::DataType; + +pub fn deep_copy_buffer(buffer: &Buffer) -> Buffer { + Buffer::from(buffer.as_slice()) +} + +pub fn deep_copy_nulls(nulls: Option<&NullBuffer>) -> Option { + let nulls = nulls?; + let bit_buffer = deep_copy_buffer(nulls.inner().inner()); + // SAFETY: `null_count` is taken from the source `NullBuffer`, which already + // upheld `NullBuffer::new_unchecked`'s invariant — the unset-bit count over + // the logical bit slice `[bit_offset, bit_offset + bit_len)`. `NullBuffer::slice` + // adjusts only `BooleanBuffer::bit_offset` / `bit_len` and never byte-advances + // the inner `Buffer`, so `deep_copy_buffer` (which copies the source `Buffer`'s + // `as_slice()` view from byte 0) reproduces the exact bit pattern at the same + // bit offsets; the unset-bit count is therefore preserved. `BooleanBuffer::new` + // panics (does not UB) if `bit_offset + bit_len > 8 * buffer.len()`, and the + // copy has the same length, so that check still passes. + Some(unsafe { + NullBuffer::new_unchecked( + BooleanBuffer::new(bit_buffer, nulls.offset(), nulls.len()), + nulls.null_count(), + ) + }) +} + +pub fn deep_copy_array_data(data: &ArrayData) -> ArrayData { + let data_type = data.data_type().clone(); + let len = data.len(); + let nulls = deep_copy_nulls(data.nulls()); + let offset = data.offset(); + let buffers = data + .buffers() + .iter() + .map(deep_copy_buffer) + .collect::>(); + let child_data = data + .child_data() + .iter() + .map(deep_copy_array_data) + .collect::>(); + // SAFETY: `build_unchecked` inherits `ArrayData::new_unchecked`'s contract — + // `(data_type, len, offset, nulls, buffers, child_data)` must form a valid + // Arrow array. This call reproduces `data` structurally: `data_type`, `len`, + // and `offset` are forwarded unchanged; each buffer is replaced by a byte- + // identical copy of its offset-applied `as_slice()` view (the output buffer + // is `MutableBuffer`-allocated, at least as aligned as the source); `nulls` + // is deep-copied with the same bit offset/length and unset-bit count (see + // `deep_copy_nulls`); `child_data` is recursively cloned with the same + // guarantee. Every value-level invariant the source upheld — UTF-8 validity, + // monotonic offsets, in-bounds dictionary indices, run-end monotonicity, + // struct child-length matching — therefore transfers to the copy. If the + // source `ArrayData` was itself constructed via `new_unchecked` with an + // invalid payload, this function faithfully reproduces that invalidity. + unsafe { + ArrayDataBuilder::new(data_type) + .len(len) + .nulls(nulls) + .offset(offset) + .buffers(buffers) + .child_data(child_data) + .build_unchecked() + } +} + +pub fn deep_copy_array(array: &dyn Array) -> Arc { + let data = array.to_data(); + let data = deep_copy_array_data(&data); + make_array(data) +} + +pub fn deep_copy_batch(batch: &RecordBatch) -> crate::Result { + let arrays = batch + .columns() + .iter() + .map(|array| deep_copy_array(array)) + .collect::>(); + RecordBatch::try_new(batch.schema(), arrays) +} + +/// Deep copy array data, extracting only the sliced portion using MutableArrayData +/// This is the most efficient and correct way to copy just the sliced data +pub fn deep_copy_array_data_sliced(data: &ArrayData) -> ArrayData { + // Use MutableArrayData to efficiently copy just the slice + let mut mutable = MutableArrayData::new(vec![data], false, data.len()); + + // Which index this takes depends on the layout, because arrow's extenders + // disagree about who applies `ArrayData::offset()`. The bit-packed one adds + // it to the raw values buffer itself, so an offset-applied index there + // reads from twice the offset and runs off the end of the buffer. Every + // other layout either reads through an already-offset-applied view or + // forwards the index to children carrying their own offsets, and takes the + // absolute index it has always been given. + let start = match data.data_type() { + DataType::Boolean => 0, + _ => data.offset(), + }; + mutable.extend(0, start, start + data.len()); + + // Freeze into immutable ArrayData + mutable.freeze() +} + +/// Deep copy an array, extracting only the sliced portion using MutableArrayData +pub fn deep_copy_array_sliced(array: &dyn Array) -> Arc { + let data = array.to_data(); + let data = deep_copy_array_data_sliced(&data); + make_array(data) +} + +/// Deep copy a RecordBatch, extracting only the sliced portion using MutableArrayData +pub fn deep_copy_batch_sliced(batch: &RecordBatch) -> crate::Result { + let arrays = batch + .columns() + .iter() + .map(|array| deep_copy_array_sliced(array)) + .collect::>(); + RecordBatch::try_new(batch.schema(), arrays) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{Array, BooleanArray, Int32Array, RecordBatch, StringArray}; + use arrow_data::ArrayDataBuilder; + use arrow_schema::{DataType, Field, Schema}; + + #[test] + fn raw_sliced_fixed_size_list_data_keeps_its_child_offset() { + // `ArrayData::slice` records the offset on the parent and leaves the + // child whole -- a shape `FixedSizeListArray::slice` never produces, + // but a legal one this public helper can be handed directly. Arrow's + // extender forwards `start * size` to that unsliced child and adds + // nothing, so the index it gets has to be the absolute one; this is + // here to stop the boolean fix from being generalised over it. + let child = Int32Array::from(vec![10, 11, 20, 21]).to_data(); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let data = ArrayDataBuilder::new(DataType::FixedSizeList(field, 2)) + .len(2) + .add_child_data(child) + .build() + .unwrap(); + let sliced = data.slice(1, 1); + + let copied = super::deep_copy_array_data_sliced(&sliced); + let copied_child = Int32Array::from(copied.child_data()[0].clone()); + assert_eq!(copied_child.values().as_ref(), &[20, 21]); + } + + #[test] + fn raw_sliced_struct_of_fixed_size_list_reaches_the_right_values() { + // Slicing raw struct data pushes the window into the struct's children, + // and a fixed-size-list child takes it as its own parent offset with + // its values left whole. Two levels of offset, and the absolute index + // has to remain right through both of them. + let values = Int32Array::from(vec![10, 11, 20, 21, 30, 31]).to_data(); + let item = Arc::new(Field::new_list_field(DataType::Int32, false)); + let list = ArrayDataBuilder::new(DataType::FixedSizeList(item, 2)) + .len(3) + .add_child_data(values) + .build() + .unwrap(); + let field = Arc::new(Field::new("l", list.data_type().clone(), false)); + let data = ArrayDataBuilder::new(DataType::Struct(vec![field].into())) + .len(3) + .add_child_data(list) + .build() + .unwrap(); + let sliced = data.slice(2, 1); + + let copied = super::deep_copy_array_data_sliced(&sliced); + let copied_values = Int32Array::from(copied.child_data()[0].child_data()[0].clone()); + assert_eq!(copied_values.values().as_ref(), &[30, 31]); + } + + #[test] + fn raw_parent_offset_struct_keeps_the_selected_child_row() { + // A struct whose children are whole and whose window is the parent + // offset: arrow's struct extender forwards the index to those children + // untouched, so it has to be the absolute one. + let field = Arc::new(Field::new("a", DataType::Int32, false)); + let data = ArrayDataBuilder::new(DataType::Struct(vec![field].into())) + .len(1) + .offset(1) + .add_child_data(Int32Array::from(vec![10, 20]).to_data()) + .build() + .unwrap(); + + let copied = super::deep_copy_array_data_sliced(&data); + let copied_child = Int32Array::from(copied.child_data()[0].clone()); + assert_eq!(copied_child.values().as_ref(), &[20]); + } + + #[test] + fn sliced_boolean_deep_copy_reads_from_the_slice() { + // A boolean slice keeps its `ArrayData::offset()` -- the buffer cannot + // advance by a fraction of a byte -- so a copy that adds the offset on + // top of what `MutableArrayData` already applies reads past the end of + // the values buffer and panics. + let array = BooleanArray::from(vec![true, false, true, false, true, false, true, false]); + for (offset, len) in [(0usize, 8usize), (1, 7), (3, 5), (7, 1)] { + let sliced = array.slice(offset, len); + let copied = super::deep_copy_array_sliced(&sliced); + let copied = copied.as_any().downcast_ref::().unwrap(); + let expected: Vec = (0..len).map(|i| sliced.value(i)).collect(); + let actual: Vec = (0..len).map(|i| copied.value(i)).collect(); + assert_eq!(actual, expected, "offset={offset} len={len}"); + } + } + + #[test] + fn sliced_boolean_deep_copy_keeps_its_nulls() { + let array = BooleanArray::from(vec![ + Some(true), + None, + Some(false), + Some(true), + None, + Some(false), + ]); + let sliced = array.slice(1, 4); + let copied = super::deep_copy_array_sliced(&sliced); + let copied = copied.as_any().downcast_ref::().unwrap(); + let expected: Vec> = (0..sliced.len()) + .map(|i| (!sliced.is_null(i)).then(|| sliced.value(i))) + .collect(); + let actual: Vec> = (0..copied.len()) + .map(|i| (!copied.is_null(i)).then(|| copied.value(i))) + .collect(); + assert_eq!(actual, expected); + } + + #[test] + fn test_deep_copy_sliced_array_with_nulls() { + let array = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + None, + Some(5), + ])); + let sliced_array = array.slice(1, 3); + let copied_array = super::deep_copy_array(&sliced_array); + assert_eq!(sliced_array.len(), copied_array.len()); + assert_eq!(sliced_array.nulls(), copied_array.nulls()); + } + + #[test] + fn test_deep_copy_array_data_sliced() { + let array = Int32Array::from((0..1000).collect::>()); + let sliced = array.slice(100, 10); + + let sliced_data = sliced.to_data(); + let copied_data = super::deep_copy_array_data_sliced(&sliced_data); + + assert_eq!(copied_data.len(), 10); + assert_eq!(copied_data.offset(), 0); + + // Verify data correctness + let copied_array = Int32Array::from(copied_data); + for i in 0..10 { + assert_eq!(copied_array.value(i), 100 + i as i32); + } + } + + #[test] + fn test_deep_copy_array_sliced() { + let array = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); + let sliced = array.slice(1, 3); + + let copied = super::deep_copy_array_sliced(&sliced); + + assert_eq!(copied.len(), 3); + let copied_int = copied.as_any().downcast_ref::().unwrap(); + assert_eq!(copied_int.value(0), 2); + assert_eq!(copied_int.value(1), 3); + assert_eq!(copied_int.value(2), 4); + } + + #[test] + fn test_deep_copy_batch_sliced() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ])); + + let id_array = Arc::new(Int32Array::from((0..100).collect::>())); + let name_array = Arc::new(StringArray::from( + (0..100) + .map(|i| format!("name_{}", i)) + .collect::>(), + )); + + let batch = RecordBatch::try_new( + schema, + vec![id_array as Arc, name_array as Arc], + ) + .unwrap(); + + let sliced = batch.slice(10, 5); + let copied = super::deep_copy_batch_sliced(&sliced).unwrap(); + + assert_eq!(copied.num_rows(), 5); + assert_eq!(copied.num_columns(), 2); + + // Verify data correctness + let id_col = copied + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let name_col = copied + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + for i in 0..5 { + assert_eq!(id_col.value(i), 10 + i as i32); + assert_eq!(name_col.value(i), format!("name_{}", 10 + i)); + } + } + + #[test] + fn test_deep_copy_array_sliced_with_nulls() { + let array = Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(3), + None, + Some(5), + ])); + let sliced = array.slice(1, 3); // [None, Some(3), None] + + let copied = super::deep_copy_array_sliced(&sliced); + + assert_eq!(copied.len(), 3); + assert_eq!(copied.null_count(), 2); // Two nulls in the slice + + let copied_int = copied.as_any().downcast_ref::().unwrap(); + assert!(!copied_int.is_valid(0)); // None + assert!(copied_int.is_valid(1)); // Some(3) + assert!(!copied_int.is_valid(2)); // None + assert_eq!(copied_int.value(1), 3); + } +} diff --git a/python/vendor/lance-arrow/src/floats.rs b/python/vendor/lance-arrow/src/floats.rs new file mode 100644 index 000000000..0024fbb02 --- /dev/null +++ b/python/vendor/lance-arrow/src/floats.rs @@ -0,0 +1,292 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Floats Array + +use std::fmt::{Debug, Display}; +use std::iter::Sum; +use std::sync::Arc; +use std::{ + fmt::Formatter, + ops::{AddAssign, DivAssign}, +}; + +use arrow_array::{ + Array, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array, + types::{Float16Type, Float32Type, Float64Type}, +}; +use arrow_schema::{DataType, Field}; +use half::{bf16, f16}; +use num_traits::{AsPrimitive, Bounded, Float, FromPrimitive}; + +use super::bfloat16::{BFloat16Array, BFloat16Type}; +use crate::Result; +use crate::bfloat16::is_bfloat16_field; + +/// Float data type. +/// +/// This helps differentiate between the different float types, +/// because bf16 is not officially supported [DataType] in arrow-rs. +#[derive(Debug)] +pub enum FloatType { + BFloat16, + Float16, + Float32, + Float64, +} + +impl std::fmt::Display for FloatType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + Self::BFloat16 => write!(f, "bfloat16"), + Self::Float16 => write!(f, "float16"), + Self::Float32 => write!(f, "float32"), + Self::Float64 => write!(f, "float64"), + } + } +} + +/// Try to convert a [DataType] to a [FloatType]. To support bfloat16, always +/// prefer using the `TryFrom<&Field>` implementation. +impl TryFrom<&DataType> for FloatType { + type Error = crate::ArrowError; + + fn try_from(value: &DataType) -> Result { + match *value { + DataType::Float16 => Ok(Self::Float16), + DataType::Float32 => Ok(Self::Float32), + DataType::Float64 => Ok(Self::Float64), + _ => Err(crate::ArrowError::InvalidArgumentError(format!( + "{:?} is not a floating type", + value + ))), + } + } +} + +impl TryFrom<&Field> for FloatType { + type Error = crate::ArrowError; + + fn try_from(field: &Field) -> Result { + match field.data_type() { + DataType::FixedSizeBinary(2) if is_bfloat16_field(field) => Ok(Self::BFloat16), + _ => Self::try_from(field.data_type()), + } + } +} + +/// Trait for float types used in Lance indexes +/// +/// This mimics the utilities provided by [`arrow_array::ArrowPrimitiveType`] +/// but applies to all float types (including bfloat16) +pub trait ArrowFloatType: Debug { + type Native: FromPrimitive + + FloatToArrayType + + AsPrimitive + + Debug + + Display; + + const FLOAT_TYPE: FloatType; + const MIN: Self::Native; + const MAX: Self::Native; + + /// Arrow Float Array Type. + type ArrayType: FloatArray; + + /// Returns empty array of this type. + fn empty_array() -> Self::ArrayType { + >::from_values(Vec::new()) + } +} + +/// Trait to be implemented by native types that have a corresponding [`ArrowFloatType`] +/// implementation. +/// +/// This helps define what operations are supported by native floats and also helps convert +/// from a native type back to the corresponding Arrow float type. +pub trait FloatToArrayType: + Float + + Bounded + + Sum + + AddAssign + + AsPrimitive + + AsPrimitive + + DivAssign + + Send + + Sync + + Copy +{ + /// The corresponding [`ArrowFloatType`] implementation for this native type + type ArrowType: ArrowFloatType; +} + +impl FloatToArrayType for bf16 { + type ArrowType = BFloat16Type; +} + +impl FloatToArrayType for f16 { + type ArrowType = Float16Type; +} + +impl FloatToArrayType for f32 { + type ArrowType = Float32Type; +} + +impl FloatToArrayType for f64 { + type ArrowType = Float64Type; +} + +impl ArrowFloatType for BFloat16Type { + type Native = bf16; + + const FLOAT_TYPE: FloatType = FloatType::BFloat16; + const MIN: Self::Native = bf16::MIN; + const MAX: Self::Native = bf16::MAX; + + type ArrayType = FixedSizeBinaryArray; +} + +impl ArrowFloatType for Float16Type { + type Native = f16; + + const FLOAT_TYPE: FloatType = FloatType::Float16; + const MIN: Self::Native = f16::MIN; + const MAX: Self::Native = f16::MAX; + + type ArrayType = Float16Array; +} + +impl ArrowFloatType for Float32Type { + type Native = f32; + + const FLOAT_TYPE: FloatType = FloatType::Float32; + const MIN: Self::Native = f32::MIN; + const MAX: Self::Native = f32::MAX; + + type ArrayType = Float32Array; +} + +impl ArrowFloatType for Float64Type { + type Native = f64; + + const FLOAT_TYPE: FloatType = FloatType::Float64; + const MIN: Self::Native = f64::MIN; + const MAX: Self::Native = f64::MAX; + + type ArrayType = Float64Array; +} + +/// [FloatArray] is a trait that is implemented by all float type arrays +/// +/// This is similar to [`arrow_array::PrimitiveArray`] but applies to all float types (including bfloat16) +/// and is implemented as a trait and not a struct +pub trait FloatArray: Array + Clone + 'static { + type FloatType: ArrowFloatType; + + /// Returns a reference to the underlying data as a slice. + /// + /// # Panics + /// + /// Implementations may panic if the array's storage shape does not match + /// the expected element layout. In particular, the `bf16` impl panics if + /// `value_length() != 2` (the `FixedSizeBinary(2)` shape required by + /// `BFloat16Array`). + /// + /// # Preconditions + /// + /// Implementations may impose additional invariants on the underlying + /// buffer. The `bf16` impl requires the value buffer to be at least + /// 2-byte aligned — satisfied automatically by every in-tree Lance + /// constructor, but external callers passing externally-built arrays + /// (FFI, IPC, `Buffer::from_custom_allocation`) must ensure alignment. + /// See the impl's docstring for details. + fn as_slice(&self) -> &[T::Native]; + + /// Construct an array from a vector of values. + fn from_values(values: Vec) -> Self; + + /// Construct an array from an iterator of values. + fn from_iter_values(values: impl IntoIterator) -> Self + where + Self: Sized, + { + Self::from_values(values.into_iter().collect()) + } +} + +impl FloatArray for Float16Array { + type FloatType = Float16Type; + + fn as_slice(&self) -> &[::Native] { + self.values() + } + + fn from_values(values: Vec<::Native>) -> Self { + Self::from(values) + } +} + +impl FloatArray for Float32Array { + type FloatType = Float32Type; + + fn as_slice(&self) -> &[::Native] { + self.values() + } + + fn from_values(values: Vec<::Native>) -> Self { + Self::from(values) + } +} + +impl FloatArray for Float64Array { + type FloatType = Float64Type; + + fn as_slice(&self) -> &[::Native] { + self.values() + } + + fn from_values(values: Vec<::Native>) -> Self { + Self::from(values) + } +} + +/// Convert a float32 array to another float array +/// +/// This is used during queries as query vectors are always provided as float32 arrays +/// and need to be converted to the appropriate float type for the index. +pub fn coerce_float_vector(input: &Float32Array, float_type: FloatType) -> Result> { + match float_type { + FloatType::BFloat16 => Ok(Arc::new( + BFloat16Array::from_iter_values(input.values().iter().map(|v| bf16::from_f32(*v))) + .into_inner(), + )), + FloatType::Float16 => Ok(Arc::new(Float16Array::from_iter_values( + input.values().iter().map(|v| f16::from_f32(*v)), + ))), + FloatType::Float32 => Ok(Arc::new(input.clone())), + FloatType::Float64 => Ok(Arc::new(Float64Array::from_iter_values( + input.values().iter().map(|v| *v as f64), + ))), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_coerce_float_vector_bfloat16() { + let input = Float32Array::from(vec![1.0f32, 2.0, 3.0]); + let array = coerce_float_vector(&input, FloatType::BFloat16).unwrap(); + + assert_eq!(array.data_type(), &DataType::FixedSizeBinary(2)); + + let fixed = array + .as_any() + .downcast_ref::() + .unwrap(); + let expected: Vec = input.values().iter().map(|v| bf16::from_f32(*v)).collect(); + assert_eq!(fixed.as_slice(), expected.as_slice()); + } +} diff --git a/python/vendor/lance-arrow/src/ipc.rs b/python/vendor/lance-arrow/src/ipc.rs new file mode 100644 index 000000000..099c362b4 --- /dev/null +++ b/python/vendor/lance-arrow/src/ipc.rs @@ -0,0 +1,621 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Zero-copy Arrow IPC stream read/write utilities. +//! +//! Provides helpers for serializing and deserializing [`RecordBatch`]es as +//! self-delimiting Arrow IPC streams using synchronous [`Read`]/[`Write`] I/O. +//! +//! These are designed for embedding IPC streams inside larger binary formats +//! (e.g. a cache entry that contains multiple IPC sections). Each stream is +//! self-delimiting (schema + batches + EOS marker) and can be read back +//! independently. +//! +//! # Zero-copy reads +//! +//! [`read_ipc_stream`] and [`read_ipc_stream_single`] take `&Bytes` and use +//! [`Bytes::slice`] to produce each message buffer. Because `Bytes::slice` +//! increments a reference count rather than copying, the resulting +//! [`Buffer`]s — and the array data decoded from them by [`FileDecoder`] — +//! are all backed by the same allocation as the input. + +use std::io::{Read, Write}; +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_buffer::Buffer; +use arrow_ipc::convert::fb_to_schema; +use arrow_ipc::reader::FileDecoder; +use arrow_ipc::root_as_message; +use arrow_ipc::writer::StreamWriter; +use arrow_schema::ArrowError; +use bytes::Bytes; + +// --------------------------------------------------------------------------- +// Length-prefixed byte utilities +// --------------------------------------------------------------------------- + +/// Write `data` prefixed by its length as a little-endian `u64`. +/// +/// Paired with [`read_len_prefixed_bytes`]. +pub fn write_len_prefixed_bytes(writer: &mut dyn Write, data: &[u8]) -> Result<(), ArrowError> { + writer + .write_all(&(data.len() as u64).to_le_bytes()) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + writer + .write_all(data) + .map_err(|e| ArrowError::IoError(e.to_string(), e)) +} + +/// Read a byte slice written by [`write_len_prefixed_bytes`]. +/// +/// Reads an 8-byte little-endian length then exactly that many bytes. +pub fn read_len_prefixed_bytes(reader: &mut dyn Read) -> Result, ArrowError> { + let mut len_buf = [0u8; 8]; + reader + .read_exact(&mut len_buf) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + let len = u64::from_le_bytes(len_buf) as usize; + let mut buf = vec![0u8; len]; + reader + .read_exact(&mut buf) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + Ok(buf) +} + +// --------------------------------------------------------------------------- +// IPC stream utilities +// --------------------------------------------------------------------------- + +// 4-byte continuation marker used by modern Arrow IPC streams. +const IPC_CONTINUATION: [u8; 4] = [0xff; 4]; + +/// Write `batch` as a single-batch Arrow IPC stream to `writer`. +pub fn write_ipc_stream(batch: &RecordBatch, writer: &mut dyn Write) -> Result<(), ArrowError> { + let mut sw = StreamWriter::try_new(&mut *writer, batch.schema_ref())?; + sw.write(batch)?; + sw.finish() +} + +/// Write all batches from `iter` as a single Arrow IPC stream to `writer`. +/// +/// `iter` must yield at least one batch; the schema is inferred from the first +/// batch. Returns `ArrowError::InvalidArgumentError` if the iterator is empty. +/// If you need to write an empty stream (schema only, no rows), construct a +/// `StreamWriter` directly. +pub fn write_ipc_stream_batches(iter: I, writer: &mut dyn Write) -> Result<(), ArrowError> +where + I: IntoIterator, +{ + let mut iter = iter.into_iter(); + let first = iter + .next() + .ok_or_else(|| ArrowError::InvalidArgumentError("no batches to serialize".into()))?; + let mut sw = StreamWriter::try_new(&mut *writer, first.schema_ref())?; + sw.write(&first)?; + for batch in iter { + sw.write(&batch)?; + } + sw.finish() +} + +/// Read one complete Arrow IPC stream message from `data` as a zero-copy [`Buffer`]. +/// +/// Parses the first message starting at byte 0 of `data`. Returns `None` on +/// EOS (size field == 0) or empty input. The returned [`Buffer`] is backed by +/// `data`'s allocation — no bytes are copied. +/// +/// The caller should advance its position by `buf.len()` after each call. +fn read_one_ipc_message(data: &Bytes) -> Result, ArrowError> { + let bytes = data.as_ref(); + + if bytes.is_empty() { + return Ok(None); + } + if bytes.len() < 4 { + return Err(ArrowError::IoError( + "IPC: truncated header".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC header"), + )); + } + + let has_continuation = bytes[..4] == IPC_CONTINUATION; + let (size_bytes, prefix_len): ([u8; 4], usize) = if has_continuation { + if bytes.len() < 8 { + return Err(ArrowError::IoError( + "IPC: truncated header after continuation".into(), + std::io::Error::new( + std::io::ErrorKind::UnexpectedEof, + "truncated after continuation", + ), + )); + } + (bytes[4..8].try_into().unwrap(), 8) + } else { + (bytes[..4].try_into().unwrap(), 4) + }; + + let meta_size = u32::from_le_bytes(size_bytes) as usize; + if meta_size == 0 { + return Ok(None); // EOS + } + + let meta_end = prefix_len + meta_size; + if bytes.len() < meta_end { + return Err(ArrowError::IoError( + "IPC: truncated metadata".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC metadata"), + )); + } + + let msg = root_as_message(&bytes[prefix_len..meta_end]) + .map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?; + let body_len = msg.bodyLength() as usize; + + let total = meta_end + body_len; + if bytes.len() < total { + return Err(ArrowError::IoError( + "IPC: truncated body".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated IPC body"), + )); + } + + // Zero-copy: Bytes::slice shares the backing allocation; Buffer::from + // wraps it without copying. + Ok(Some(Buffer::from(data.slice(0..total)))) +} + +/// Read a length-prefixed byte slice at `offset` in `data`, advancing `offset`. +/// +/// Reads an 8-byte little-endian length, then slices exactly that many bytes +/// from `data`. The returned [`Bytes`] is zero-copy (shares `data`'s allocation). +pub fn read_len_prefixed_bytes_at(data: &Bytes, offset: &mut usize) -> Result { + let bytes = data.as_ref(); + let len_end = offset + .checked_add(8) + .filter(|&e| e <= bytes.len()) + .ok_or_else(|| { + ArrowError::IoError( + "length-prefixed bytes: truncated length field".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated length"), + ) + })?; + let len = u64::from_le_bytes(bytes[*offset..len_end].try_into().unwrap()) as usize; + *offset = len_end; + let data_end = offset + .checked_add(len) + .filter(|&e| e <= bytes.len()) + .ok_or_else(|| { + ArrowError::IoError( + "length-prefixed bytes: truncated data".into(), + std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "truncated data"), + ) + })?; + let result = data.slice(*offset..data_end); + *offset = data_end; + Ok(result) +} + +/// Read all [`RecordBatch`]es from one Arrow IPC stream starting at `offset`, +/// advancing `offset` past the stream (including the EOS marker). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation. +pub fn read_ipc_stream_at( + data: &Bytes, + offset: &mut usize, +) -> Result, ArrowError> { + let batches = read_ipc_stream(&data.slice(*offset..))?; + + // Recompute how many bytes were consumed by re-parsing message sizes. + // We can't get this from read_ipc_stream directly, so we re-walk the + // message headers (metadata only, no body re-read) to sum up lengths. + let slice = &data.as_ref()[*offset..]; + let mut consumed = 0usize; + loop { + let rem = &slice[consumed..]; + if rem.is_empty() { + break; + } + let has_cont = rem.len() >= 4 && rem[..4] == IPC_CONTINUATION; + let (size_bytes, prefix_len): ([u8; 4], usize) = if has_cont { + if rem.len() < 8 { + break; + } + (rem[4..8].try_into().unwrap(), 8) + } else { + if rem.len() < 4 { + break; + } + (rem[..4].try_into().unwrap(), 4) + }; + let meta_size = u32::from_le_bytes(size_bytes) as usize; + if meta_size == 0 { + // EOS — consume it and stop. + consumed += prefix_len; + break; + } + let meta_end = prefix_len + meta_size; + if rem.len() < meta_end { + break; + } + let msg = root_as_message(&rem[prefix_len..meta_end]) + .map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?; + let body_len = msg.bodyLength() as usize; + consumed += meta_end + body_len; + } + *offset += consumed; + + Ok(batches) +} + +/// Read exactly one [`RecordBatch`] from one Arrow IPC stream starting at `offset`, +/// advancing `offset` past the stream (including the EOS marker). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation. +pub fn read_ipc_stream_single_at( + data: &Bytes, + offset: &mut usize, +) -> Result { + let mut batches = read_ipc_stream_at(data, offset)?; + match batches.len() { + 1 => Ok(batches.remove(0)), + n => Err(ArrowError::ParseError(format!( + "expected exactly 1 IPC record batch, got {n}" + ))), + } +} + +/// Extract the prefix length and metadata size from a raw IPC message buffer. +/// +/// Modern IPC streams have an 8-byte prefix `[continuation: 4][size: 4]`. +/// Legacy streams have a 4-byte prefix `[size: 4]`. Returns `(prefix_len, meta_size)`. +fn parse_ipc_message_prefix(buf: &Buffer) -> Result<(usize, usize), ArrowError> { + let has_continuation = buf.len() >= 4 && buf[..4] == IPC_CONTINUATION; + if has_continuation { + if buf.len() < 8 { + return Err(ArrowError::ParseError( + "IPC message buffer too short".into(), + )); + } + let meta_size = u32::from_le_bytes(buf[4..8].try_into().unwrap()) as usize; + Ok((8, meta_size)) + } else { + if buf.len() < 4 { + return Err(ArrowError::ParseError( + "IPC message buffer too short".into(), + )); + } + let meta_size = u32::from_le_bytes(buf[..4].try_into().unwrap()) as usize; + Ok((4, meta_size)) + } +} + +/// Read all [`RecordBatch`]es from one Arrow IPC stream. +/// +/// Zero-copy: each batch's array data buffers are borrowed from the input +/// message buffer(s) and not copied during decoding. +/// +/// Uses [`FileDecoder`] directly (rather than `StreamDecoder`) to avoid a +/// known edge case where `StreamDecoder` does not produce a batch for messages +/// with a zero-length body when the message exactly fills the decode buffer. +pub fn read_ipc_stream(data: &Bytes) -> Result, ArrowError> { + let mut offset = 0usize; + + let schema_buf = read_one_ipc_message(&data.slice(offset..))?.ok_or_else(|| { + ArrowError::ParseError("IPC stream: expected schema message, got EOS".into()) + })?; + offset += schema_buf.len(); + + let (prefix_len, meta_size) = parse_ipc_message_prefix(&schema_buf)?; + let schema_msg = root_as_message(&schema_buf[prefix_len..prefix_len + meta_size]) + .map_err(|e| ArrowError::ParseError(format!("IPC schema parse error: {e}")))?; + let schema = Arc::new(fb_to_schema(schema_msg.header_as_schema().ok_or_else( + || ArrowError::ParseError("IPC stream: first message is not a schema".into()), + )?)); + let mut decoder = FileDecoder::new(schema, schema_msg.version()); + + let mut batches = Vec::new(); + + loop { + let Some(buf) = read_one_ipc_message(&data.slice(offset..))? else { + break; + }; + offset += buf.len(); + + let (prefix_len, meta_size) = parse_ipc_message_prefix(&buf)?; + let msg = root_as_message(&buf[prefix_len..prefix_len + meta_size]) + .map_err(|e| ArrowError::ParseError(format!("IPC message parse error: {e}")))?; + let body_len = msg.bodyLength() as usize; + + // Block offset = 0 since the buffer starts at the message boundary. + // metaDataLength = prefix_len + meta_size (prefix + flatbuf + padding). + let block = arrow_ipc::Block::new(0, (prefix_len + meta_size) as i32, body_len as i64); + + match msg.header_type() { + arrow_ipc::MessageHeader::RecordBatch => { + if let Some(batch) = decoder.read_record_batch(&block, &buf)? { + batches.push(batch); + } + } + arrow_ipc::MessageHeader::DictionaryBatch => { + decoder.read_dictionary(&block, &buf)?; + } + _ => break, + } + } + + Ok(batches) +} + +/// Read exactly one [`RecordBatch`] from one Arrow IPC stream. +pub fn read_ipc_stream_single(data: &Bytes) -> Result { + let mut batches = read_ipc_stream(data)?; + match batches.len() { + 1 => Ok(batches.remove(0)), + n => Err(ArrowError::ParseError(format!( + "expected exactly 1 IPC record batch, got {n}" + ))), + } +} + +// --------------------------------------------------------------------------- +// Aligned IPC sections +// --------------------------------------------------------------------------- + +/// Byte alignment that each IPC section's stream start is padded to. +/// +/// When several IPC streams are concatenated into one larger blob (e.g. a +/// cache entry), a section that starts at an arbitrary offset would leave its +/// array data misaligned. [`FileDecoder`] with `require_alignment = false` +/// then silently copies each buffer into a freshly aligned allocation on +/// every read, defeating zero-copy. Padding each section start to a 64-byte +/// boundary keeps the decoded buffers borrowed directly from the input. +pub const IPC_SECTION_ALIGNMENT: usize = 64; + +/// Number of zero-padding bytes needed to advance `pos` to the next +/// [`IPC_SECTION_ALIGNMENT`] boundary. +fn section_padding(pos: usize) -> usize { + (IPC_SECTION_ALIGNMENT - (pos % IPC_SECTION_ALIGNMENT)) % IPC_SECTION_ALIGNMENT +} + +/// A [`Write`] adapter that counts the bytes written through it. +struct CountingWriter<'a> { + inner: &'a mut dyn Write, + count: usize, +} + +impl Write for CountingWriter<'_> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let n = self.inner.write(buf)?; + self.count += n; + Ok(n) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +/// Write zero padding so the next byte lands on an [`IPC_SECTION_ALIGNMENT`] +/// boundary, advancing `pos` past it. +fn write_section_padding(writer: &mut dyn Write, pos: &mut usize) -> Result<(), ArrowError> { + let pad = section_padding(*pos); + if pad > 0 { + const ZEROS: [u8; IPC_SECTION_ALIGNMENT] = [0u8; IPC_SECTION_ALIGNMENT]; + writer + .write_all(&ZEROS[..pad]) + .map_err(|e| ArrowError::IoError(e.to_string(), e))?; + *pos += pad; + } + Ok(()) +} + +/// Write `batch` as a 64-byte-aligned single-batch Arrow IPC section. +/// +/// `pos` is the absolute byte offset of `writer` within the enclosing blob. +/// Zero padding is written first so the IPC stream begins on an +/// [`IPC_SECTION_ALIGNMENT`] boundary, then the stream itself. `pos` is +/// advanced past both the padding and the stream so the caller can write +/// further aligned sections. +/// +/// Paired with [`read_ipc_section_at`]. For the decoded buffers to be borrowed +/// zero-copy, the blob must ultimately be read back from a buffer whose base +/// address is at least 64-byte aligned. +pub fn write_ipc_section( + writer: &mut dyn Write, + pos: &mut usize, + batch: &RecordBatch, +) -> Result<(), ArrowError> { + write_section_padding(writer, pos)?; + + let mut counting = CountingWriter { + inner: writer, + count: 0, + }; + write_ipc_stream(batch, &mut counting)?; + *pos += counting.count; + Ok(()) +} + +/// Read a single [`RecordBatch`] from an aligned IPC section at `offset`. +/// +/// Skips the alignment padding written by [`write_ipc_section`], then reads +/// the stream, advancing `offset` past the section (padding + stream + EOS). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base +/// address is at least 64-byte aligned (see [`write_ipc_section`]). +pub fn read_ipc_section_at(data: &Bytes, offset: &mut usize) -> Result { + *offset += section_padding(*offset); + read_ipc_stream_single_at(data, offset) +} + +/// Write `batches` as a single 64-byte-aligned multi-batch Arrow IPC section. +/// +/// Like [`write_ipc_section`] but emits every batch from `iter` into one IPC +/// stream (schema + N batches + EOS). `iter` must yield at least one batch. +/// Paired with [`read_ipc_section_batches_at`]. +pub fn write_ipc_section_batches( + writer: &mut dyn Write, + pos: &mut usize, + iter: I, +) -> Result<(), ArrowError> +where + I: IntoIterator, +{ + write_section_padding(writer, pos)?; + + let mut counting = CountingWriter { + inner: writer, + count: 0, + }; + write_ipc_stream_batches(iter, &mut counting)?; + *pos += counting.count; + Ok(()) +} + +/// Read all [`RecordBatch`]es from an aligned multi-batch IPC section at +/// `offset`, advancing `offset` past the section (padding + stream + EOS). +/// +/// Zero-copy: array buffers borrow from `data`'s allocation when `data`'s base +/// address is at least 64-byte aligned (see [`write_ipc_section_batches`]). +pub fn read_ipc_section_batches_at( + data: &Bytes, + offset: &mut usize, +) -> Result, ArrowError> { + *offset += section_padding(*offset); + read_ipc_stream_at(data, offset) +} + +#[cfg(test)] +mod tests { + use arrow_array::{ArrayRef, record_batch}; + + use super::*; + + #[test] + fn test_ipc_roundtrip() { + let batch1 = record_batch!( + ("int", Int32, [1, 2, 3]), + ("str", Utf8, ["foo", "bar", "baz"]) + ) + .unwrap(); + let batch2 = record_batch!(("int", Int32, [4, 5]), ("str", Utf8, ["qux", "quux"])).unwrap(); + let batches = vec![batch1.clone(), batch2.clone()]; + + let mut buf = Vec::new(); + write_ipc_stream_batches(batches, &mut buf).unwrap(); + + let data = Bytes::from(buf); + + let batches = read_ipc_stream(&data).unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(batches[0], batch1); + assert_eq!(batches[1], batch2); + + let data_base = data.as_ptr() as usize; + let data_end = data_base + data.len(); + let assert_col_zero_copy = |array: &ArrayRef| { + for buffer in array.to_data().buffers() { + let ptr = buffer.as_ptr() as usize; + assert!( + ptr >= data_base && ptr < data_end, + "buffer at {ptr:#x} is not backed by the input Bytes allocation \ + [{data_base:#x}..{data_end:#x})" + ); + } + }; + + for batch in &batches { + assert_eq!(batch.schema(), batch1.schema()); + assert_col_zero_copy(batch.column(0)); + assert_col_zero_copy(batch.column(1)); + } + } + + /// Allocate a [`Bytes`] whose base address is 64-byte aligned, modelling a + /// backend that reads cache entries into an aligned buffer. A plain + /// `Bytes::from(vec)` only guarantees the allocator's alignment for `u8`. + fn aligned_bytes(payload: &[u8]) -> Bytes { + let mut v = vec![0u8; payload.len() + IPC_SECTION_ALIGNMENT]; + let pad = section_padding(v.as_ptr() as usize); + v[pad..pad + payload.len()].copy_from_slice(payload); + Bytes::from(v).slice(pad..pad + payload.len()) + } + + #[test] + fn test_aligned_ipc_sections_are_zero_copy() { + // A LargeBinary column exercises the i64-offset buffer whose 8-byte + // alignment requirement triggers a realigning memcpy when misaligned. + let blocks = arrow_array::LargeBinaryArray::from_vec(vec![&b"hello"[..], b"world"]); + let section_a = RecordBatch::try_from_iter([("a", Arc::new(blocks) as ArrayRef)]).unwrap(); + let section_b = record_batch!(("b", Int64, [10i64, 20, 30, 40, 50])).unwrap(); + + let mut buf = Vec::new(); + // Arbitrary, deliberately non-64-aligned preamble so the first section + // must be padded rather than landing at offset 0 by luck. + buf.extend_from_slice(&[0xABu8; 7]); + let mut pos = buf.len(); + // The first section's stream begins after padding the 7-byte preamble + // up to the next 64-byte boundary. + assert_eq!(7 + section_padding(7), IPC_SECTION_ALIGNMENT); + write_ipc_section(&mut buf, &mut pos, §ion_a).unwrap(); + write_ipc_section(&mut buf, &mut pos, §ion_b).unwrap(); + + let data = aligned_bytes(&buf); + assert_eq!( + section_padding(data.as_ptr() as usize), + 0, + "base not aligned" + ); + + let mut offset = 7; + let read_a = read_ipc_section_at(&data, &mut offset).unwrap(); + let read_b = read_ipc_section_at(&data, &mut offset).unwrap(); + assert_eq!(read_a, section_a); + assert_eq!(read_b, section_b); + + let data_base = data.as_ptr() as usize; + let data_end = data_base + data.len(); + for batch in [&read_a, &read_b] { + for buffer in batch.column(0).to_data().buffers() { + let ptr = buffer.as_ptr() as usize; + assert!( + ptr >= data_base && ptr < data_end, + "section buffer at {ptr:#x} was realigned out of the input \ + [{data_base:#x}..{data_end:#x}) — misaligned section", + ); + } + } + } + + #[test] + fn test_aligned_multi_batch_section_roundtrip_zero_copy() { + // A multi-batch section (e.g. IVF SQ storage chunks) must round-trip + // every batch and decode the first batch's buffers zero-copy. + let b1 = record_batch!(("v", Int64, [1i64, 2, 3])).unwrap(); + let b2 = record_batch!(("v", Int64, [4i64, 5])).unwrap(); + let b3 = record_batch!(("v", Int64, [6i64])).unwrap(); + + let mut buf = vec![0xCDu8; 5]; + let mut pos = buf.len(); + write_ipc_section_batches(&mut buf, &mut pos, [b1.clone(), b2.clone(), b3.clone()]) + .unwrap(); + + let data = aligned_bytes(&buf); + let mut offset = 5; + let read = read_ipc_section_batches_at(&data, &mut offset).unwrap(); + assert_eq!(read, vec![b1, b2, b3]); + assert_eq!(offset, buf.len(), "offset should land at section end"); + + let data_base = data.as_ptr() as usize; + let data_end = data_base + data.len(); + for buffer in read[0].column(0).to_data().buffers() { + let ptr = buffer.as_ptr() as usize; + assert!( + ptr >= data_base && ptr < data_end, + "first batch buffer at {ptr:#x} was realigned out of the input", + ); + } + } +} diff --git a/python/vendor/lance-arrow/src/json.rs b/python/vendor/lance-arrow/src/json.rs new file mode 100644 index 000000000..df908c4cb --- /dev/null +++ b/python/vendor/lance-arrow/src/json.rs @@ -0,0 +1,1538 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! JSON support for Apache Arrow. + +use std::convert::TryFrom; +use std::sync::Arc; + +use arrow_array::builder::{GenericStringBuilder, LargeBinaryBuilder}; +use arrow_array::cast::AsArray; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, GenericStringArray, LargeBinaryArray, LargeListArray, + LargeStringArray, ListArray, MapArray, OffsetSizeTrait, RecordBatch, StringArray, StructArray, +}; +use arrow_schema::{ArrowError, DataType, Field as ArrowField, Fields, Schema}; +use lance_jsonb::OwnedJsonb; +use lance_jsonb::jsonpath::{JsonPath, Selector}; + +use crate::ARROW_EXT_NAME_KEY; + +/// Arrow extension type name for JSON data (Lance internal) +pub const JSON_EXT_NAME: &str = "lance.json"; + +/// Arrow extension type name for JSON data (Arrow official) +pub const ARROW_JSON_EXT_NAME: &str = "arrow.json"; + +/// Check if a field is a JSON extension field (Lance internal JSONB storage) +pub fn is_json_field(field: &ArrowField) -> bool { + field.data_type() == &DataType::LargeBinary + && field + .metadata() + .get(ARROW_EXT_NAME_KEY) + .map(|name| name == JSON_EXT_NAME) + .unwrap_or_default() +} + +/// Check if a field is an Arrow JSON extension field (PyArrow pa.json() type) +pub fn is_arrow_json_field(field: &ArrowField) -> bool { + // Arrow JSON extension type uses Utf8 or LargeUtf8 as storage type + (field.data_type() == &DataType::Utf8 || field.data_type() == &DataType::LargeUtf8) + && field + .metadata() + .get(ARROW_EXT_NAME_KEY) + .map(|name| name == ARROW_JSON_EXT_NAME) + .unwrap_or_default() +} + +/// Check if a field or any of its descendants is a JSON field +pub fn has_json_fields(field: &ArrowField) -> bool { + if is_json_field(field) { + return true; + } + + match field.data_type() { + DataType::Struct(fields) => fields.iter().any(|f| has_json_fields(f)), + DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => { + has_json_fields(f) + } + DataType::Map(f, _) => has_json_fields(f), + _ => false, + } +} + +/// Check if a field or any of its descendants is an Arrow JSON field +pub fn has_arrow_json_fields(field: &ArrowField) -> bool { + if is_arrow_json_field(field) { + return true; + } + + match field.data_type() { + DataType::Struct(fields) => fields.iter().any(|f| has_arrow_json_fields(f)), + DataType::List(f) | DataType::LargeList(f) | DataType::FixedSizeList(f, _) => { + has_arrow_json_fields(f) + } + DataType::Map(f, _) => has_arrow_json_fields(f), + _ => false, + } +} + +/// The Arrow representation holding a JSON field's values. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum JsonEncoding { + /// JSONB bytes: [`JSON_EXT_NAME`] over `LargeBinary`. Lance stores JSON + /// this way. + Jsonb, + /// JSON text: [`ARROW_JSON_EXT_NAME`] over `Utf8` or `LargeUtf8`. Callers + /// write JSON this way, and reads return it this way by default. + Text, +} + +impl JsonEncoding { + /// The encoding of `field`'s values, or `None` when `field` does not hold + /// JSON. + pub fn of_field(field: &ArrowField) -> Option { + if is_json_field(field) { + Some(Self::Jsonb) + } else if is_arrow_json_field(field) { + Some(Self::Text) + } else { + None + } + } +} + +/// The values of a JSON column, readable as JSONB or as text whichever +/// encoding they arrived in. +/// +/// Code that computes on JSON, such as tokenizers and indices, reads values +/// through this type, so its behavior does not depend on whether they come +/// from storage or from a caller. +#[derive(Debug, Clone, Copy)] +pub struct JsonValues<'a> { + encoding: JsonEncoding, + array: &'a ArrayRef, +} + +impl<'a> JsonValues<'a> { + /// Read `array` as the values of `field`, or `None` when `field` does not + /// hold JSON. + pub fn try_new(field: &ArrowField, array: &'a ArrayRef) -> Option { + JsonEncoding::of_field(field).map(|encoding| Self { encoding, array }) + } + + pub fn encoding(&self) -> JsonEncoding { + self.encoding + } + + /// The values as JSONB. Text values are parsed, so invalid JSON fails. + pub fn to_jsonb(&self) -> Result { + match self.encoding { + JsonEncoding::Jsonb => Ok(self.array.as_binary::().clone()), + JsonEncoding::Text => Ok(JsonArray::try_from(self.array.clone())?.into_inner()), + } + } + + /// The values as JSON text: the input array when it already holds text, + /// otherwise `LargeUtf8` decoded from JSONB. + pub fn to_text(&self) -> ArrayRef { + match self.encoding { + JsonEncoding::Jsonb => Arc::new(decode_jsonb_array::(self.array.as_binary())), + JsonEncoding::Text => self.array.clone(), + } + } +} + +/// Create a JSON field with the appropriate extension metadata +pub fn json_field(name: &str, nullable: bool) -> ArrowField { + let mut field = ArrowField::new(name, DataType::LargeBinary, nullable); + let mut metadata = std::collections::HashMap::new(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), JSON_EXT_NAME.to_string()); + field.set_metadata(metadata); + field +} + +/// A specialized array for JSON data stored as JSONB binary format +#[derive(Debug, Clone)] +pub struct JsonArray { + inner: LargeBinaryArray, +} + +impl JsonArray { + /// Create a new JsonArray from an iterator of JSON strings + pub fn try_from_iter(iter: I) -> Result + where + I: IntoIterator>, + S: AsRef, + { + let mut builder = LargeBinaryBuilder::new(); + + for json_str in iter { + match json_str { + Some(s) => { + let encoded = encode_json(s.as_ref()).map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e)) + })?; + builder.append_value(&encoded); + } + None => builder.append_null(), + } + } + + Ok(Self { + inner: builder.finish(), + }) + } + + /// Get the underlying LargeBinaryArray + pub fn into_inner(self) -> LargeBinaryArray { + self.inner + } + + /// Get a reference to the underlying LargeBinaryArray + pub fn inner(&self) -> &LargeBinaryArray { + &self.inner + } + + /// Get the value at index i as decoded JSON string + pub fn value(&self, i: usize) -> Result { + if self.inner.is_null(i) { + return Err(ArrowError::InvalidArgumentError( + "Value is null".to_string(), + )); + } + + let jsonb_bytes = self.inner.value(i); + Ok(decode_json(jsonb_bytes)) + } + + /// Get the value at index i as raw JSONB bytes + pub fn value_bytes(&self, i: usize) -> &[u8] { + self.inner.value(i) + } + + /// Get JSONPath value from the JSON at index i + pub fn json_path(&self, i: usize, path: &str) -> Result, ArrowError> { + if self.inner.is_null(i) { + return Ok(None); + } + + let path = parse_json_path(path)?; + Ok(select_json_path(self.inner.value(i), &path)?.map(|value| value.to_string())) + } + + /// Convert to Arrow string array (JSON as UTF-8) + pub fn to_arrow_json(&self) -> ArrayRef { + Arc::new(decode_jsonb_array::(&self.inner)) + } + + pub fn len(&self) -> usize { + self.inner.len() + } + + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + pub fn is_null(&self, i: usize) -> bool { + self.inner.is_null(i) + } +} + +// TryFrom implementations for string arrays +impl TryFrom for JsonArray { + type Error = ArrowError; + + fn try_from(array: StringArray) -> Result { + Self::try_from(&array) + } +} + +impl TryFrom<&StringArray> for JsonArray { + type Error = ArrowError; + + fn try_from(array: &StringArray) -> Result { + encode_json_text_array(array) + } +} + +impl TryFrom for JsonArray { + type Error = ArrowError; + + fn try_from(array: LargeStringArray) -> Result { + Self::try_from(&array) + } +} + +impl TryFrom<&LargeStringArray> for JsonArray { + type Error = ArrowError; + + fn try_from(array: &LargeStringArray) -> Result { + encode_json_text_array(array) + } +} + +fn encode_json_text_array( + array: &GenericStringArray, +) -> Result { + let mut builder = LargeBinaryBuilder::with_capacity(array.len(), array.value_data().len()); + for value in array { + match value { + Some(json_str) => { + let encoded = encode_json(json_str).map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to encode JSON: {}", e)) + })?; + builder.append_value(&encoded); + } + None => builder.append_null(), + } + } + Ok(JsonArray { + inner: builder.finish(), + }) +} + +fn decode_jsonb_array(array: &LargeBinaryArray) -> GenericStringArray { + let mut builder = + GenericStringBuilder::::with_capacity(array.len(), array.value_data().len()); + for value in array { + match value { + Some(jsonb_bytes) => builder.append_value(decode_json(jsonb_bytes)), + None => builder.append_null(), + } + } + builder.finish() +} + +impl TryFrom for JsonArray { + type Error = ArrowError; + + fn try_from(array_ref: ArrayRef) -> Result { + match array_ref.data_type() { + DataType::Utf8 => { + // Downcast is guaranteed to succeed after matching on DataType::Utf8 + let string_array = array_ref + .as_any() + .downcast_ref::() + .expect("DataType::Utf8 array must be StringArray"); + Self::try_from(string_array) + } + DataType::LargeUtf8 => { + // Downcast is guaranteed to succeed after matching on DataType::LargeUtf8 + let large_string_array = array_ref + .as_any() + .downcast_ref::() + .expect("DataType::LargeUtf8 array must be LargeStringArray"); + Self::try_from(large_string_array) + } + dt => Err(ArrowError::InvalidArgumentError(format!( + "Unsupported array type for JSON: {:?}. Expected Utf8 or LargeUtf8", + dt + ))), + } + } +} + +/// Encode JSON string to JSONB format +pub fn encode_json(json_str: &str) -> Result, Box> { + let value = lance_jsonb::parse_value(json_str.as_bytes())?; + Ok(value.to_vec()) +} + +/// Decode JSONB bytes to JSON string +pub fn decode_json(jsonb_bytes: &[u8]) -> String { + let raw_jsonb = lance_jsonb::RawJsonb::new(jsonb_bytes); + raw_jsonb.to_string() +} + +/// Parse a JSONPath expression. +pub fn parse_json_path(path: &str) -> Result, ArrowError> { + lance_jsonb::jsonpath::parse_json_path(path.as_bytes()) + .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid JSONPath '{path}': {e}"))) +} + +/// Select the value at `path` in a JSONB value. +/// +/// A path matching several values selects them together as one JSON array, +/// so no match is silently dropped. +pub fn select_json_path( + jsonb_bytes: &[u8], + path: &JsonPath<'_>, +) -> Result, ArrowError> { + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) + .select_value(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) + }) +} + +/// Whether `path` matches any value in a JSONB value. +pub fn json_path_exists(jsonb_bytes: &[u8], path: &JsonPath<'_>) -> Result { + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) + .exists(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to match JSONPath {path}: {e}")) + }) +} + +/// Select every value matched by `path` in a JSONB value. +pub fn select_json_path_values( + jsonb_bytes: &[u8], + path: &JsonPath<'_>, +) -> Result, ArrowError> { + Selector::new(lance_jsonb::RawJsonb::new(jsonb_bytes)) + .select_values(path) + .map_err(|e| { + ArrowError::InvalidArgumentError(format!("Failed to select JSONPath {path}: {e}")) + }) +} + +/// Convert an Arrow JSON field to Lance JSON field (with JSONB storage) +pub fn arrow_json_to_lance_json(field: &ArrowField) -> ArrowField { + if is_arrow_json_field(field) { + return field_with_extension(field, DataType::LargeBinary, JSON_EXT_NAME); + } + + let data_type = match field.data_type() { + DataType::Struct(fields) => { + let fields = fields + .iter() + .map(|field| Arc::new(arrow_json_to_lance_json(field))) + .collect::>(); + DataType::Struct(Fields::from(fields)) + } + DataType::List(item) => DataType::List(Arc::new(arrow_json_to_lance_json(item))), + DataType::LargeList(item) => DataType::LargeList(Arc::new(arrow_json_to_lance_json(item))), + DataType::FixedSizeList(item, size) => { + DataType::FixedSizeList(Arc::new(arrow_json_to_lance_json(item)), *size) + } + DataType::Map(entries, keys_sorted) => { + DataType::Map(Arc::new(arrow_json_to_lance_json(entries)), *keys_sorted) + } + _ => return field.clone(), + }; + + field_with_data_type(field, data_type) +} + +/// Convert a Lance JSON field to Arrow JSON field. +pub fn lance_json_to_arrow_json(field: &ArrowField) -> ArrowField { + if is_json_field(field) { + return field_with_extension(field, DataType::Utf8, ARROW_JSON_EXT_NAME); + } + + let data_type = match field.data_type() { + DataType::Struct(fields) => { + let fields = fields + .iter() + .map(|field| Arc::new(lance_json_to_arrow_json(field))) + .collect::>(); + DataType::Struct(Fields::from(fields)) + } + DataType::List(item) => DataType::List(Arc::new(lance_json_to_arrow_json(item))), + DataType::LargeList(item) => DataType::LargeList(Arc::new(lance_json_to_arrow_json(item))), + DataType::FixedSizeList(item, size) => { + DataType::FixedSizeList(Arc::new(lance_json_to_arrow_json(item)), *size) + } + DataType::Map(entries, keys_sorted) => { + DataType::Map(Arc::new(lance_json_to_arrow_json(entries)), *keys_sorted) + } + _ => return field.clone(), + }; + + field_with_data_type(field, data_type) +} + +fn field_with_data_type(field: &ArrowField, data_type: DataType) -> ArrowField { + ArrowField::new(field.name(), data_type, field.is_nullable()) + .with_metadata(field.metadata().clone()) +} + +fn field_with_extension( + field: &ArrowField, + data_type: DataType, + extension_name: &str, +) -> ArrowField { + let mut metadata = field.metadata().clone(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), extension_name.to_string()); + ArrowField::new(field.name(), data_type, field.is_nullable()).with_metadata(metadata) +} + +fn convert_json_array( + field: &ArrowField, + array: &ArrayRef, + convert_leaf: &F, +) -> Result<(ArrowField, ArrayRef, bool), ArrowError> +where + F: Fn(&ArrowField, &ArrayRef) -> Result, ArrowError>, +{ + if let Some((field, array)) = convert_leaf(field, array)? { + return Ok((field, array, true)); + } + + match field.data_type() { + DataType::Struct(fields) => { + let struct_array = array.as_struct(); + let mut new_fields = Vec::with_capacity(fields.len()); + let mut new_columns = Vec::with_capacity(fields.len()); + let mut changed = false; + + for (field, column) in fields.iter().zip(struct_array.columns()) { + let (new_field, new_column, field_changed) = + convert_json_array(field, column, convert_leaf)?; + changed |= field_changed; + new_fields.push(Arc::new(new_field)); + new_columns.push(new_column); + } + + if changed { + let fields = Fields::from(new_fields); + let new_field = field_with_data_type(field, DataType::Struct(fields.clone())); + let new_array = + StructArray::new(fields, new_columns, struct_array.nulls().cloned()); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::List(item) => { + let list_array: &ListArray = array.as_list(); + let (new_item, new_values, changed) = + convert_json_array(item, list_array.values(), convert_leaf)?; + if changed { + let new_field = + field_with_data_type(field, DataType::List(Arc::new(new_item.clone()))); + let new_array = ListArray::new( + Arc::new(new_item), + list_array.offsets().clone(), + new_values, + list_array.nulls().cloned(), + ); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::LargeList(item) => { + let list_array: &LargeListArray = array.as_list(); + let (new_item, new_values, changed) = + convert_json_array(item, list_array.values(), convert_leaf)?; + if changed { + let new_field = + field_with_data_type(field, DataType::LargeList(Arc::new(new_item.clone()))); + let new_array = LargeListArray::new( + Arc::new(new_item), + list_array.offsets().clone(), + new_values, + list_array.nulls().cloned(), + ); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::FixedSizeList(item, size) => { + let list_array: &FixedSizeListArray = array.as_fixed_size_list(); + let (new_item, new_values, changed) = + convert_json_array(item, list_array.values(), convert_leaf)?; + if changed { + let new_field = field_with_data_type( + field, + DataType::FixedSizeList(Arc::new(new_item.clone()), *size), + ); + let new_array = FixedSizeListArray::try_new_with_length( + Arc::new(new_item), + *size, + new_values, + list_array.nulls().cloned(), + list_array.len(), + )?; + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + DataType::Map(entries, keys_sorted) => { + let map_array = array + .as_any() + .downcast_ref::() + .expect("DataType::Map array must be MapArray"); + let entries_array = Arc::new(map_array.entries().clone()) as ArrayRef; + let (new_entries, new_entries_array, changed) = + convert_json_array(entries, &entries_array, convert_leaf)?; + if changed { + let entries_struct = new_entries_array + .as_any() + .downcast_ref::() + .expect("Map entries must be StructArray") + .clone(); + let new_field = field_with_data_type( + field, + DataType::Map(Arc::new(new_entries.clone()), *keys_sorted), + ); + let new_array = MapArray::new( + Arc::new(new_entries), + map_array.offsets().clone(), + entries_struct, + map_array.nulls().cloned(), + *keys_sorted, + ); + Ok((new_field, Arc::new(new_array) as ArrayRef, true)) + } else { + Ok((field.clone(), array.clone(), false)) + } + } + _ => Ok((field.clone(), array.clone(), false)), + } +} + +fn convert_arrow_json_array( + field: &ArrowField, + array: &ArrayRef, +) -> Result<(ArrowField, ArrayRef, bool), ArrowError> { + convert_json_array(field, array, &|field, array| { + if is_arrow_json_field(field) { + let json_array = JsonArray::try_from(array.clone())?; + Ok(Some(( + arrow_json_to_lance_json(field), + Arc::new(json_array.into_inner()) as ArrayRef, + ))) + } else { + Ok(None) + } + }) +} + +fn convert_lance_json_array( + field: &ArrowField, + array: &ArrayRef, +) -> Result<(ArrowField, ArrayRef, bool), ArrowError> { + convert_json_array(field, array, &|field, array| { + if is_json_field(field) { + Ok(Some(( + lance_json_to_arrow_json(field), + Arc::new(decode_jsonb_array::(array.as_binary())) as ArrayRef, + ))) + } else { + Ok(None) + } + }) +} + +/// Convert a RecordBatch with Lance JSON columns (JSONB) back to Arrow JSON format (strings) +pub fn convert_lance_json_to_arrow( + batch: &arrow_array::RecordBatch, +) -> Result { + let schema = batch.schema(); + let mut needs_conversion = false; + let mut new_fields = Vec::with_capacity(schema.fields().len()); + let mut new_columns = Vec::with_capacity(batch.num_columns()); + + for (i, field) in schema.fields().iter().enumerate() { + let column = batch.column(i); + let (new_field, new_column, changed) = convert_lance_json_array(field, column)?; + + needs_conversion |= changed; + new_fields.push(new_field); + new_columns.push(new_column); + } + + if needs_conversion { + let new_schema = Arc::new(Schema::new_with_metadata( + new_fields, + schema.metadata().clone(), + )); + RecordBatch::try_new(new_schema, new_columns) + } else { + // No conversion needed, return original batch + Ok(batch.clone()) + } +} + +/// Convert a RecordBatch with Arrow JSON columns to Lance JSON format (JSONB) +pub fn convert_json_columns( + batch: &arrow_array::RecordBatch, +) -> Result { + let schema = batch.schema(); + let mut needs_conversion = false; + let mut new_fields = Vec::with_capacity(schema.fields().len()); + let mut new_columns = Vec::with_capacity(batch.num_columns()); + + for (i, field) in schema.fields().iter().enumerate() { + let column = batch.column(i); + let (new_field, new_column, changed) = convert_arrow_json_array(field, column)?; + + needs_conversion |= changed; + new_fields.push(new_field); + new_columns.push(new_column); + } + + if needs_conversion { + let new_schema = Arc::new(Schema::new_with_metadata( + new_fields, + schema.metadata().clone(), + )); + RecordBatch::try_new(new_schema, new_columns) + } else { + // No conversion needed, return original batch + Ok(batch.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_json_field_creation() { + let field = json_field("data", true); + assert_eq!(field.name(), "data"); + assert_eq!(field.data_type(), &DataType::LargeBinary); + assert!(field.is_nullable()); + assert!(is_json_field(&field)); + } + + #[test] + fn test_json_array_from_strings() { + let json_strings = vec![ + Some(r#"{"name": "Alice", "age": 30}"#), + None, + Some(r#"{"name": "Bob", "age": 25}"#), + ]; + + let array = JsonArray::try_from_iter(json_strings).unwrap(); + assert_eq!(array.len(), 3); + assert!(!array.is_null(0)); + assert!(array.is_null(1)); + assert!(!array.is_null(2)); + + let decoded = array.value(0).unwrap(); + assert!(decoded.contains("Alice")); + } + + #[test] + fn test_json_array_from_string_array() { + let string_array = StringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + None, + ]); + + let json_array = JsonArray::try_from(string_array).unwrap(); + assert_eq!(json_array.len(), 3); + assert!(!json_array.is_null(0)); + assert!(!json_array.is_null(1)); + assert!(json_array.is_null(2)); + } + + #[test] + fn test_json_path_extraction() { + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"user": {"name": "Alice", "age": 30}}"#), + Some(r#"{"user": {"name": "Bob"}}"#), + ]) + .unwrap(); + + let name = json_array.json_path(0, "$.user.name").unwrap(); + assert_eq!(name, Some("\"Alice\"".to_string())); + + let age = json_array.json_path(1, "$.user.age").unwrap(); + assert_eq!(age, None); + + // Several matches are selected together, as the SQL JSON functions do. + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"a": [{"b": 1}, {"b": 2}]}"#)]).unwrap(); + let bs = json_array.json_path(0, "$.a[*].b").unwrap(); + assert_eq!(bs, Some("[1,2]".to_string())); + } + + #[test] + fn test_json_values_read_either_encoding() { + let docs = vec![Some(r#"{"a":1}"#), None]; + let text: ArrayRef = Arc::new(StringArray::from(docs.clone())); + let jsonb: ArrayRef = Arc::new(JsonArray::try_from_iter(docs).unwrap().into_inner()); + let mut text_field = ArrowField::new("j", DataType::Utf8, true); + text_field.set_metadata(std::collections::HashMap::from([( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + )])); + let jsonb_field = json_field("j", true); + + let from_text = JsonValues::try_new(&text_field, &text).unwrap(); + let from_jsonb = JsonValues::try_new(&jsonb_field, &jsonb).unwrap(); + assert_eq!(from_text.encoding(), JsonEncoding::Text); + assert_eq!(from_jsonb.encoding(), JsonEncoding::Jsonb); + + assert_eq!(&from_text.to_jsonb().unwrap(), jsonb.as_binary::()); + assert_eq!(&from_jsonb.to_jsonb().unwrap(), jsonb.as_binary::()); + assert_eq!(&from_text.to_text(), &text); + let decoded = from_jsonb.to_text(); + let decoded = decoded.as_string::(); + assert_eq!(decoded.value(0), r#"{"a":1}"#); + assert!(decoded.is_null(1)); + + let plain = ArrowField::new("j", DataType::Utf8, true); + assert!(JsonValues::try_new(&plain, &text).is_none()); + } + + #[test] + fn test_convert_json_columns() { + // Create a batch with Arrow JSON column + let json_strings = vec![Some(r#"{"name": "Alice"}"#), Some(r#"{"name": "Bob"}"#)]; + let json_arr = StringArray::from(json_strings); + + // Create field with arrow.json extension + let mut field = ArrowField::new("data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(json_arr) as ArrayRef]).unwrap(); + + // Convert the batch + let converted = convert_json_columns(&batch).unwrap(); + + // Check the converted schema + assert_eq!(converted.num_columns(), 1); + let converted_schema = converted.schema(); + let converted_field = converted_schema.field(0); + assert_eq!(converted_field.data_type(), &DataType::LargeBinary); + assert_eq!( + converted_field.metadata().get(ARROW_EXT_NAME_KEY), + Some(&JSON_EXT_NAME.to_string()) + ); + + // Check the data was converted + let converted_column = converted.column(0); + assert_eq!(converted_column.data_type(), &DataType::LargeBinary); + assert_eq!(converted_column.len(), 2); + + // Verify the data is valid JSONB + let binary_array = converted_column + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..binary_array.len() { + let jsonb_bytes = binary_array.value(i); + let decoded = decode_json(jsonb_bytes); + assert!(decoded.contains("name")); + } + } + + #[test] + fn test_convert_nested_json_columns() { + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + + let uri_field = Arc::new(ArrowField::new("uri", DataType::Utf8, false)); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let extra_field = + Arc::new(ArrowField::new("extra", DataType::Utf8, true).with_metadata(metadata)); + let item_fields = Fields::from(vec![uri_field, extra_field]); + + let values = StructArray::new( + item_fields.clone(), + vec![ + Arc::new(StringArray::from(vec![Some("a.jpg"), Some("b.jpg")])) as ArrayRef, + Arc::new(StringArray::from(vec![ + Some(r#"{"codec":"h264"}"#), + None::<&str>, + ])) as ArrayRef, + ], + None, + ); + let item = Arc::new(ArrowField::new("item", DataType::Struct(item_fields), true)); + let media = ListArray::new( + item, + OffsetBuffer::new(ScalarBuffer::from(vec![0, 1, 2])), + Arc::new(values), + None, + ); + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "media", + media.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(media) as ArrayRef]).unwrap(); + + assert!(has_arrow_json_fields(batch.schema().field(0))); + + let converted = convert_json_columns(&batch).unwrap(); + let converted_schema = converted.schema(); + let DataType::List(item) = converted_schema.field(0).data_type() else { + panic!("expected list field"); + }; + let DataType::Struct(fields) = item.data_type() else { + panic!("expected struct item"); + }; + assert!(is_json_field(&fields[1])); + + let list_array: &ListArray = converted.column(0).as_list(); + let values = list_array.values().as_struct(); + let extra = values + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(decode_json(extra.value(0)).contains("h264")); + assert!(extra.is_null(1)); + + let logical = convert_lance_json_to_arrow(&converted).unwrap(); + let logical_schema = logical.schema(); + let DataType::List(item) = logical_schema.field(0).data_type() else { + panic!("expected list field"); + }; + let DataType::Struct(fields) = item.data_type() else { + panic!("expected struct item"); + }; + assert!(is_arrow_json_field(&fields[1])); + + let list_array: &ListArray = logical.column(0).as_list(); + let values = list_array.values().as_struct(); + let extra = values + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(extra.value(0).contains("h264")); + assert!(extra.is_null(1)); + } + + #[test] + fn test_convert_fixed_size_list_zero_json_preserves_length() { + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + let item = Arc::new(ArrowField::new("item", DataType::Utf8, true).with_metadata(metadata)); + let values = Arc::new(StringArray::from(Vec::>::new())) as ArrayRef; + let lists = FixedSizeListArray::try_new_with_length(item, 0, values, None, 3).unwrap(); + let schema = Arc::new(Schema::new(vec![ArrowField::new( + "lists", + lists.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(lists) as ArrayRef]).unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_rows(), 3); + assert_eq!(converted.column(0).len(), 3); + + let converted_schema = converted.schema(); + let DataType::FixedSizeList(item, size) = converted_schema.field(0).data_type() else { + panic!("expected fixed size list field"); + }; + assert_eq!(*size, 0); + assert!(is_json_field(item)); + + let logical = convert_lance_json_to_arrow(&converted).unwrap(); + assert_eq!(logical.num_rows(), 3); + assert_eq!(logical.column(0).len(), 3); + + let logical_schema = logical.schema(); + let DataType::FixedSizeList(item, size) = logical_schema.field(0).data_type() else { + panic!("expected fixed size list field"); + }; + assert_eq!(*size, 0); + assert!(is_arrow_json_field(item)); + } + + #[test] + fn test_has_json_fields() { + // Test direct JSON field + let json_f = json_field("data", true); + assert!(has_json_fields(&json_f)); + + // Test non-JSON field + let non_json = ArrowField::new("data", DataType::Utf8, true); + assert!(!has_json_fields(&non_json)); + + // Test struct containing JSON field + let struct_field = ArrowField::new( + "struct", + DataType::Struct(vec![json_field("nested_json", true)].into()), + true, + ); + assert!(has_json_fields(&struct_field)); + + // Test struct without JSON field + let struct_no_json = ArrowField::new( + "struct", + DataType::Struct(vec![ArrowField::new("text", DataType::Utf8, true)].into()), + true, + ); + assert!(!has_json_fields(&struct_no_json)); + + // Test List containing JSON field + let list_field = ArrowField::new( + "list", + DataType::List(Arc::new(json_field("item", true))), + true, + ); + assert!(has_json_fields(&list_field)); + + // Test LargeList containing JSON field + let large_list_field = ArrowField::new( + "large_list", + DataType::LargeList(Arc::new(json_field("item", true))), + true, + ); + assert!(has_json_fields(&large_list_field)); + + // Test FixedSizeList containing JSON field + let fixed_list_field = ArrowField::new( + "fixed_list", + DataType::FixedSizeList(Arc::new(json_field("item", true)), 3), + true, + ); + assert!(has_json_fields(&fixed_list_field)); + + // Test Map containing JSON field + let map_field = ArrowField::new( + "map", + DataType::Map( + Arc::new(ArrowField::new( + "entries", + DataType::Struct( + vec![ + ArrowField::new("key", DataType::Utf8, false), + json_field("value", true), + ] + .into(), + ), + false, + )), + false, + ), + true, + ); + assert!(has_json_fields(&map_field)); + } + + #[test] + fn test_json_array_inner() { + let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap(); + let inner = json_array.inner(); + assert_eq!(inner.len(), 1); + } + + #[test] + fn test_json_array_value_null_error() { + let json_array = JsonArray::try_from_iter(vec![None::<&str>]).unwrap(); + let result = json_array.value(0); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("null")); + } + + #[test] + fn test_json_array_value_bytes() { + let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap(); + let bytes = json_array.value_bytes(0); + assert!(!bytes.is_empty()); + } + + #[test] + fn test_json_path_with_null() { + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"user": {"name": "Alice"}}"#), None::<&str>]) + .unwrap(); + + let result = json_array.json_path(1, "$.user.name").unwrap(); + assert_eq!(result, None); + } + + #[test] + fn test_to_arrow_json() { + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"name": "Alice"}"#), + None::<&str>, + Some(r#"{"name": "Bob"}"#), + ]) + .unwrap(); + + let arrow_json = json_array.to_arrow_json(); + assert_eq!(arrow_json.len(), 3); + assert!(!arrow_json.is_null(0)); + assert!(arrow_json.is_null(1)); + assert!(!arrow_json.is_null(2)); + + let string_array = arrow_json.as_any().downcast_ref::().unwrap(); + assert!(string_array.value(0).contains("Alice")); + assert!(string_array.value(2).contains("Bob")); + } + + #[test] + fn test_json_array_trait_methods() { + let json_array = + JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#), Some(r#"{"b": 2}"#)]).unwrap(); + + // Wrapper methods + assert_eq!(json_array.len(), 2); + assert!(!json_array.is_empty()); + assert!(!json_array.is_null(0)); + + // Underlying Arrow array + assert_eq!(json_array.inner().data_type(), &DataType::LargeBinary); + assert_eq!(json_array.inner().len(), 2); + } + + #[test] + fn test_json_array_empty() { + let json_array = JsonArray::try_from_iter(Vec::>::new()).unwrap(); + assert!(json_array.is_empty()); + assert_eq!(json_array.len(), 0); + } + + #[test] + fn test_try_from_large_string_array() { + let large_string_array = LargeStringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + None, + ]); + + // Test TryFrom<&LargeStringArray> + let json_array = JsonArray::try_from(&large_string_array).unwrap(); + assert_eq!(json_array.len(), 3); + assert!(!json_array.is_null(0)); + assert!(!json_array.is_null(1)); + assert!(json_array.is_null(2)); + + // Test TryFrom (owned) + let large_string_array2 = LargeStringArray::from(vec![Some(r#"{"x": 1}"#)]); + let json_array2 = JsonArray::try_from(large_string_array2).unwrap(); + assert_eq!(json_array2.len(), 1); + } + + #[test] + fn test_try_from_array_ref() { + // Test with Utf8 + let string_array: ArrayRef = Arc::new(StringArray::from(vec![ + Some(r#"{"a": 1}"#), + Some(r#"{"b": 2}"#), + ])); + let json_array = JsonArray::try_from(string_array).unwrap(); + assert_eq!(json_array.len(), 2); + + // Test with LargeUtf8 + let large_string_array: ArrayRef = Arc::new(LargeStringArray::from(vec![ + Some(r#"{"c": 3}"#), + Some(r#"{"d": 4}"#), + ])); + let json_array2 = JsonArray::try_from(large_string_array).unwrap(); + assert_eq!(json_array2.len(), 2); + + // Test with unsupported type + let int_array: ArrayRef = Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3])); + let result = JsonArray::try_from(int_array); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Unsupported")); + } + + #[test] + fn test_arrow_json_to_lance_json_non_json_field() { + // Test that non-JSON fields are returned unchanged + let field = ArrowField::new("text", DataType::Utf8, true); + let converted = arrow_json_to_lance_json(&field); + assert_eq!(converted.data_type(), &DataType::Utf8); + assert_eq!(converted.name(), "text"); + } + + #[test] + fn test_convert_lance_json_to_arrow() { + // Create a batch with Lance JSON column (JSONB) + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"name": "Alice"}"#), + None::<&str>, + Some(r#"{"name": "Bob"}"#), + ]) + .unwrap(); + + let lance_json_field = json_field("data", true); + let schema = Arc::new(Schema::new(vec![lance_json_field])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(json_array.into_inner()) as ArrayRef]) + .unwrap(); + + // Convert back to Arrow JSON + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + + // Check schema + let converted_schema = converted.schema(); + let converted_field = converted_schema.field(0); + assert_eq!(converted_field.data_type(), &DataType::Utf8); + assert_eq!( + converted_field.metadata().get(ARROW_EXT_NAME_KEY), + Some(&ARROW_JSON_EXT_NAME.to_string()) + ); + + // Check data + let string_array = converted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(!string_array.is_null(0)); + assert!(string_array.is_null(1)); + assert!(!string_array.is_null(2)); + assert!(string_array.value(0).contains("Alice")); + assert!(string_array.value(2).contains("Bob")); + } + + #[test] + fn test_convert_lance_json_to_arrow_empty_batch() { + // Create an empty batch with Lance JSON column + let lance_json_field = json_field("data", true); + let schema = Arc::new(Schema::new(vec![lance_json_field])); + let empty_binary = LargeBinaryBuilder::new().finish(); + let batch = RecordBatch::try_new(schema, vec![Arc::new(empty_binary) as ArrayRef]).unwrap(); + + // Convert back to Arrow JSON + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + assert_eq!(converted.num_rows(), 0); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_convert_lance_json_to_arrow_no_json_columns() { + // Create a batch without JSON columns + let field = ArrowField::new("text", DataType::Utf8, true); + let schema = Arc::new(Schema::new(vec![field])); + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef]).unwrap(); + + // Convert - should return the same batch + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + assert_eq!(converted.num_columns(), 1); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_convert_json_columns_empty_batch() { + // Create an empty batch with Arrow JSON column + let mut field = ArrowField::new("data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let empty_strings = arrow_array::builder::StringBuilder::new().finish(); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(empty_strings) as ArrayRef]).unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_rows(), 0); + assert_eq!( + converted.schema().field(0).data_type(), + &DataType::LargeBinary + ); + } + + #[test] + fn test_convert_json_columns_large_string() { + // Create a batch with Arrow JSON column using LargeUtf8 + let json_strings = LargeStringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + ]); + + let mut field = ArrowField::new("data", DataType::LargeUtf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(json_strings) as ArrayRef]).unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_columns(), 1); + assert_eq!( + converted.schema().field(0).data_type(), + &DataType::LargeBinary + ); + assert_eq!(converted.num_rows(), 2); + } + + #[test] + fn test_convert_json_columns_no_json_columns() { + // Create a batch without JSON columns + let field = ArrowField::new("text", DataType::Utf8, true); + let schema = Arc::new(Schema::new(vec![field])); + let string_array = StringArray::from(vec![Some("hello"), Some("world")]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef]).unwrap(); + + // Convert - should return the same batch + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_columns(), 1); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + } + + #[test] + fn test_convert_json_columns_mixed_columns() { + // Create a batch with both JSON and non-JSON columns + let json_strings = StringArray::from(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + ]); + let text_strings = StringArray::from(vec![Some("hello"), Some("world")]); + + let mut json_field = ArrowField::new("json_data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + json_field.set_metadata(metadata); + + let text_field = ArrowField::new("text_data", DataType::Utf8, true); + + let schema = Arc::new(Schema::new(vec![json_field, text_field])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(json_strings) as ArrayRef, + Arc::new(text_strings) as ArrayRef, + ], + ) + .unwrap(); + + let converted = convert_json_columns(&batch).unwrap(); + assert_eq!(converted.num_columns(), 2); + assert_eq!( + converted.schema().field(0).data_type(), + &DataType::LargeBinary + ); + assert_eq!(converted.schema().field(1).data_type(), &DataType::Utf8); + } + + #[test] + fn test_is_arrow_json_field_large_utf8() { + // Test with LargeUtf8 storage type + let mut field = ArrowField::new("data", DataType::LargeUtf8, true); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + assert!(is_arrow_json_field(&field)); + } + + #[test] + fn test_encode_json_invalid() { + // Test encoding invalid JSON + let result = encode_json("not valid json {"); + assert!(result.is_err()); + } + + #[test] + fn test_json_array_from_invalid_json() { + // Test creating JsonArray from invalid JSON strings + let result = JsonArray::try_from_iter(vec![Some("invalid json {")]); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Failed to encode")); + } + + #[test] + fn test_try_from_string_array_invalid_json() { + let string_array = StringArray::from(vec![Some("invalid json {")]); + let result = JsonArray::try_from(string_array); + assert!(result.is_err()); + } + + #[test] + fn test_try_from_large_string_array_invalid_json() { + let large_string_array = LargeStringArray::from(vec![Some("invalid json {")]); + let result = JsonArray::try_from(large_string_array); + assert!(result.is_err()); + } + + #[test] + fn test_convert_lance_json_to_arrow_mixed_columns() { + // Create a batch with both JSON and non-JSON columns + let json_array = JsonArray::try_from_iter(vec![ + Some(r#"{"name": "Alice"}"#), + Some(r#"{"name": "Bob"}"#), + ]) + .unwrap(); + let text_strings = StringArray::from(vec![Some("hello"), Some("world")]); + + let json_f = json_field("json_data", true); + let text_field = ArrowField::new("text_data", DataType::Utf8, true); + + let schema = Arc::new(Schema::new(vec![json_f, text_field])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(json_array.into_inner()) as ArrayRef, + Arc::new(text_strings) as ArrayRef, + ], + ) + .unwrap(); + + let converted = convert_lance_json_to_arrow(&batch).unwrap(); + assert_eq!(converted.num_columns(), 2); + assert_eq!(converted.schema().field(0).data_type(), &DataType::Utf8); + assert_eq!(converted.schema().field(1).data_type(), &DataType::Utf8); + } + + #[test] + fn test_json_path_invalid_path() { + let json_array = JsonArray::try_from_iter(vec![Some(r#"{"a": 1}"#)]).unwrap(); + // Invalid JSONPath syntax should return error + let result = json_array.json_path(0, "invalid path without $"); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Invalid JSONPath")); + } + + #[test] + fn test_convert_json_columns_invalid_storage_type() { + // Create a batch with Arrow JSON field but wrong storage type (Int32 instead of Utf8) + let int_array = arrow_array::Int32Array::from(vec![1, 2, 3]); + + let mut field = ArrowField::new("data", DataType::Int32, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(int_array) as ArrayRef]).unwrap(); + + // This should succeed since Int32 doesn't match is_arrow_json_field check + // (is_arrow_json_field requires Utf8 or LargeUtf8) + let result = convert_json_columns(&batch); + assert!(result.is_ok()); + } + + #[test] + fn test_is_json_field_wrong_extension() { + // LargeBinary field without the correct extension metadata + let field = ArrowField::new("data", DataType::LargeBinary, true); + assert!(!is_json_field(&field)); + + // LargeBinary field with wrong extension name + let mut field2 = ArrowField::new("data", DataType::LargeBinary, true); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + "other.extension".to_string(), + ); + field2.set_metadata(metadata); + assert!(!is_json_field(&field2)); + } + + #[test] + fn test_is_arrow_json_field_wrong_extension() { + // Utf8 field without extension metadata + let field = ArrowField::new("data", DataType::Utf8, true); + assert!(!is_arrow_json_field(&field)); + + // Utf8 field with wrong extension name + let mut field2 = ArrowField::new("data", DataType::Utf8, true); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + "other.extension".to_string(), + ); + field2.set_metadata(metadata); + assert!(!is_arrow_json_field(&field2)); + + // Wrong type entirely + let field3 = ArrowField::new("data", DataType::Int32, true); + assert!(!is_arrow_json_field(&field3)); + } + + #[test] + fn test_convert_json_columns_invalid_json_utf8() { + // Test error propagation when converting invalid JSON (Utf8) + let invalid_json = StringArray::from(vec![Some("invalid json {")]); + + let mut field = ArrowField::new("data", DataType::Utf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(invalid_json) as ArrayRef]).unwrap(); + + let result = convert_json_columns(&batch); + assert!(result.is_err()); + } + + #[test] + fn test_convert_json_columns_invalid_json_large_utf8() { + // Test error propagation when converting invalid JSON (LargeUtf8) + let invalid_json = LargeStringArray::from(vec![Some("invalid json {")]); + + let mut field = ArrowField::new("data", DataType::LargeUtf8, false); + let mut metadata = std::collections::HashMap::new(); + metadata.insert( + ARROW_EXT_NAME_KEY.to_string(), + ARROW_JSON_EXT_NAME.to_string(), + ); + field.set_metadata(metadata); + + let schema = Arc::new(Schema::new(vec![field])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(invalid_json) as ArrayRef]).unwrap(); + + let result = convert_json_columns(&batch); + assert!(result.is_err()); + } + + #[test] + fn test_json_path_on_corrupted_jsonb() { + // Create corrupted JSONB bytes directly + let corrupted_bytes: &[u8] = &[0xFF, 0xFE, 0x00, 0x01, 0x02]; + let corrupted_binary = LargeBinaryArray::from(vec![Some(corrupted_bytes)]); + + // Wrap in JsonArray + let corrupted_json = JsonArray { + inner: corrupted_binary, + }; + + // Try to use json_path on corrupted data - the selector might fail or return unexpected results + // This exercises the code path but may not produce an error depending on jsonb library behavior + let _result = corrupted_json.json_path(0, "$.a"); + // We don't assert on the result as the behavior depends on the jsonb library + } + + #[test] + fn test_decode_json_on_various_inputs() { + // Test decode_json with various inputs + let valid_jsonb = encode_json(r#"{"key": "value"}"#).unwrap(); + let decoded = decode_json(&valid_jsonb); + assert!(decoded.contains("key")); + + // Empty bytes - jsonb library handles this gracefully + let decoded_empty = decode_json(&[]); + // Just verify it doesn't panic + let _ = decoded_empty; + + // Random bytes - jsonb library handles this gracefully + let decoded_random = decode_json(&[0xFF, 0xFE, 0x00]); + // Just verify it doesn't panic + let _ = decoded_random; + } +} diff --git a/python/vendor/lance-arrow/src/lib.rs b/python/vendor/lance-arrow/src/lib.rs new file mode 100644 index 000000000..f26e1f3e3 --- /dev/null +++ b/python/vendor/lance-arrow/src/lib.rs @@ -0,0 +1,3102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extend Arrow Functionality +//! +//! To improve Arrow-RS ergonomic + +#![warn(clippy::undocumented_unsafe_blocks)] + +// lance-arrow reinterprets value bytes as native numeric types in +// `FloatArray::as_slice` for `bf16` (rust/lance-arrow/src/bfloat16.rs), which +// requires the host byte order to match the on-disk byte order Lance writes. +// Lance writes little-endian; building on a big-endian target would silently +// produce wrong numeric values. +#[cfg(not(target_endian = "little"))] +compile_error!("lance-arrow only supports little-endian targets"); + +use std::sync::Arc; +use std::{collections::HashMap, ptr::NonNull}; + +use arrow_array::{ + Array, ArrayRef, ArrowNumericType, FixedSizeBinaryArray, FixedSizeListArray, GenericListArray, + LargeListArray, ListArray, MapArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, StructArray, + UInt8Array, UInt32Array, cast::AsArray, +}; +use arrow_array::{ + Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, new_null_array, +}; +use arrow_buffer::MutableBuffer; +use arrow_data::ArrayDataBuilder; +use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema, SortOptions}; +use arrow_select::{interleave::interleave, take::take}; +use rand::prelude::*; + +pub mod deepcopy; +pub mod schema; +pub use schema::*; +pub mod bfloat16; +pub mod floats; +use crate::list::ListArrayExt; +pub use floats::*; + +pub mod ipc; +pub mod json; +pub mod list; +pub mod memory; +pub mod scalar; +pub mod stream; +pub mod r#struct; + +/// Arrow extension metadata key for extension name +pub const ARROW_EXT_NAME_KEY: &str = "ARROW:extension:name"; + +/// Arrow extension metadata key for extension metadata +pub const ARROW_EXT_META_KEY: &str = "ARROW:extension:metadata"; + +/// Key used by lance to mark a field as a blob +/// TODO: Use Arrow extension mechanism instead? +pub const BLOB_META_KEY: &str = "lance-encoding:blob"; +/// Arrow extension type name for Lance blob v2 columns +pub const BLOB_V2_EXT_NAME: &str = "lance.blob.v2"; +/// Metadata key for overriding the dedicated blob size threshold (in bytes) +pub const BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY: &str = + "lance-encoding:blob-dedicated-size-threshold"; +/// Metadata key for overriding the inline blob size threshold (in bytes) +pub const BLOB_INLINE_SIZE_THRESHOLD_META_KEY: &str = "lance-encoding:blob-inline-size-threshold"; +/// Metadata key for overriding the maximum size (in bytes) of a packed blob sidecar file +pub const BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY: &str = + "lance-encoding:blob-pack-file-size-threshold"; + +type Result = std::result::Result; + +pub trait DataTypeExt { + /// Returns true if the data type is binary-like, such as Utf8, Binary, or the large and/or view variants. + fn is_binary_like(&self) -> bool; + + /// Returns true if the data type is a struct. + fn is_struct(&self) -> bool; + + /// Check whether the given Arrow DataType is fixed stride. + /// + /// A fixed stride type has the same byte width for all array elements + /// This includes all PrimitiveType's Boolean, FixedSizeList, FixedSizeBinary, and Decimals + fn is_fixed_stride(&self) -> bool; + + /// Returns true if the [DataType] is a dictionary type. + fn is_dictionary(&self) -> bool; + + /// Returns the byte width of the data type + /// Panics if the data type is not fixed stride. + fn byte_width(&self) -> usize; + + /// Returns the byte width of the data type, if it is fixed stride. + /// Returns None if the data type is not fixed stride. + fn byte_width_opt(&self) -> Option; +} + +impl DataTypeExt for DataType { + fn is_binary_like(&self) -> bool { + use DataType::*; + matches!( + self, + Utf8 | Binary | LargeUtf8 | LargeBinary | Utf8View | BinaryView + ) + } + + fn is_struct(&self) -> bool { + matches!(self, Self::Struct(_)) + } + + fn is_fixed_stride(&self) -> bool { + use DataType::*; + matches!( + self, + Boolean + | UInt8 + | UInt16 + | UInt32 + | UInt64 + | Int8 + | Int16 + | Int32 + | Int64 + | Float16 + | Float32 + | Float64 + | Decimal128(_, _) + | Decimal256(_, _) + | FixedSizeList(_, _) + | FixedSizeBinary(_) + | Duration(_) + | Timestamp(_, _) + | Date32 + | Date64 + | Time32(_) + | Time64(_) + ) + } + + fn is_dictionary(&self) -> bool { + matches!(self, Self::Dictionary(_, _)) + } + + fn byte_width_opt(&self) -> Option { + match self { + Self::Int8 => Some(1), + Self::Int16 => Some(2), + Self::Int32 => Some(4), + Self::Int64 => Some(8), + Self::UInt8 => Some(1), + Self::UInt16 => Some(2), + Self::UInt32 => Some(4), + Self::UInt64 => Some(8), + Self::Float16 => Some(2), + Self::Float32 => Some(4), + Self::Float64 => Some(8), + Self::Date32 => Some(4), + Self::Date64 => Some(8), + Self::Time32(_) => Some(4), + Self::Time64(_) => Some(8), + Self::Timestamp(_, _) => Some(8), + Self::Duration(_) => Some(8), + Self::Decimal128(_, _) => Some(16), + Self::Decimal256(_, _) => Some(32), + Self::Interval(unit) => match unit { + IntervalUnit::YearMonth => Some(4), + IntervalUnit::DayTime => Some(8), + IntervalUnit::MonthDayNano => Some(16), + }, + Self::FixedSizeBinary(s) => Some(*s as usize), + Self::FixedSizeList(dt, s) => dt + .data_type() + .byte_width_opt() + .map(|width| width * *s as usize), + _ => None, + } + } + + fn byte_width(&self) -> usize { + self.byte_width_opt() + .unwrap_or_else(|| panic!("Expecting fixed stride data type, found {:?}", self)) + } +} + +/// Create an [`GenericListArray`] from values and offsets. +/// +/// ``` +/// use arrow_array::{Int32Array, Int64Array, ListArray}; +/// use arrow_array::types::Int64Type; +/// use lance_arrow::try_new_generic_list_array; +/// +/// let offsets = Int32Array::from_iter([0, 2, 7, 10]); +/// let int_values = Int64Array::from_iter(0..10); +/// let list_arr = try_new_generic_list_array(int_values, &offsets).unwrap(); +/// assert_eq!(list_arr, +/// ListArray::from_iter_primitive::(vec![ +/// Some(vec![Some(0), Some(1)]), +/// Some(vec![Some(2), Some(3), Some(4), Some(5), Some(6)]), +/// Some(vec![Some(7), Some(8), Some(9)]), +/// ])) +/// ``` +pub fn try_new_generic_list_array( + values: T, + offsets: &PrimitiveArray, +) -> Result> +where + Offset::Native: OffsetSizeTrait, +{ + let data_type = if Offset::Native::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + values.data_type().clone(), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + values.data_type().clone(), + true, + ))) + }; + let data = ArrayDataBuilder::new(data_type) + .len(offsets.len() - 1) + .add_buffer(offsets.into_data().buffers()[0].clone()) + .add_child_data(values.into_data()) + .build()?; + + Ok(GenericListArray::from(data)) +} + +pub fn fixed_size_list_type(list_width: i32, inner_type: DataType) -> DataType { + DataType::FixedSizeList(Arc::new(Field::new("item", inner_type, true)), list_width) +} + +pub trait FixedSizeListArrayExt { + /// Create an [`FixedSizeListArray`] from values and list size. + /// + /// ``` + /// use arrow_array::{Int64Array, FixedSizeListArray}; + /// use arrow_array::types::Int64Type; + /// use lance_arrow::FixedSizeListArrayExt; + /// + /// let int_values = Int64Array::from_iter(0..10); + /// let fixed_size_list_arr = FixedSizeListArray::try_new_from_values(int_values, 2).unwrap(); + /// assert_eq!(fixed_size_list_arr, + /// FixedSizeListArray::from_iter_primitive::(vec![ + /// Some(vec![Some(0), Some(1)]), + /// Some(vec![Some(2), Some(3)]), + /// Some(vec![Some(4), Some(5)]), + /// Some(vec![Some(6), Some(7)]), + /// Some(vec![Some(8), Some(9)]) + /// ], 2)) + /// ``` + fn try_new_from_values( + values: T, + list_size: i32, + ) -> Result; + + /// Sample `n` rows from the [FixedSizeListArray] + /// + /// ``` + /// use arrow_array::{Int64Array, FixedSizeListArray, Array}; + /// use lance_arrow::FixedSizeListArrayExt; + /// + /// let int_values = Int64Array::from_iter(0..256); + /// let fixed_size_list_arr = FixedSizeListArray::try_new_from_values(int_values, 16).unwrap(); + /// let sampled = fixed_size_list_arr.sample(10).unwrap(); + /// assert_eq!(sampled.len(), 10); + /// assert_eq!(sampled.value_length(), 16); + /// assert_eq!(sampled.values().len(), 160); + /// ``` + fn sample(&self, n: usize) -> Result; + + /// Ensure the [FixedSizeListArray] of Float16, Float32, Float64, + /// Int8, Int16, Int32, Int64, UInt8, UInt32 type to its closest floating point type. + fn convert_to_floating_point(&self) -> Result; +} + +impl FixedSizeListArrayExt for FixedSizeListArray { + fn try_new_from_values(values: T, list_size: i32) -> Result { + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let values = Arc::new(values); + + Self::try_new(field, list_size, values, None) + } + + fn sample(&self, n: usize) -> Result { + if n >= self.len() { + return Ok(self.clone()); + } + let mut rng = SmallRng::from_os_rng(); + let chosen = (0..self.len() as u32).choose_multiple(&mut rng, n); + take(self, &UInt32Array::from(chosen), None).map(|arr| arr.as_fixed_size_list().clone()) + } + + fn convert_to_floating_point(&self) -> Result { + match self.data_type() { + DataType::FixedSizeList(field, size) => match field.data_type() { + DataType::Float16 | DataType::Float32 | DataType::Float64 => Ok(self.clone()), + DataType::Int8 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float32, + field.is_nullable(), + )), + *size, + Arc::new(Float32Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int8Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f32)), + )), + self.nulls().cloned(), + )), + DataType::Int16 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float32, + field.is_nullable(), + )), + *size, + Arc::new(Float32Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int16Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f32)), + )), + self.nulls().cloned(), + )), + DataType::Int32 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float32, + field.is_nullable(), + )), + *size, + Arc::new(Float32Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int32Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f32)), + )), + self.nulls().cloned(), + )), + DataType::Int64 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float64, + field.is_nullable(), + )), + *size, + Arc::new(Float64Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to Int64Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f64)), + )), + self.nulls().cloned(), + )), + DataType::UInt8 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float64, + field.is_nullable(), + )), + *size, + Arc::new(Float64Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to UInt8Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f64)), + )), + self.nulls().cloned(), + )), + DataType::UInt32 => Ok(Self::new( + Arc::new(arrow_schema::Field::new( + field.name(), + DataType::Float64, + field.is_nullable(), + )), + *size, + Arc::new(Float64Array::from_iter( + self.values() + .as_any() + .downcast_ref::() + .ok_or(ArrowError::ParseError( + "Fail to cast primitive array to UInt32Type".to_string(), + ))? + .into_iter() + .map(|x| x.map(|y| y as f64)), + )), + self.nulls().cloned(), + )), + data_type => Err(ArrowError::ParseError(format!( + "Expect either floating type or integer got {:?}", + data_type + ))), + }, + data_type => Err(ArrowError::ParseError(format!( + "Expect either FixedSizeList got {:?}", + data_type + ))), + } + } +} + +/// Force downcast of an [`Array`], such as an [`ArrayRef`], to +/// [`FixedSizeListArray`], panic'ing on failure. +pub fn as_fixed_size_list_array(arr: &dyn Array) -> &FixedSizeListArray { + arr.as_any().downcast_ref::().unwrap() +} + +pub trait FixedSizeBinaryArrayExt { + /// Create an [`FixedSizeBinaryArray`] from values and stride. + /// + /// ``` + /// use arrow_array::{UInt8Array, FixedSizeBinaryArray}; + /// use arrow_array::types::UInt8Type; + /// use lance_arrow::FixedSizeBinaryArrayExt; + /// + /// let int_values = UInt8Array::from_iter(0..10); + /// let fixed_size_list_arr = FixedSizeBinaryArray::try_new_from_values(&int_values, 2).unwrap(); + /// assert_eq!(fixed_size_list_arr, + /// FixedSizeBinaryArray::try_from_iter( + /// [[0u8, 1], [2, 3], [4, 5], [6, 7], [8, 9]].into_iter() + /// ).unwrap()) + /// ``` + fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result; +} + +impl FixedSizeBinaryArrayExt for FixedSizeBinaryArray { + fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result { + let data_type = DataType::FixedSizeBinary(stride); + let data = ArrayDataBuilder::new(data_type) + .len(values.len() / stride as usize) + .add_buffer(values.into_data().buffers()[0].clone()) + .build()?; + Ok(Self::from(data)) + } +} + +pub fn as_fixed_size_binary_array(arr: &dyn Array) -> &FixedSizeBinaryArray { + arr.as_any().downcast_ref::().unwrap() +} + +pub fn iter_str_array(arr: &dyn Array) -> Box> + Send + '_> { + match arr.data_type() { + DataType::Utf8 => Box::new(arr.as_string::().iter()), + DataType::LargeUtf8 => Box::new(arr.as_string::().iter()), + DataType::Utf8View => Box::new(arr.as_string_view().iter()), + _ => panic!( + "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}", + arr.data_type() + ), + } +} + +pub fn iter_binary_array( + arr: &dyn Array, +) -> Result> + Send + '_>> { + match arr.data_type() { + DataType::Binary => Ok(Box::new(arr.as_binary::().iter())), + DataType::LargeBinary => Ok(Box::new(arr.as_binary::().iter())), + DataType::BinaryView => Ok(Box::new(arr.as_binary_view().iter())), + DataType::FixedSizeBinary(_) => Ok(Box::new(arr.as_fixed_size_binary().iter())), + data_type => Err(ArrowError::InvalidArgumentError(format!( + "Expecting a binary type, found {data_type}" + ))), + } +} + +/// Extends Arrow's [RecordBatch]. +pub trait RecordBatchExt { + /// Append a new column to this [`RecordBatch`] and returns a new RecordBatch. + /// + /// ``` + /// use std::sync::Arc; + /// use arrow_array::{RecordBatch, Int32Array, StringArray}; + /// use arrow_schema::{Schema, Field, DataType}; + /// use lance_arrow::*; + /// + /// let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + /// let int_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + /// let record_batch = RecordBatch::try_new(schema, vec![int_arr.clone()]).unwrap(); + /// + /// let new_field = Field::new("s", DataType::Utf8, true); + /// let str_arr = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + /// let new_record_batch = record_batch.try_with_column(new_field, str_arr.clone()).unwrap(); + /// + /// assert_eq!( + /// new_record_batch, + /// RecordBatch::try_new( + /// Arc::new(Schema::new( + /// vec![ + /// Field::new("a", DataType::Int32, true), + /// Field::new("s", DataType::Utf8, true) + /// ]) + /// ), + /// vec![int_arr, str_arr], + /// ).unwrap() + /// ) + /// ``` + fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result; + + /// Created a new RecordBatch with column at index. + fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result; + + /// Creates a new [`RecordBatch`] from the provided [`StructArray`]. + /// + /// The fields on the [`StructArray`] need to match this [`RecordBatch`] schema + fn try_new_from_struct_array(&self, arr: StructArray) -> Result; + + /// Merge with another [`RecordBatch`] and returns a new one. + /// + /// Fields are merged based on name. First we iterate the left columns. If a matching + /// name is found in the right then we merge the two columns. If there is no match then + /// we add the left column to the output. + /// + /// To merge two columns we consider the type. If both arrays are struct arrays we recurse. + /// Otherwise we use the left array. + /// + /// Afterwards we add all non-matching right columns to the output. + /// + /// Note: This method likely does not handle nested fields correctly and you may want to consider + /// using [`Self::merge_with_schema`] instead. + /// ``` + /// use std::sync::Arc; + /// use arrow_array::*; + /// use arrow_schema::{Schema, Field, DataType}; + /// use lance_arrow::*; + /// + /// let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + /// let int_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + /// let left = RecordBatch::try_new(left_schema, vec![int_arr.clone()]).unwrap(); + /// + /// let right_schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + /// let str_arr = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + /// let right = RecordBatch::try_new(right_schema, vec![str_arr.clone()]).unwrap(); + /// + /// let new_record_batch = left.merge(&right).unwrap(); + /// + /// assert_eq!( + /// new_record_batch, + /// RecordBatch::try_new( + /// Arc::new(Schema::new( + /// vec![ + /// Field::new("a", DataType::Int32, true), + /// Field::new("s", DataType::Utf8, true) + /// ]) + /// ), + /// vec![int_arr, str_arr], + /// ).unwrap() + /// ) + /// ``` + /// + /// TODO: add merge nested fields support. + fn merge(&self, other: &RecordBatch) -> Result; + + /// Create a batch by merging columns between two batches with a given schema. + /// + /// A reference schema is used to determine the proper ordering of nested fields. + /// + /// For each field in the reference schema we look for corresponding fields in + /// the left and right batches. If a field is found in both batches we recursively merge + /// it. + /// + /// If a field is only in the left or right batch we take it as it is. + fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result; + + /// Drop one column specified with the name and return the new [`RecordBatch`]. + /// + /// If the named column does not exist, it returns a copy of this [`RecordBatch`]. + fn drop_column(&self, name: &str) -> Result; + + /// Replace a column (specified by name) and return the new [`RecordBatch`]. + fn replace_column_by_name(&self, name: &str, column: Arc) -> Result; + + /// Replace a column schema (specified by name) and return the new [`RecordBatch`]. + fn replace_column_schema_by_name( + &self, + name: &str, + new_data_type: DataType, + column: Arc, + ) -> Result; + + /// Rename a column at a given index. + fn rename_column(&self, index: usize, new_name: &str) -> Result; + + /// Get (potentially nested) column by qualified name. + fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef>; + + /// Project the schema over the [RecordBatch]. + fn project_by_schema(&self, schema: &Schema) -> Result; + + /// metadata of the schema. + fn metadata(&self) -> &HashMap; + + /// Add metadata to the schema. + fn add_metadata(&self, key: String, value: String) -> Result { + let mut metadata = self.metadata().clone(); + metadata.insert(key, value); + self.with_metadata(metadata) + } + + /// Replace the schema metadata with the provided one. + fn with_metadata(&self, metadata: HashMap) -> Result; + + /// Take selected rows from the [RecordBatch]. + fn take(&self, indices: &UInt32Array) -> Result; + + /// Create a new RecordBatch with compacted memory after slicing. + fn shrink_to_fit(&self) -> Result; + + /// Helper method to sort the RecordBatch by a column + fn sort_by_column(&self, column: usize, options: Option) -> Result; +} + +impl RecordBatchExt for RecordBatch { + fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result { + let new_schema = Arc::new(self.schema().as_ref().try_with_column(field)?); + let mut new_columns = self.columns().to_vec(); + new_columns.push(arr); + Self::try_new(new_schema, new_columns) + } + + fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result { + let new_schema = Arc::new(self.schema().as_ref().try_with_column_at(index, field)?); + let mut new_columns = self.columns().to_vec(); + new_columns.insert(index, arr); + Self::try_new(new_schema, new_columns) + } + + fn try_new_from_struct_array(&self, arr: StructArray) -> Result { + let schema = Arc::new(Schema::new_with_metadata( + arr.fields().to_vec(), + self.schema().metadata.clone(), + )); + let batch = Self::from(arr); + batch.with_schema(schema) + } + + fn merge(&self, other: &Self) -> Result { + if self.num_rows() != other.num_rows() { + return Err(ArrowError::InvalidArgumentError(format!( + "Attempt to merge two RecordBatch with different sizes: {} != {}", + self.num_rows(), + other.num_rows() + ))); + } + let left_struct_array: StructArray = self.clone().into(); + let right_struct_array: StructArray = other.clone().into(); + self.try_new_from_struct_array(merge(&left_struct_array, &right_struct_array)) + } + + fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result { + if self.num_rows() != other.num_rows() { + return Err(ArrowError::InvalidArgumentError(format!( + "Attempt to merge two RecordBatch with different sizes: {} != {}", + self.num_rows(), + other.num_rows() + ))); + } + let left_struct_array: StructArray = self.clone().into(); + let right_struct_array: StructArray = other.clone().into(); + self.try_new_from_struct_array(merge_with_schema( + &left_struct_array, + &right_struct_array, + schema.fields(), + )) + } + + fn drop_column(&self, name: &str) -> Result { + let mut fields = vec![]; + let mut columns = vec![]; + for i in 0..self.schema().fields.len() { + if self.schema().field(i).name() != name { + fields.push(self.schema().field(i).clone()); + columns.push(self.column(i).clone()); + } + } + Self::try_new( + Arc::new(Schema::new_with_metadata( + fields, + self.schema().metadata().clone(), + )), + columns, + ) + } + + fn rename_column(&self, index: usize, new_name: &str) -> Result { + let mut fields = self.schema().fields().to_vec(); + if index >= fields.len() { + return Err(ArrowError::InvalidArgumentError(format!( + "Index out of bounds: {}", + index + ))); + } + fields[index] = Arc::new(Field::new( + new_name, + fields[index].data_type().clone(), + fields[index].is_nullable(), + )); + Self::try_new( + Arc::new(Schema::new_with_metadata( + fields, + self.schema().metadata().clone(), + )), + self.columns().to_vec(), + ) + } + + fn replace_column_by_name(&self, name: &str, column: Arc) -> Result { + let mut columns = self.columns().to_vec(); + let field_i = self + .schema() + .fields() + .iter() + .position(|f| f.name() == name) + .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?; + columns[field_i] = column; + Self::try_new(self.schema(), columns) + } + + fn replace_column_schema_by_name( + &self, + name: &str, + new_data_type: DataType, + column: Arc, + ) -> Result { + let fields = self + .schema() + .fields() + .iter() + .map(|x| { + if x.name() != name { + x.clone() + } else { + let new_field = Field::new(name, new_data_type.clone(), x.is_nullable()); + Arc::new(new_field) + } + }) + .collect::>(); + let schema = Schema::new_with_metadata(fields, self.schema().metadata.clone()); + let mut columns = self.columns().to_vec(); + let field_i = self + .schema() + .fields() + .iter() + .position(|f| f.name() == name) + .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?; + columns[field_i] = column; + Self::try_new(Arc::new(schema), columns) + } + + fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef> { + let split = name.split('.').collect::>(); + if split.is_empty() { + return None; + } + + self.column_by_name(split[0]) + .and_then(|arr| get_sub_array(arr, &split[1..])) + } + + fn project_by_schema(&self, schema: &Schema) -> Result { + let struct_array: StructArray = self.clone().into(); + self.try_new_from_struct_array(project(&struct_array, schema.fields())?) + } + + fn metadata(&self) -> &HashMap { + self.schema_ref().metadata() + } + + fn with_metadata(&self, metadata: HashMap) -> Result { + let mut schema = self.schema_ref().as_ref().clone(); + schema.metadata = metadata; + Self::try_new(schema.into(), self.columns().into()) + } + + fn take(&self, indices: &UInt32Array) -> Result { + let struct_array: StructArray = self.clone().into(); + let taken = take(&struct_array, indices, None)?; + self.try_new_from_struct_array(taken.as_struct().clone()) + } + + fn shrink_to_fit(&self) -> Result { + // Deep copy the sliced record batch, instead of whole batch + crate::deepcopy::deep_copy_batch_sliced(self) + } + + fn sort_by_column(&self, column: usize, options: Option) -> Result { + if column >= self.num_columns() { + return Err(ArrowError::InvalidArgumentError(format!( + "Column index out of bounds: {}", + column + ))); + } + let column = self.column(column); + let sorted = arrow_ord::sort::sort_to_indices(column, options, None)?; + self.take(&sorted) + } +} + +/// Recursively projects an array to match the target field's structure. +/// This handles reordering fields inside nested List types. +fn project_array(array: &ArrayRef, target_field: &Field) -> Result { + match target_field.data_type() { + DataType::Struct(subfields) => { + let struct_arr = array.as_struct(); + let projected = project(struct_arr, subfields)?; + Ok(Arc::new(projected)) + } + DataType::List(inner_field) => { + let list_arr: &ListArray = array.as_list(); + let projected_values = project_array(list_arr.values(), inner_field.as_ref())?; + Ok(Arc::new(ListArray::new( + inner_field.clone(), + list_arr.offsets().clone(), + projected_values, + list_arr.nulls().cloned(), + ))) + } + DataType::LargeList(inner_field) => { + let list_arr: &LargeListArray = array.as_list(); + let projected_values = project_array(list_arr.values(), inner_field.as_ref())?; + Ok(Arc::new(LargeListArray::new( + inner_field.clone(), + list_arr.offsets().clone(), + projected_values, + list_arr.nulls().cloned(), + ))) + } + DataType::FixedSizeList(inner_field, size) => { + let list_arr = array.as_fixed_size_list(); + let projected_values = project_array(list_arr.values(), inner_field.as_ref())?; + Ok(Arc::new(FixedSizeListArray::new( + inner_field.clone(), + *size, + projected_values, + list_arr.nulls().cloned(), + ))) + } + // A nullable entries field fails MapArray::try_new unconditionally, + // so a (schema-invalid) map declared that way keeps the clone + // fallthrough it always had rather than gaining a new error. + DataType::Map(entries_field, sorted) if !entries_field.is_nullable() => { + let map_arr = array.as_map(); + let DataType::Struct(entry_fields) = entries_field.data_type() else { + return Err(ArrowError::SchemaError(format!( + "Map entries field must be a struct, got {}", + entries_field.data_type() + ))); + }; + let projected_entries = project(map_arr.entries(), entry_fields)?; + // try_new re-checks the entries invariants (a non-null entries + // struct, two entry columns, offset bounds); null keys are ruled + // out one level down, by the struct rebuild against the + // non-nullable key field. + Ok(Arc::new(MapArray::try_new( + entries_field.clone(), + map_arr.offsets().clone(), + projected_entries, + map_arr.nulls().cloned(), + *sorted, + )?)) + } + _ => Ok(array.clone()), + } +} + +fn project(struct_array: &StructArray, fields: &Fields) -> Result { + if fields.is_empty() { + return Ok(StructArray::new_empty_fields( + struct_array.len(), + struct_array.nulls().cloned(), + )); + } + let mut columns: Vec = vec![]; + for field in fields.iter() { + if let Some(col) = struct_array.column_by_name(field.name()) { + let projected = project_array(col, field.as_ref())?; + columns.push(projected); + } else { + return Err(ArrowError::SchemaError(format!( + "field {} does not exist in the RecordBatch", + field.name() + ))); + } + } + // Preserve the struct's validity when projecting + StructArray::try_new(fields.clone(), columns, struct_array.nulls().cloned()) +} + +/// Whether `data_type` has a struct somewhere below it, through lists. +fn holds_struct(data_type: &DataType) -> bool { + match data_type { + DataType::Struct(_) => true, + DataType::List(item) | DataType::LargeList(item) => holds_struct(item.data_type()), + _ => false, + } +} + +/// Merge the values of two lists: structs field by field, nested lists +/// recursively, anything else keeps the left side. +fn merge_list_values(left: &ArrayRef, right: &ArrayRef) -> ArrayRef { + match (left.data_type(), right.data_type()) { + (DataType::Struct(_), DataType::Struct(_)) => { + Arc::new(merge(left.as_struct(), right.as_struct())) + } + (DataType::List(_), DataType::List(_)) + | (DataType::LargeList(_), DataType::LargeList(_)) + if left.data_type() != right.data_type() + && holds_struct(left.data_type()) + && holds_struct(right.data_type()) => + { + merge_list_struct(left, right) + } + _ => left.clone(), + } +} + +fn merge_list_struct_helper( + left: &dyn Array, + right: &dyn Array, +) -> Arc { + let left_list: &GenericListArray = left.as_list(); + let right_list: &GenericListArray = right.as_list(); + let item_name = match left.data_type() { + DataType::List(item) | DataType::LargeList(item) => item.name().clone(), + _ => unreachable!(), + }; + // One side all-null (a placeholder for columns this fragment lacks) has + // no values, so stand in nulls shaped like the other side's. + let null_left = left_list.null_count() == left_list.len(); + let null_right = !null_left && right_list.null_count() == right_list.len(); + let (merged_items, list, items_nullable) = if null_left || null_right { + let not_null = if null_left { right_list } else { left_list }; + let values_len = not_null.values().len(); + let values = |list: &GenericListArray, is_null: bool| { + if is_null { + new_null_array(list.values().data_type(), values_len) + } else { + list.values().clone() + } + }; + let merged = merge_list_values( + &values(left_list, null_left), + &values(right_list, null_right), + ); + (merged, not_null, true) + } else { + if left_list.offsets().inner() != right_list.offsets().inner() { + panic!("Attempt to merge list struct arrays which do not have same offsets"); + } + let merged = merge_list_values(left_list.values(), right_list.values()); + let nullable = match left.data_type() { + DataType::List(item) | DataType::LargeList(item) => item.is_nullable(), + _ => unreachable!(), + }; + (merged, left_list, nullable) + }; + let items_field = Arc::new(Field::new( + item_name, + merged_items.data_type().clone(), + items_nullable, + )); + Arc::new(GenericListArray::::new( + items_field, + list.offsets().clone(), + merged_items, + list.nulls().cloned(), + )) +} + +/// Merge two lists whose items hold structs. Only succeeds if both lists have +/// the same offsets or one of them is all-null; panics otherwise. +fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc { + match (left.data_type(), right.data_type()) { + (DataType::List(_), DataType::List(_)) => merge_list_struct_helper::(left, right), + (DataType::LargeList(_), DataType::LargeList(_)) => { + merge_list_struct_helper::(left, right) + } + _ => unreachable!(), + } +} + +/// Helper function to merge validity buffers from two struct arrays. +/// +/// A row is valid if it is valid in either input. +/// An absent validity buffer means all rows are valid, an all-null buffer acts as the identity for this merge. +fn merge_struct_validity( + left_validity: Option<&arrow_buffer::NullBuffer>, + right_validity: Option<&arrow_buffer::NullBuffer>, +) -> Option { + match (left_validity, right_validity) { + // Fast paths: no computation needed + (None, _) | (_, None) => None, + (Some(left), Some(right)) => { + if left.null_count() == 0 || right.null_count() == 0 { + return None; + } + if left.null_count() == left.len() { + return Some(right.clone()); + } + if right.null_count() == right.len() { + return Some(left.clone()); + } + + let left_buffer = left.inner(); + let right_buffer = right.inner(); + + // Perform bitwise OR directly on BooleanBuffers + // This preserves the correct semantics: 1 = valid, 0 = null + let merged_buffer = left_buffer | right_buffer; + + Some(arrow_buffer::NullBuffer::from(merged_buffer)) + } + } +} + +fn merge_list_child_values( + child_field: &Field, + left_values: ArrayRef, + right_values: ArrayRef, +) -> ArrayRef { + match child_field.data_type() { + DataType::Struct(child_fields) => Arc::new(merge_with_schema( + left_values.as_struct(), + right_values.as_struct(), + child_fields, + )) as ArrayRef, + DataType::List(grandchild) => { + let left_list = left_values + .as_any() + .downcast_ref::() + .expect("left list values should be ListArray"); + let right_list = right_values + .as_any() + .downcast_ref::() + .expect("right list values should be ListArray"); + let merged_values = merge_list_child_values( + grandchild.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + Arc::new(ListArray::new( + grandchild.clone(), + left_list.offsets().clone(), + merged_values, + merged_validity, + )) as ArrayRef + } + DataType::LargeList(grandchild) => { + let left_list = left_values + .as_any() + .downcast_ref::() + .expect("left list values should be LargeListArray"); + let right_list = right_values + .as_any() + .downcast_ref::() + .expect("right list values should be LargeListArray"); + let merged_values = merge_list_child_values( + grandchild.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + Arc::new(LargeListArray::new( + grandchild.clone(), + left_list.offsets().clone(), + merged_values, + merged_validity, + )) as ArrayRef + } + DataType::FixedSizeList(grandchild, list_size) => { + let left_list = left_values + .as_any() + .downcast_ref::() + .expect("left list values should be FixedSizeListArray"); + let right_list = right_values + .as_any() + .downcast_ref::() + .expect("right list values should be FixedSizeListArray"); + let merged_values = merge_list_child_values( + grandchild.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls()); + Arc::new(FixedSizeListArray::new( + grandchild.clone(), + *list_size, + merged_values, + merged_validity, + )) as ArrayRef + } + _ => left_values.clone(), + } +} + +// Helper function to adjust child array validity based on parent struct validity +// When parent struct is null, propagates null to child array +// Optimized with fast paths and SIMD operations +fn adjust_child_validity( + child: &ArrayRef, + parent_validity: Option<&arrow_buffer::NullBuffer>, +) -> ArrayRef { + // Fast path: no parent validity means no adjustment needed + let parent_validity = match parent_validity { + None => return child.clone(), + Some(p) if p.null_count() == 0 => return child.clone(), // No nulls to propagate + Some(p) => p, + }; + + // Fast path: DataType::Null arrays are always entirely null by definition and cannot + // carry an explicit null bitmap (Arrow rejects it). No adjustment is needed. + if child.data_type() == &DataType::Null { + return child.clone(); + } + + let child_validity = child.nulls(); + + // Compute the new validity: child_validity AND parent_validity + let new_validity = match child_validity { + None => { + // Fast path: child has no existing validity, just use parent's + parent_validity.clone() + } + Some(child_nulls) => { + let child_buffer = child_nulls.inner(); + let parent_buffer = parent_validity.inner(); + + // Perform bitwise AND directly on BooleanBuffers + // This preserves the correct semantics: 1 = valid, 0 = null + let merged_buffer = child_buffer & parent_buffer; + + arrow_buffer::NullBuffer::from(merged_buffer) + } + }; + + // ArrayData keeps the child's offset for its value buffers. Its null bitmap + // must cover that offset too, even though the merged validity describes only + // the visible rows. Copy logical bits into an offset-aligned bitmap when needed. + let null_bit_buffer = if child.offset() == 0 && new_validity.inner().offset() == 0 { + new_validity.into_inner().into_inner() + } else { + let mut builder = + arrow_buffer::builder::BooleanBufferBuilder::new(child.offset() + child.len()); + builder.append_n(child.offset(), true); + builder.append_buffer(new_validity.inner()); + builder.build().into_inner() + }; + + // Create new array with adjusted validity + arrow_array::make_array( + arrow_data::ArrayData::try_new( + child.data_type().clone(), + child.len(), + Some(null_bit_buffer), + child.offset(), + child.to_data().buffers().to_vec(), + child.to_data().child_data().to_vec(), + ) + .unwrap(), + ) +} + +fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray { + let mut fields: Vec = vec![]; + let mut columns: Vec = vec![]; + let right_fields = right_struct_array.fields(); + let right_columns = right_struct_array.columns(); + + // Get the validity buffers from both structs + let left_validity = left_struct_array.nulls(); + let right_validity = right_struct_array.nulls(); + + // Compute merged validity + let merged_validity = merge_struct_validity(left_validity, right_validity); + + // iterate through the fields on the left hand side + for (left_field, left_column) in left_struct_array + .fields() + .iter() + .zip(left_struct_array.columns().iter()) + { + match right_fields + .iter() + .position(|f| f.name() == left_field.name()) + { + // if the field exists on the right hand side, merge them recursively if appropriate + Some(right_index) => { + let right_field = right_fields.get(right_index).unwrap(); + let right_column = right_columns.get(right_index).unwrap(); + // if both fields are struct, merge them recursively + match (left_field.data_type(), right_field.data_type()) { + (DataType::Struct(_), DataType::Struct(_)) => { + let left_sub_array = left_column.as_struct(); + let right_sub_array = right_column.as_struct(); + let merged_sub_array = merge(left_sub_array, right_sub_array); + fields.push(Field::new( + left_field.name(), + merged_sub_array.data_type().clone(), + left_field.is_nullable(), + )); + columns.push(Arc::new(merged_sub_array) as ArrayRef); + } + (DataType::List(left_list), DataType::List(right_list)) + | (DataType::LargeList(left_list), DataType::LargeList(right_list)) + if holds_struct(left_list.data_type()) + && holds_struct(right_list.data_type()) => + { + // Identical inner types: nothing to merge, keep the left column. + if left_list.data_type() == right_list.data_type() { + fields.push(left_field.as_ref().clone()); + columns.push(left_column.clone()); + } else { + // The struct fields differ, so merge them structurally. merge_list_struct + // only succeeds when both lists share offsets or one side is all-null; + // it panics otherwise. + let merged_sub_array = merge_list_struct(&left_column, &right_column); + + fields.push(Field::new( + left_field.name(), + merged_sub_array.data_type().clone(), + left_field.is_nullable(), + )); + columns.push(merged_sub_array); + } + } + // otherwise, just use the field on the left hand side + _ => { + // TODO handle list-of-struct and other types + fields.push(left_field.as_ref().clone()); + // Adjust the column validity: if left struct was null, propagate to child + let adjusted_column = adjust_child_validity(left_column, left_validity); + columns.push(adjusted_column); + } + } + } + None => { + fields.push(left_field.as_ref().clone()); + // Adjust the column validity: if left struct was null, propagate to child + let adjusted_column = adjust_child_validity(left_column, left_validity); + columns.push(adjusted_column); + } + } + } + + // now iterate through the fields on the right hand side + right_fields + .iter() + .zip(right_columns.iter()) + .for_each(|(field, column)| { + // add new columns on the right + if !left_struct_array + .fields() + .iter() + .any(|f| f.name() == field.name()) + { + fields.push(field.as_ref().clone()); + // This field doesn't exist on the left + // We use the right's column but need to adjust for struct validity + let adjusted_column = adjust_child_validity(column, right_validity); + columns.push(adjusted_column); + } + }); + + StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap() +} + +fn merge_with_schema( + left_struct_array: &StructArray, + right_struct_array: &StructArray, + fields: &Fields, +) -> StructArray { + // Helper function that returns true if both types are struct or both are non-struct + fn same_type_kind(left: &DataType, right: &DataType) -> bool { + match (left, right) { + (DataType::Struct(_), DataType::Struct(_)) => true, + (DataType::Struct(_), _) => false, + (_, DataType::Struct(_)) => false, + _ => true, + } + } + + let mut output_fields: Vec = Vec::with_capacity(fields.len()); + let mut columns: Vec = Vec::with_capacity(fields.len()); + + let left_fields = left_struct_array.fields(); + let left_columns = left_struct_array.columns(); + let right_fields = right_struct_array.fields(); + let right_columns = right_struct_array.columns(); + + // Get the validity buffers from both structs + let left_validity = left_struct_array.nulls(); + let right_validity = right_struct_array.nulls(); + + // Compute merged validity + let merged_validity = merge_struct_validity(left_validity, right_validity); + + for field in fields { + let left_match_idx = left_fields.iter().position(|f| { + f.name() == field.name() && same_type_kind(f.data_type(), field.data_type()) + }); + let right_match_idx = right_fields.iter().position(|f| { + f.name() == field.name() && same_type_kind(f.data_type(), field.data_type()) + }); + + match (left_match_idx, right_match_idx) { + (None, Some(right_idx)) => { + output_fields.push(right_fields[right_idx].as_ref().clone()); + // Adjust validity if the right struct was null + let adjusted_column = + adjust_child_validity(&right_columns[right_idx], right_validity); + columns.push(adjusted_column); + } + (Some(left_idx), None) => { + output_fields.push(left_fields[left_idx].as_ref().clone()); + // Adjust validity if the left struct was null + let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity); + columns.push(adjusted_column); + } + (Some(left_idx), Some(right_idx)) => { + match field.data_type() { + DataType::Struct(child_fields) => { + let left_sub_array = left_columns[left_idx].as_struct(); + let right_sub_array = right_columns[right_idx].as_struct(); + let merged_sub_array = + merge_with_schema(left_sub_array, right_sub_array, child_fields); + output_fields.push(Field::new( + field.name(), + merged_sub_array.data_type().clone(), + field.is_nullable(), + )); + columns.push(Arc::new(merged_sub_array) as ArrayRef); + } + DataType::List(child_field) => { + let left_list = left_columns[left_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let right_list = right_columns[right_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let merged_values = merge_list_child_values( + child_field.as_ref(), + left_list.trimmed_values(), + right_list.trimmed_values(), + ); + let merged_validity = + merge_struct_validity(left_list.nulls(), right_list.nulls()); + // `trimmed_values` starts at the first used value, so offsets + // must be shifted to match or `ListArray::new` panics when the + // input list was sliced (e.g. from a filtered batch). + let merged_list = ListArray::new( + child_field.clone(), + left_list.trimmed_offsets(), + merged_values, + merged_validity, + ); + output_fields.push(field.as_ref().clone()); + columns.push(Arc::new(merged_list) as ArrayRef); + } + DataType::LargeList(child_field) => { + let left_list = left_columns[left_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let right_list = right_columns[right_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let merged_values = merge_list_child_values( + child_field.as_ref(), + left_list.trimmed_values(), + right_list.trimmed_values(), + ); + let merged_validity = + merge_struct_validity(left_list.nulls(), right_list.nulls()); + let merged_list = LargeListArray::new( + child_field.clone(), + left_list.trimmed_offsets(), + merged_values, + merged_validity, + ); + output_fields.push(field.as_ref().clone()); + columns.push(Arc::new(merged_list) as ArrayRef); + } + DataType::FixedSizeList(child_field, list_size) => { + let left_list = left_columns[left_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let right_list = right_columns[right_idx] + .as_any() + .downcast_ref::() + .unwrap(); + let merged_values = merge_list_child_values( + child_field.as_ref(), + left_list.values().clone(), + right_list.values().clone(), + ); + let merged_validity = + merge_struct_validity(left_list.nulls(), right_list.nulls()); + let merged_list = FixedSizeListArray::new( + child_field.clone(), + *list_size, + merged_values, + merged_validity, + ); + output_fields.push(field.as_ref().clone()); + columns.push(Arc::new(merged_list) as ArrayRef); + } + _ => { + output_fields.push(left_fields[left_idx].as_ref().clone()); + // For fields that exist in both, use left but adjust validity + let adjusted_column = + adjust_child_validity(&left_columns[left_idx], left_validity); + columns.push(adjusted_column); + } + } + } + (None, None) => { + // The field will not be included in the output + } + } + } + + StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap() +} + +fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> { + if components.is_empty() { + return Some(array); + } + if !matches!(array.data_type(), DataType::Struct(_)) { + return None; + } + let struct_arr = array.as_struct(); + struct_arr + .column_by_name(components[0]) + .and_then(|arr| get_sub_array(arr, &components[1..])) +} + +/// Interleave multiple RecordBatches into a single RecordBatch. +/// +/// Behaves like [`arrow_select::interleave::interleave`], but for RecordBatches. +pub fn interleave_batches( + batches: &[RecordBatch], + indices: &[(usize, usize)], +) -> Result { + let first_batch = batches.first().ok_or_else(|| { + ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string()) + })?; + let schema = first_batch.schema(); + let num_columns = first_batch.num_columns(); + let mut columns = Vec::with_capacity(num_columns); + let mut chunks = Vec::with_capacity(batches.len()); + + for i in 0..num_columns { + for batch in batches { + chunks.push(batch.column(i).as_ref()); + } + let new_column = interleave(&chunks, indices)?; + columns.push(new_column); + chunks.clear(); + } + + RecordBatch::try_new(schema, columns) +} + +pub trait BufferExt { + /// Create an `arrow_buffer::Buffer`` from a `bytes::Bytes` object + /// + /// The alignment must be specified (as `bytes_per_value`) since we want to make + /// sure we can safely reinterpret the buffer. + /// + /// If the buffer is properly aligned this will be zero-copy. If not, a copy + /// will be made and an owned buffer returned. + /// + /// If `bytes_per_value` is not a power of two, then we assume the buffer is + /// never going to be reinterpreted into another type and we can safely + /// ignore the alignment. + /// + /// Yes, the method name is odd. It's because there is already a `from_bytes` + /// which converts from `arrow_buffer::bytes::Bytes` (not `bytes::Bytes`) + fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self; + + /// Allocates a new properly aligned arrow buffer and copies `bytes` into it + /// + /// `size_bytes` can be larger than `bytes` and, if so, the trailing bytes will + /// be zeroed out. + /// + /// # Panics + /// + /// Panics if `size_bytes` is less than `bytes.len()` + fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self; +} + +fn is_pwr_two(n: u64) -> bool { + n & (n - 1) == 0 +} + +impl BufferExt for arrow_buffer::Buffer { + fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self { + if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0 + { + // The original buffer is not aligned, cannot zero-copy + let size_bytes = bytes.len(); + Self::copy_bytes_bytes(bytes, size_bytes) + } else { + // The original buffer is aligned, can zero-copy + // SAFETY: the alignment is correct we can make this conversion + unsafe { + Self::from_custom_allocation( + NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"), + bytes.len(), + Arc::new(bytes), + ) + } + } + } + + fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self { + assert!(size_bytes >= bytes.len()); + let mut buf = MutableBuffer::with_capacity(size_bytes); + let to_fill = size_bytes - bytes.len(); + buf.extend(bytes); + buf.extend(std::iter::repeat_n(0_u8, to_fill)); + + // FIX for issue #4512: Shrink buffer to actual size before converting to immutable + // This reduces memory overhead from capacity over-allocation + buf.shrink_to_fit(); + + Self::from(buf) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::{ + BinaryArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, Float32Array, Int32Array, + LargeBinaryArray, ListArray, NullArray, StringArray, StructArray, new_empty_array, + new_null_array, + }; + use arrow_buffer::OffsetBuffer; + + #[test] + fn test_adjust_child_validity_with_sliced_child() { + let child: ArrayRef = Arc::new(BooleanArray::from_iter((0..1030).map(|index| { + if index == 1027 { + None + } else { + Some(index % 2 == 1) + } + }))); + let child = child.slice(1025, 4); + assert_eq!(child.offset(), 1025); + + let parent_validity = arrow_buffer::NullBuffer::from(vec![true, false, true, true]); + let adjusted = adjust_child_validity(&child, Some(&parent_validity)); + assert_eq!(adjusted.offset(), child.offset()); + let adjusted = adjusted.as_any().downcast_ref::().unwrap(); + assert_eq!( + adjusted.iter().collect::>(), + vec![Some(true), None, None, Some(false)] + ); + } + + #[test] + fn test_adjust_child_validity_with_sliced_parent_bitmap() { + let child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let parent_validity: arrow_buffer::NullBuffer = + (0..1030).map(|index| index != 1026).collect(); + let parent_validity = parent_validity.slice(1025, 4); + assert_eq!(parent_validity.inner().offset(), 1025); + + let adjusted = adjust_child_validity(&child, Some(&parent_validity)); + let adjusted = adjusted.as_any().downcast_ref::().unwrap(); + assert_eq!( + adjusted.iter().collect::>(), + vec![Some(1), None, Some(3), Some(4)] + ); + } + + #[test] + fn test_convert_to_floating_point_preserves_inner_nulls() { + // A FixedSizeList with a null inner element must convert to a + // FixedSizeList with the null kept in place. Dropping it would + // shorten the values array and shift every later element (and, when the + // remaining count is not a multiple of the list size, panic). + let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]); + let fsl = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Int8, true)), + 2, + Arc::new(values), + None, + ); + + let converted = fsl.convert_to_floating_point().unwrap(); + + assert_eq!(converted.len(), 2); + let conv_values = converted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(conv_values.len(), 4); + assert_eq!(conv_values.value(0), 1.0); + assert!(conv_values.is_null(1)); + assert_eq!(conv_values.value(2), 3.0); + assert_eq!(conv_values.value(3), 4.0); + } + + #[test] + fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() { + // The Float64-producing arms (Int64/UInt8/UInt32) share the same fix as the + // Float32 arms; cover one representative (UInt8 -> Float64) so both branch + // families are exercised. + let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]); + let fsl = FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::UInt8, true)), + 2, + Arc::new(values), + None, + ); + + let converted = fsl.convert_to_floating_point().unwrap(); + + assert_eq!(converted.len(), 2); + let conv_values = converted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(conv_values.len(), 4); + assert_eq!(conv_values.value(0), 10.0); + assert!(conv_values.is_null(1)); + assert_eq!(conv_values.value(2), 30.0); + assert_eq!(conv_values.value(3), 40.0); + } + + #[test] + fn test_merge_recursive() { + let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]); + let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]); + let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]); + let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]); + + let left_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()), + true, + ), + ]); + let left_batch = RecordBatch::try_new( + Arc::new(left_schema), + vec![ + Arc::new(a_array.clone()), + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("c", DataType::Int32, true)), + Arc::new(c_array.clone()) as ArrayRef, + )])), + ], + ) + .unwrap(); + + let right_schema = Schema::new(vec![ + Field::new("e", DataType::Int32, true), + Field::new( + "b", + DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()), + true, + ), + ]); + let right_batch = RecordBatch::try_new( + Arc::new(right_schema), + vec![ + Arc::new(e_array.clone()), + Arc::new(StructArray::from(vec![( + Arc::new(Field::new("d", DataType::Utf8, true)), + Arc::new(d_array.clone()) as ArrayRef, + )])) as ArrayRef, + ], + ) + .unwrap(); + + let merged_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::Struct( + vec![ + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Utf8, true), + ] + .into(), + ), + true, + ), + Field::new("e", DataType::Int32, true), + ]); + let merged_batch = RecordBatch::try_new( + Arc::new(merged_schema), + vec![ + Arc::new(a_array) as ArrayRef, + Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("c", DataType::Int32, true)), + Arc::new(c_array) as ArrayRef, + ), + ( + Arc::new(Field::new("d", DataType::Utf8, true)), + Arc::new(d_array) as ArrayRef, + ), + ])) as ArrayRef, + Arc::new(e_array) as ArrayRef, + ], + ) + .unwrap(); + + let result = left_batch.merge(&right_batch).unwrap(); + assert_eq!(result, merged_batch); + } + + #[test] + fn test_merge_with_schema() { + fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) { + let fields: Fields = names + .iter() + .zip(types) + .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false)) + .collect(); + let schema = Schema::new(vec![Field::new( + "struct", + DataType::Struct(fields.clone()), + false, + )]); + let children = types.iter().map(new_empty_array).collect::>(); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef], + ); + (schema, batch.unwrap()) + } + + let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]); + let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]); + let (output_schema, _) = test_batch( + &["b", "a", "c"], + &[DataType::Int64, DataType::Int32, DataType::Int32], + ); + + // If we use merge_with_schema the schema is respected + let merged = left_batch + .merge_with_schema(&right_batch, &output_schema) + .unwrap(); + assert_eq!(merged.schema().as_ref(), &output_schema); + + // If we use merge we get first-come first-serve based on the left batch + let (naive_schema, _) = test_batch( + &["a", "b", "c"], + &[DataType::Int32, DataType::Int64, DataType::Int32], + ); + let merged = left_batch.merge(&right_batch).unwrap(); + assert_eq!(merged.schema().as_ref(), &naive_schema); + } + + #[test] + fn test_merge_list_struct() { + let x_field = Arc::new(Field::new("x", DataType::Int32, true)); + let y_field = Arc::new(Field::new("y", DataType::Int32, true)); + let x_struct_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone()])), + true, + )); + let y_struct_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![y_field.clone()])), + true, + )); + let both_struct_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])), + true, + )); + let left_schema = Schema::new(vec![Field::new( + "list_struct", + DataType::List(x_struct_field.clone()), + true, + )]); + let right_schema = Schema::new(vec![Field::new( + "list_struct", + DataType::List(y_struct_field.clone()), + true, + )]); + let both_schema = Schema::new(vec![Field::new( + "list_struct", + DataType::List(both_struct_field.clone()), + true, + )]); + + let x = Arc::new(Int32Array::from(vec![1])); + let y = Arc::new(Int32Array::from(vec![2])); + let x_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone()]), + vec![x.clone()], + None, + )); + let y_struct = Arc::new(StructArray::new( + Fields::from(vec![y_field.clone()]), + vec![y.clone()], + None, + )); + let both_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone(), y_field.clone()]), + vec![x.clone(), y], + None, + )); + let both_null_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field, y_field]), + vec![x, Arc::new(new_null_array(&DataType::Int32, 1))], + None, + )); + let offsets = OffsetBuffer::from_lengths([1]); + let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None); + let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None); + let both_list = ListArray::new( + both_struct_field.clone(), + offsets.clone(), + both_struct, + None, + ); + let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None); + let x_batch = + RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap(); + let y_batch = RecordBatch::try_new( + Arc::new(right_schema.clone()), + vec![Arc::new(y_s_list.clone())], + ) + .unwrap(); + let merged = x_batch.merge(&y_batch).unwrap(); + let expected = + RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap(); + assert_eq!(merged, expected); + + let y_null_list = new_null_array(y_s_list.data_type(), 1); + let y_null_batch = + RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())]) + .unwrap(); + let expected = + RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap(); + let merged = x_batch.merge(&y_null_batch).unwrap(); + assert_eq!(merged, expected); + } + + #[test] + fn test_merge_all_null_list_adds_nested_child() { + // An all-null placeholder (what a fragment reader merges in for children + // missing from its files) adds `b` to a struct below one or two lists. + let strukt = |cols: Vec<(&str, ArrayRef)>| -> ArrayRef { + let fields = cols + .iter() + .map(|(name, col)| Field::new(*name, col.data_type().clone(), true)) + .collect::>(); + Arc::new(StructArray::new( + fields.into(), + cols.into_iter().map(|(_, col)| col).collect(), + None, + )) + }; + // Two rows: `[values]` and null. + let list = |values: ArrayRef, large: bool| -> ArrayRef { + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let lengths = [values.len(), 0]; + let nulls = Some(arrow_buffer::NullBuffer::from(vec![true, false])); + if large { + Arc::new(LargeListArray::new( + field, + OffsetBuffer::from_lengths(lengths), + values, + nulls, + )) + } else { + Arc::new(ListArray::new( + field, + OffsetBuffer::from_lengths(lengths), + values, + nulls, + )) + } + }; + let a = || Arc::new(Int32Array::from(vec![7])) as ArrayRef; + let b = || new_null_array(&DataType::Utf8, 1); + for large in [false, true] { + // list>>, built with or without `t`. + let under_list = |inner: ArrayRef, with_t: bool| { + let mut cols = vec![("in", inner)]; + if with_t { + cols.insert(0, ("t", Arc::new(Int32Array::from(vec![1])) as ArrayRef)); + } + list(strukt(cols), large) + }; + // list>> + let under_two_lists = |inner: ArrayRef, _: bool| { + let inner_list: ArrayRef = if large { + let field = Arc::new(Field::new("item", inner.data_type().clone(), true)); + Arc::new(LargeListArray::new( + field, + OffsetBuffer::from_lengths([1]), + inner, + None, + )) + } else { + let field = Arc::new(Field::new("item", inner.data_type().clone(), true)); + Arc::new(ListArray::new( + field, + OffsetBuffer::from_lengths([1]), + inner, + None, + )) + }; + list(inner_list, large) + }; + let shapes: [&dyn Fn(ArrayRef, bool) -> ArrayRef; 2] = [&under_list, &under_two_lists]; + for shape in shapes { + let data = shape(strukt(vec![("a", a())]), true); + let expected = shape(strukt(vec![("a", a()), ("b", b())]), true); + let placeholder = + new_null_array(shape(strukt(vec![("b", b())]), false).data_type(), 2); + let batch = |col: ArrayRef| { + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "c", + col.data_type().clone(), + true, + )])), + vec![col], + ) + .unwrap() + }; + let expected = batch(expected); + // Either side may be the placeholder; field order follows the left. + for merged in [ + batch(data.clone()) + .merge(&batch(placeholder.clone())) + .unwrap(), + batch(placeholder.clone()) + .merge(&batch(data.clone())) + .unwrap(), + ] { + let merged = merged.project_by_schema(expected.schema_ref()).unwrap(); + assert_eq!(merged, expected, "large={large}"); + } + } + } + } + + #[test] + fn test_merge_list_struct_identical_schema() { + // Merging two batches whose `List` columns have identical types + // should yield a single column equal to the input (there is nothing to + // merge), not a struct with the field pushed twice. + let x_field = Arc::new(Field::new("x", DataType::Int32, true)); + let item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone()])), + true, + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "list_struct", + DataType::List(item_field.clone()), + true, + )])); + + let build_list = |values: Vec| { + let len = values.len(); + let item_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone()]), + vec![Arc::new(Int32Array::from(values))], + None, + )); + ListArray::new( + item_field.clone(), + OffsetBuffer::from_lengths([len]), + item_struct, + None, + ) + }; + + // Distinct values so the equality assertion proves the left column is kept, + // not that some column with the same values happens to be present. + let left = + RecordBatch::try_new(schema.clone(), vec![Arc::new(build_list(vec![1, 2]))]).unwrap(); + let right = RecordBatch::try_new(schema, vec![Arc::new(build_list(vec![3, 4]))]).unwrap(); + + let merged = left.merge(&right).unwrap(); + + // Exactly one column, equal to the left input: the field must not be duplicated. + assert_eq!(merged.num_columns(), 1); + assert_eq!(merged, left); + } + + #[test] + fn test_merge_nested_list_struct_identical_schema() { + // A `List` nested inside a struct reaches the identical-type + // short-circuit through the recursive `merge` call. The recursion must + // keep the field exactly once, equal to the left input. + let x_field = Arc::new(Field::new("x", DataType::Int32, true)); + let item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![x_field.clone()])), + true, + )); + let companies_field = Arc::new(Field::new( + "companies", + DataType::List(item_field.clone()), + true, + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "outer", + DataType::Struct(Fields::from(vec![companies_field.clone()])), + true, + )])); + + let build_outer = |x: i32| { + // One list row: [{x}]. + let item_struct = Arc::new(StructArray::new( + Fields::from(vec![x_field.clone()]), + vec![Arc::new(Int32Array::from(vec![x]))], + None, + )); + let companies = Arc::new(ListArray::new( + item_field.clone(), + OffsetBuffer::from_lengths([1]), + item_struct, + None, + )); + StructArray::new( + Fields::from(vec![companies_field.clone()]), + vec![companies], + None, + ) + }; + + // Distinct values so equality proves the left column is kept. + let left = RecordBatch::try_new(schema.clone(), vec![Arc::new(build_outer(10))]).unwrap(); + let right = RecordBatch::try_new(schema, vec![Arc::new(build_outer(20))]).unwrap(); + + let merged = left.merge(&right).unwrap(); + + // The recursive identical-type merge keeps exactly one `companies` field, + // with no double-push inside the nested struct. + assert_eq!(merged.column(0).as_struct().num_columns(), 1); + assert_eq!(merged, left); + } + + #[test] + fn test_byte_width_opt() { + assert_eq!(DataType::Int32.byte_width_opt(), Some(4)); + assert_eq!(DataType::Int64.byte_width_opt(), Some(8)); + assert_eq!(DataType::Float32.byte_width_opt(), Some(4)); + assert_eq!(DataType::Float64.byte_width_opt(), Some(8)); + assert_eq!(DataType::Utf8.byte_width_opt(), None); + assert_eq!(DataType::Binary.byte_width_opt(), None); + assert_eq!( + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(), + None + ); + assert_eq!( + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3) + .byte_width_opt(), + Some(12) + ); + assert_eq!( + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4) + .byte_width_opt(), + Some(16) + ); + assert_eq!( + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5) + .byte_width_opt(), + None + ); + } + + #[test] + fn test_take_record_batch() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(StringArray::from_iter_values( + (0..20).map(|i| format!("str-{}", i)), + )), + ], + ) + .unwrap(); + let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap(); + assert_eq!( + taken, + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 5, 10])), + Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])), + ], + ) + .unwrap() + ) + } + + #[test] + fn test_schema_project_by_schema() { + let metadata = [("key".to_string(), "value".to_string())]; + let schema = Arc::new( + Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]) + .with_metadata(metadata.clone().into()), + ); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(StringArray::from_iter_values( + (0..20).map(|i| format!("str-{}", i)), + )), + ], + ) + .unwrap(); + + // Empty schema + let empty_schema = Schema::empty(); + let empty_projected = batch.project_by_schema(&empty_schema).unwrap(); + let expected_schema = empty_schema.with_metadata(metadata.clone().into()); + assert_eq!( + empty_projected, + RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None)) + .with_schema(Arc::new(expected_schema)) + .unwrap() + ); + + // Re-ordered schema + let reordered_schema = Schema::new(vec![ + Field::new("b", DataType::Utf8, true), + Field::new("a", DataType::Int32, true), + ]); + let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap(); + let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into())); + assert_eq!( + reordered_projected, + RecordBatch::try_new( + expected_schema, + vec![ + Arc::new(StringArray::from_iter_values( + (0..20).map(|i| format!("str-{}", i)), + )), + Arc::new(Int32Array::from_iter_values(0..20)), + ], + ) + .unwrap() + ); + + // Sub schema + let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let sub_projected = batch.project_by_schema(&sub_schema).unwrap(); + let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into())); + assert_eq!( + sub_projected, + RecordBatch::try_new( + expected_schema, + vec![Arc::new(Int32Array::from_iter_values(0..20))], + ) + .unwrap() + ); + } + + #[test] + fn test_project_rebuilds_sliced_map() { + // A sliced MapArray keeps its full entries array behind sliced + // offsets and validity; the Map projection arm must rebuild it + // without renormalizing either. + let entry_fields = Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ]); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(entry_fields.clone()), + false, + )); + let entries = StructArray::new( + entry_fields, + vec![ + Arc::new(StringArray::from(vec!["k0", "k1", "k2"])) as ArrayRef, + Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef, + ], + None, + ); + let map = MapArray::new( + entries_field.clone(), + OffsetBuffer::new(vec![0, 1, 1, 3].into()), + entries, + Some(arrow_buffer::NullBuffer::from(vec![true, false, true])), + false, + ); + let schema = Arc::new(Schema::new(vec![Field::new( + "m", + DataType::Map(entries_field, false), + true, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(map) as ArrayRef]).unwrap(); + + // Rows 1..3: a null slot and a two-entry slot, offsets not zero-based. + let projected = batch + .slice(1, 2) + .project_by_schema(schema.as_ref()) + .unwrap(); + let map = projected.column(0).as_map(); + assert!(map.is_null(0)); + assert!(map.is_valid(1)); + assert_eq!(map.value_length(1), 2); + assert_eq!( + map.value(1) + .column(1) + .as_primitive::() + .values(), + &[2, 3] + ); + } + + #[test] + fn test_project_preserves_struct_validity() { + // Test that projecting a struct array preserves its validity (fix for issue #4385) + let fields = Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Float32, true), + ]); + + // Create a struct array with validity + let id_array = Int32Array::from(vec![1, 2, 3]); + let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]); + let struct_array = StructArray::new( + fields.clone(), + vec![ + Arc::new(id_array) as ArrayRef, + Arc::new(value_array) as ArrayRef, + ], + Some(vec![true, false, true].into()), // Second struct is null + ); + + // Project the struct array + let projected = project(&struct_array, &fields).unwrap(); + + // Verify the validity is preserved + assert_eq!(projected.null_count(), 1); + assert!(!projected.is_null(0)); + assert!(projected.is_null(1)); + assert!(!projected.is_null(2)); + } + + #[test] + fn test_merge_struct_with_different_validity() { + // Test case from Weston's review comment + // File 1 has height field with some nulls + let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]); + let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]); + let left_struct = StructArray::new( + left_fields, + vec![Arc::new(height_array) as ArrayRef], + Some(vec![true, false, true, false].into()), // Rows 2 and 4 are null structs + ); + + // File 2 has width field with some nulls + let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]); + let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]); + let right_struct = StructArray::new( + right_fields, + vec![Arc::new(width_array) as ArrayRef], + Some(vec![true, true, false, false].into()), // Rows 3 and 4 are null structs + ); + + // Merge the two structs + let merged = merge(&left_struct, &right_struct); + + // Expected: + // Row 1: both non-null -> {width: 300, height: 500} + // Row 2: left null, right non-null -> {width: 200, height: null} + // Row 3: left non-null, right null -> {width: null, height: 600} + // Row 4: both null -> null struct + + assert_eq!(merged.null_count(), 1); // Only row 4 is null + assert!(!merged.is_null(0)); + assert!(!merged.is_null(1)); + assert!(!merged.is_null(2)); + assert!(merged.is_null(3)); + + // Check field values + let height_col = merged.column_by_name("height").unwrap(); + let height_values = height_col.as_any().downcast_ref::().unwrap(); + assert_eq!(height_values.value(0), 500); + assert!(height_values.is_null(1)); // height is null when left struct was null + assert_eq!(height_values.value(2), 600); + + let width_col = merged.column_by_name("width").unwrap(); + let width_values = width_col.as_any().downcast_ref::().unwrap(); + assert_eq!(width_values.value(0), 300); + assert_eq!(width_values.value(1), 200); + assert!(width_values.is_null(2)); // width is null when right struct was null + + // An all-null validity buffer is data, not a placeholder meaning "this side has no + // validity": merging two of them keeps the rows null. + let all_null_left = StructArray::new( + Fields::from(vec![Field::new("height", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef], + Some(vec![false, false].into()), + ); + let all_null_right = StructArray::new( + Fields::from(vec![Field::new("width", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef], + Some(vec![false, false].into()), + ); + + let merged = merge(&all_null_left, &all_null_right); + assert_eq!(merged.null_count(), 2); + + // An all-null side is the identity of the merge, so the other side decides each row. + let partial_left = StructArray::new( + Fields::from(vec![Field::new("height", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef], + Some(vec![true, false].into()), + ); + let merged = merge(&partial_left, &all_null_right); + assert!(!merged.is_null(0)); + assert!(merged.is_null(1)); + + // A missing validity buffer means all rows are valid, which absorbs an all-null side. + let all_valid_left = StructArray::new( + Fields::from(vec![Field::new("height", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef], + None, + ); + let merged = merge(&all_valid_left, &all_null_right); + assert_eq!(merged.null_count(), 0); + + // An explicit all-valid buffer has the same semantics as a missing buffer. + let all_valid: arrow_buffer::NullBuffer = vec![true, true].into(); + let partial: arrow_buffer::NullBuffer = vec![true, false].into(); + assert!(merge_struct_validity(Some(&all_valid), Some(&partial)).is_none()); + assert!(merge_struct_validity(Some(&partial), Some(&all_valid)).is_none()); + } + + #[test] + fn test_merge_null_typed_column_with_parent_validity() { + // Reproduces ENT-990: panic in adjust_child_validity when a Null-typed column + // exists on one side and the parent struct has null rows. + // Arrow's Null type has no null bitmap, so passing one to ArrayData::try_new panics. + let left_struct = StructArray::new( + Fields::from(vec![Field::new("a", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef], + Some(vec![true, false].into()), + ); + let right_struct = StructArray::new( + Fields::from(vec![Field::new("b", DataType::Null, true)]), + vec![Arc::new(NullArray::new(2)) as ArrayRef], + Some(vec![true, false].into()), + ); + + // Previously panicked: "Arrays of type Null cannot contain a null bitmask" + let merged = merge(&left_struct, &right_struct); + assert_eq!(merged.len(), 2); + let b_col = merged.column_by_name("b").unwrap(); + // DataType::Null implies all-null by definition; no null bitmap is stored. + assert_eq!(b_col.data_type(), &DataType::Null); + assert_eq!(b_col.len(), 2); + } + + #[test] + fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() { + // left_list setup + let left_company_id = Arc::new(Int32Array::from(vec![None, None])); + let left_count = Arc::new(Int32Array::from(vec![None, None])); + let left_struct = Arc::new(StructArray::new( + Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("count", DataType::Int32, true), + ]), + vec![left_company_id, left_count], + None, + )); + let left_list = Arc::new(ListArray::new( + Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([2]), + left_struct, + None, + )); + + // Right List Setup + let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"])); + let right_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]), + vec![right_company_name], + None, + )); + let right_list = Arc::new(ListArray::new( + Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([2]), + right_struct, + None, + )); + + let target_fields = Fields::from(vec![Field::new( + "companies", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("company_name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + true, + ))), + true, + )]); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "companies", + left_list.data_type().clone(), + true, + )])), + vec![left_list as ArrayRef], + ) + .unwrap(); + + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "companies", + right_list.data_type().clone(), + true, + )])), + vec![right_list as ArrayRef], + ) + .unwrap(); + + let merged = left_batch + .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec())) + .unwrap(); + + // Verify the merged structure + let merged_list = merged + .column_by_name("companies") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let merged_struct = merged_list.values().as_struct(); + + // Should have all 3 fields + assert_eq!(merged_struct.num_columns(), 3); + assert!(merged_struct.column_by_name("company_id").is_some()); + assert!(merged_struct.column_by_name("company_name").is_some()); + assert!(merged_struct.column_by_name("count").is_some()); + + // Verify values + let company_id = merged_struct + .column_by_name("company_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(company_id.is_null(0)); + assert!(company_id.is_null(1)); + + let company_name = merged_struct + .column_by_name("company_name") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(company_name.value(0), "Google"); + assert_eq!(company_name.value(1), "Microsoft"); + + let count = merged_struct + .column_by_name("count") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(count.is_null(0)); + assert!(count.is_null(1)); + } + + #[test] + fn test_merge_struct_lists() { + test_merge_struct_lists_generic::(); + } + + #[test] + fn test_merge_struct_large_lists() { + test_merge_struct_lists_generic::(); + } + + fn test_merge_struct_lists_generic() { + // left_list setup + let left_company_id = Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), + Some(10), + Some(11), + Some(12), + Some(13), + Some(14), + Some(15), + Some(16), + Some(17), + Some(18), + Some(19), + Some(20), + ])); + let left_count = Arc::new(Int32Array::from(vec![ + Some(10), + Some(20), + Some(30), + Some(40), + Some(50), + Some(60), + Some(70), + Some(80), + Some(90), + Some(100), + Some(110), + Some(120), + Some(130), + Some(140), + Some(150), + Some(160), + Some(170), + Some(180), + Some(190), + Some(200), + ])); + let left_struct = Arc::new(StructArray::new( + Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("count", DataType::Int32, true), + ]), + vec![left_company_id, left_count], + None, + )); + + let left_list = Arc::new(GenericListArray::::new( + Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([3, 1]), + left_struct.clone(), + None, + )); + + let left_list_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new( + "companies", + if O::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(left_struct.fields().clone()), + true, + ))) + }, + true, + )]), + vec![left_list as ArrayRef], + None, + )); + + // right_list setup + let right_company_name = Arc::new(StringArray::from(vec![ + "Google", + "Microsoft", + "Apple", + "Facebook", + ])); + let right_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]), + vec![right_company_name], + None, + )); + let right_list = Arc::new(GenericListArray::::new( + Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + )), + OffsetBuffer::from_lengths([3, 1]), + right_struct.clone(), + None, + )); + + let right_list_struct = Arc::new(StructArray::new( + Fields::from(vec![Field::new( + "companies", + if O::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(right_struct.fields().clone()), + true, + ))) + }, + true, + )]), + vec![right_list as ArrayRef], + None, + )); + + // prepare schema + let target_fields = Fields::from(vec![Field::new( + "companies", + if O::IS_LARGE { + DataType::LargeList(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("company_name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + true, + ))) + } else { + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("company_id", DataType::Int32, true), + Field::new("company_name", DataType::Utf8, true), + Field::new("count", DataType::Int32, true), + ])), + true, + ))) + }, + true, + )]); + + // merge left_list and right_list + let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields); + assert_eq!(merged_array.len(), 2); + } + + #[test] + fn test_merge_with_schema_sliced_list_struct() { + test_merge_with_schema_sliced_list_struct_generic::(); + } + + #[test] + fn test_merge_with_schema_sliced_large_list_struct() { + test_merge_with_schema_sliced_list_struct_generic::(); + } + + // Regression for #6580: merge_with_schema panicked when the left list was a + // sliced view whose offsets did not start at zero (common after a filtered + // scan). Cloning those offsets alongside `trimmed_values` produced offsets + // larger than the trimmed child, panicking in `(Large)ListArray::new`. + fn test_merge_with_schema_sliced_list_struct_generic() { + let make_list_dtype = |item_field: Arc| { + if O::IS_LARGE { + DataType::LargeList(item_field) + } else { + DataType::List(item_field) + } + }; + + // Build a List with two rows of 5 items each, then slice away + // the first row so the remaining list's offsets start at 5, not 0. + let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]); + let left_values = Arc::new(StructArray::new( + struct_fields_a.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef], + None, + )); + let full_list = GenericListArray::::new( + Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)), + OffsetBuffer::::from_lengths([5, 5]), + left_values, + None, + ); + let sliced_left = full_list.slice(1, 1); + assert_eq!(sliced_left.offsets()[0].as_usize(), 5); + assert_eq!(sliced_left.offsets()[1].as_usize(), 10); + + let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]); + let right_values = Arc::new(StructArray::new( + struct_fields_b.clone(), + vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef], + None, + )); + let right_list = GenericListArray::::new( + Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)), + OffsetBuffer::::from_lengths([5]), + right_values, + None, + ); + + let target_item_field = Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])), + true, + )); + let target_fields = Fields::from(vec![Field::new( + "items", + make_list_dtype(target_item_field), + true, + )]); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + sliced_left.data_type().clone(), + true, + )])), + vec![Arc::new(sliced_left) as ArrayRef], + ) + .unwrap(); + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "items", + right_list.data_type().clone(), + true, + )])), + vec![Arc::new(right_list) as ArrayRef], + ) + .unwrap(); + + let merged = left_batch + .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec())) + .unwrap(); + + let merged_list = merged + .column_by_name("items") + .unwrap() + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(merged_list.len(), 1); + assert_eq!(merged_list.value_length(0).as_usize(), 5); + let merged_struct = merged_list.values().as_struct(); + assert_eq!(merged_struct.num_columns(), 2); + let a = merged_struct + .column_by_name("a") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + // After shifting offsets to zero, values 5..10 should be first. + let a_vals: Vec = a.iter().map(|v| v.unwrap()).collect(); + assert_eq!(a_vals, vec![5, 6, 7, 8, 9]); + } + + #[test] + fn test_project_by_schema_list_struct_reorder() { + // Test that project_by_schema correctly reorders fields inside List + // This is a regression test for issue #5702 + + // Source schema with inner struct fields in order: c, b, a + let source_inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("c", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + Field::new("a", DataType::Utf8, true), + ])); + let source_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "data", + DataType::List(Arc::new(Field::new( + "item", + source_inner_struct.clone(), + true, + ))), + true, + ), + ])); + + // Create source data with c, b, a order + let c_array = StringArray::from(vec!["c1", "c2"]); + let b_array = StringArray::from(vec!["b1", "b2"]); + let a_array = StringArray::from(vec!["a1", "a2"]); + let inner_struct = StructArray::from(vec![ + ( + Arc::new(Field::new("c", DataType::Utf8, true)), + Arc::new(c_array) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, true)), + Arc::new(b_array) as ArrayRef, + ), + ( + Arc::new(Field::new("a", DataType::Utf8, true)), + Arc::new(a_array) as ArrayRef, + ), + ]); + + let list_array = ListArray::new( + Arc::new(Field::new("item", source_inner_struct, true)), + OffsetBuffer::from_lengths([1, 1]), + Arc::new(inner_struct), + None, + ); + + let batch = RecordBatch::try_new( + source_schema, + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)], + ) + .unwrap(); + + // Target schema with inner struct fields in order: a, b, c + let target_inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Utf8, true), + ])); + let target_schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "data", + DataType::List(Arc::new(Field::new("item", target_inner_struct, true))), + true, + ), + ]); + + // Project should reorder the inner struct fields + let projected = batch.project_by_schema(&target_schema).unwrap(); + + // Verify the schema is correct + assert_eq!(projected.schema().as_ref(), &target_schema); + + // Verify the data is correct by checking inner struct field order + let projected_list = projected.column(1).as_list::(); + let projected_struct = projected_list.values().as_struct(); + + // Fields should now be in order: a, b, c + assert_eq!( + projected_struct.column_by_name("a").unwrap().as_ref(), + &StringArray::from(vec!["a1", "a2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column_by_name("b").unwrap().as_ref(), + &StringArray::from(vec!["b1", "b2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column_by_name("c").unwrap().as_ref(), + &StringArray::from(vec!["c1", "c2"]) as &dyn Array + ); + + // Also verify positional access matches expected order (a=0, b=1, c=2) + assert_eq!( + projected_struct.column(0).as_ref(), + &StringArray::from(vec!["a1", "a2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column(1).as_ref(), + &StringArray::from(vec!["b1", "b2"]) as &dyn Array + ); + assert_eq!( + projected_struct.column(2).as_ref(), + &StringArray::from(vec!["c1", "c2"]) as &dyn Array + ); + } + + #[test] + fn test_project_by_schema_nested_list_struct() { + // Test deeply nested List>> projection + let inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("y", DataType::Int32, true), + Field::new("x", DataType::Int32, true), + ])); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "outer", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("b", DataType::Utf8, true), + Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))), + true, + ), + Field::new("a", DataType::Utf8, true), + ])), + true, + ))), + true, + )])); + + // Create deeply nested data + let y_array = Int32Array::from(vec![1, 2]); + let x_array = Int32Array::from(vec![3, 4]); + let innermost_struct = StructArray::from(vec![ + ( + Arc::new(Field::new("y", DataType::Int32, true)), + Arc::new(y_array) as ArrayRef, + ), + ( + Arc::new(Field::new("x", DataType::Int32, true)), + Arc::new(x_array) as ArrayRef, + ), + ]); + let inner_list = ListArray::new( + Arc::new(Field::new("item", inner_struct.clone(), true)), + OffsetBuffer::from_lengths([2]), + Arc::new(innermost_struct), + None, + ); + + let b_array = StringArray::from(vec!["b1"]); + let a_array = StringArray::from(vec!["a1"]); + let middle_struct = StructArray::from(vec![ + ( + Arc::new(Field::new("b", DataType::Utf8, true)), + Arc::new(b_array) as ArrayRef, + ), + ( + Arc::new(Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", inner_struct, true))), + true, + )), + Arc::new(inner_list) as ArrayRef, + ), + ( + Arc::new(Field::new("a", DataType::Utf8, true)), + Arc::new(a_array) as ArrayRef, + ), + ]); + + let outer_list = ListArray::new( + Arc::new(Field::new("item", middle_struct.data_type().clone(), true)), + OffsetBuffer::from_lengths([1]), + Arc::new(middle_struct), + None, + ); + + let batch = + RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap(); + + // Target schema with reordered fields at all levels + let target_inner_struct = DataType::Struct(Fields::from(vec![ + Field::new("x", DataType::Int32, true), // x before y now + Field::new("y", DataType::Int32, true), + ])); + let target_schema = Schema::new(vec![Field::new( + "outer", + DataType::List(Arc::new(Field::new( + "item", + DataType::Struct(Fields::from(vec![ + Field::new("a", DataType::Utf8, true), // a before b now + Field::new( + "inner_list", + DataType::List(Arc::new(Field::new("item", target_inner_struct, true))), + true, + ), + Field::new("b", DataType::Utf8, true), + ])), + true, + ))), + true, + )]); + + let projected = batch.project_by_schema(&target_schema).unwrap(); + + // Verify schema + assert_eq!(projected.schema().as_ref(), &target_schema); + + // Verify deeply nested data is reordered correctly + let outer_list = projected.column(0).as_list::(); + let middle_struct = outer_list.values().as_struct(); + + // Middle struct should have a first, then inner_list, then b + assert_eq!( + middle_struct.column(0).as_ref(), + &StringArray::from(vec!["a1"]) as &dyn Array + ); + assert_eq!( + middle_struct.column(2).as_ref(), + &StringArray::from(vec!["b1"]) as &dyn Array + ); + + // Inner list's struct should have x first, then y + let inner_list = middle_struct.column(1).as_list::(); + let innermost_struct = inner_list.values().as_struct(); + assert_eq!( + innermost_struct.column(0).as_ref(), + &Int32Array::from(vec![3, 4]) as &dyn Array + ); + assert_eq!( + innermost_struct.column(1).as_ref(), + &Int32Array::from(vec![1, 2]) as &dyn Array + ); + } + + #[test] + fn test_iter_binary_array_accepts_binary_variants() { + let binary = BinaryArray::from(vec![b"a".as_slice(), b"bc"]); + assert_eq!( + iter_binary_array(&binary).unwrap().collect::>(), + vec![Some(b"a".as_slice()), Some(b"bc".as_slice())] + ); + + let large_binary = LargeBinaryArray::from(vec![b"x".as_slice(), b"yz"]); + assert_eq!( + iter_binary_array(&large_binary) + .unwrap() + .collect::>(), + vec![Some(b"x".as_slice()), Some(b"yz".as_slice())] + ); + + let binary_view = BinaryViewArray::from(vec![b"1".as_slice(), b"23"]); + assert_eq!( + iter_binary_array(&binary_view).unwrap().collect::>(), + vec![Some(b"1".as_slice()), Some(b"23".as_slice())] + ); + + let fixed_size = + FixedSizeBinaryArray::try_from_iter([b"abcd", b"efgh"].into_iter()).unwrap(); + assert_eq!( + iter_binary_array(&fixed_size).unwrap().collect::>(), + vec![Some(b"abcd".as_slice()), Some(b"efgh".as_slice())] + ); + } + + #[test] + fn test_iter_binary_array_rejects_non_binary() { + let int_array = Int32Array::from(vec![1, 2, 3]); + let Err(error) = iter_binary_array(&int_array) else { + panic!("expected an error for non-binary array"); + }; + assert!(error.to_string().contains("Expecting a binary type")); + } +} diff --git a/python/vendor/lance-arrow/src/list.rs b/python/vendor/lance-arrow/src/list.rs new file mode 100644 index 000000000..750b54c93 --- /dev/null +++ b/python/vendor/lance-arrow/src/list.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use arrow_array::{Array, BooleanArray, GenericListArray, OffsetSizeTrait}; +use arrow_buffer::{BooleanBufferBuilder, OffsetBuffer, ScalarBuffer}; +use arrow_schema::Field; + +pub trait ListArrayExt { + /// Filters out masked null items from the list array + /// + /// It is legal for a list array to have a null entry with a non-zero length. The + /// values inside the entry are "garbage" and should be ignored. This function + /// filters the values array to remove the garbage values. + /// + /// The output list will always have zero-length nulls. + fn filter_garbage_nulls(&self) -> Self; + /// Returns a copy of the list's values array that has been sliced to size + /// + /// It is legal for a list array's offsets to not start with zero. It's also legal + /// for a list array's offsets to not extend to the entire values array. This function + /// behaves similarly to `values()` except it slices the array so that it starts at + /// the first list offset and ends at the last list offset. + fn trimmed_values(&self) -> Arc; + /// The offset type of the underlying list array. + type Offset: OffsetSizeTrait; + /// Returns offsets shifted so the first offset is zero, matching + /// [`Self::trimmed_values`]. + /// + /// Sliced list arrays (e.g. a filtered batch) keep offsets that reference the + /// original values buffer, so combining them with trimmed values produces + /// offsets that exceed the values length. Use this together with + /// `trimmed_values` when constructing a new list array. + fn trimmed_offsets(&self) -> OffsetBuffer; +} + +impl ListArrayExt for GenericListArray { + fn filter_garbage_nulls(&self) -> Self { + if self.is_empty() { + return self.clone(); + } + let Some(validity) = self.nulls().cloned() else { + return self.clone(); + }; + + let mut should_keep = BooleanBufferBuilder::new(self.values().len()); + + // Handle case where offsets do not start at 0 + let preamble_len = self.offsets().first().unwrap().to_usize().unwrap(); + should_keep.append_n(preamble_len, false); + + let mut new_offsets: Vec = Vec::with_capacity(self.len() + 1); + new_offsets.push(OffsetSize::zero()); + let mut cur_len = OffsetSize::zero(); + for (offset, is_valid) in self.offsets().windows(2).zip(validity.iter()) { + let len = offset[1] - offset[0]; + if is_valid { + cur_len += len; + should_keep.append_n(len.to_usize().unwrap(), true); + new_offsets.push(cur_len); + } else { + should_keep.append_n(len.to_usize().unwrap(), false); + new_offsets.push(cur_len); + } + } + + // Offsets may not reference entire values buffer + let trailer = self.values().len() - should_keep.len(); + should_keep.append_n(trailer, false); + + let should_keep = should_keep.finish(); + let should_keep = BooleanArray::new(should_keep, None); + let new_values = arrow_select::filter::filter(self.values(), &should_keep).unwrap(); + let new_offsets = ScalarBuffer::from(new_offsets); + let new_offsets = OffsetBuffer::new(new_offsets); + + Self::new( + Arc::new(Field::new( + "item", + self.value_type(), + self.values().is_nullable(), + )), + new_offsets, + new_values, + Some(validity), + ) + } + + fn trimmed_values(&self) -> Arc { + let first_value = self + .offsets() + .first() + .map(|v| v.to_usize().unwrap()) + .unwrap_or(0); + let last_value = self + .offsets() + .last() + .map(|v| v.to_usize().unwrap()) + .unwrap_or(0); + self.values().slice(first_value, last_value - first_value) + } + + type Offset = OffsetSize; + + fn trimmed_offsets(&self) -> OffsetBuffer { + let offsets = self.offsets(); + let Some(&first) = offsets.first() else { + return offsets.clone(); + }; + if first == OffsetSize::zero() { + return offsets.clone(); + } + let shifted: Vec = offsets.iter().map(|&o| o - first).collect(); + OffsetBuffer::new(ScalarBuffer::from(shifted)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ListArray, UInt64Array}; + use arrow_buffer::{BooleanBuffer, NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field}; + + use super::ListArrayExt; + + #[test] + fn test_filter_garbage_nulls() { + let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + let offsets = ScalarBuffer::::from(vec![2, 5, 8, 9]); + let offsets = OffsetBuffer::new(offsets); + let list_validity = NullBuffer::new(BooleanBuffer::from(vec![true, false, true])); + let list_arr = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, true)), + offsets, + Arc::new(items), + Some(list_validity.clone()), + ); + + let filtered = list_arr.filter_garbage_nulls(); + + let expected_items = UInt64Array::from(vec![2, 3, 4, 8]); + let offsets = ScalarBuffer::::from(vec![0, 3, 3, 4]); + let expected = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, false)), + OffsetBuffer::new(offsets), + Arc::new(expected_items), + Some(list_validity), + ); + + assert_eq!(filtered, expected); + } + + #[test] + fn test_trim_values() { + let items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + let offsets = ScalarBuffer::::from(vec![2, 5, 6, 8, 9]); + let offsets = OffsetBuffer::new(offsets); + let list_arr = ListArray::new( + Arc::new(Field::new("item", DataType::UInt64, true)), + offsets, + Arc::new(items), + None, + ); + let list_arr = list_arr.slice(1, 2); + + let trimmed = list_arr.trimmed_values(); + + let expected_items = UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); + let expected_items = expected_items.slice(5, 3); + + assert_eq!(trimmed.as_ref(), &expected_items); + } +} diff --git a/python/vendor/lance-arrow/src/memory.rs b/python/vendor/lance-arrow/src/memory.rs new file mode 100644 index 000000000..1038a3f91 --- /dev/null +++ b/python/vendor/lance-arrow/src/memory.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::collections::HashSet; + +use arrow_array::{Array, RecordBatch}; +use arrow_data::ArrayData; + +/// Counts memory used by buffers of Arrow arrays and RecordBatches. +/// +/// This is meant to capture how much memory is being used by the Arrow data +/// structures as they are. It does not represent the memory used if the data +/// were to be serialized and then deserialized. In particular: +/// +/// * This does not double count memory used by buffers shared by multiple +/// arrays or batches. Round-tripped data may use more memory because of this. +/// * This counts the **total** size of the buffers, even if the array is a slice. +/// Round-tripped data may use less memory because of this. +#[derive(Default)] +pub struct MemoryAccumulator { + seen: HashSet, + total: usize, +} + +impl MemoryAccumulator { + pub fn record_array(&mut self, array: &dyn Array) { + let data = array.to_data(); + self.record_array_data(&data); + } + + fn record_array_data(&mut self, data: &ArrayData) { + for buffer in data.buffers() { + let ptr = buffer.as_ptr(); + if self.seen.insert(ptr as usize) { + self.total += buffer.capacity(); + } + } + + if let Some(nulls) = data.nulls() { + let null_buf = nulls.inner().inner(); + let ptr = null_buf.as_ptr(); + if self.seen.insert(ptr as usize) { + self.total += null_buf.capacity(); + } + } + + for child in data.child_data() { + self.record_array_data(child); + } + } + + pub fn record_batch(&mut self, batch: &RecordBatch) { + for array in batch.columns() { + self.record_array(array); + } + } + + pub fn total(&self) -> usize { + self.total + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::Int32Array; + use arrow_schema::{DataType, Field, Schema}; + + use super::*; + + #[test] + fn test_memory_accumulator() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + let slice = batch.slice(1, 2); + + let mut acc = MemoryAccumulator::default(); + + // Should record whole buffer, not just slice + acc.record_batch(&slice); + assert_eq!(acc.total(), 3 * std::mem::size_of::()); + + // Should not double count + acc.record_batch(&slice); + assert_eq!(acc.total(), 3 * std::mem::size_of::()); + } +} diff --git a/python/vendor/lance-arrow/src/scalar.rs b/python/vendor/lance-arrow/src/scalar.rs new file mode 100644 index 000000000..412253d12 --- /dev/null +++ b/python/vendor/lance-arrow/src/scalar.rs @@ -0,0 +1,302 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_array::{ArrayRef, UInt64Array, make_array}; +use arrow_buffer::Buffer; +use arrow_data::ArrayDataBuilder; +use arrow_schema::{ArrowError, DataType}; +use arrow_select::take::take; + +use crate::DataTypeExt; + +type Result = std::result::Result; + +pub const INLINE_VALUE_MAX_BYTES: usize = 32; + +pub fn extract_scalar_value(array: &ArrayRef, idx: usize) -> Result { + if idx >= array.len() { + return Err(ArrowError::InvalidArgumentError( + "Scalar index out of bounds".to_string(), + )); + } + + take(array.as_ref(), &UInt64Array::from(vec![idx as u64]), None) +} + +fn read_u32(buf: &[u8], offset: &mut usize) -> Result { + if *offset + 4 > buf.len() { + return Err(ArrowError::InvalidArgumentError( + "Invalid scalar value buffer: unexpected EOF".to_string(), + )); + } + let bytes = [ + buf[*offset], + buf[*offset + 1], + buf[*offset + 2], + buf[*offset + 3], + ]; + *offset += 4; + Ok(u32::from_le_bytes(bytes)) +} + +fn read_bytes<'a>(buf: &'a [u8], offset: &mut usize, len: usize) -> Result<&'a [u8]> { + if *offset + len > buf.len() { + return Err(ArrowError::InvalidArgumentError( + "Invalid scalar value buffer: unexpected EOF".to_string(), + )); + } + let slice = &buf[*offset..*offset + len]; + *offset += len; + Ok(slice) +} + +fn write_u32(out: &mut Vec, v: u32) { + out.extend_from_slice(&v.to_le_bytes()); +} + +fn write_bytes(out: &mut Vec, bytes: &[u8]) { + out.extend_from_slice(bytes); +} + +pub fn encode_scalar_value_buffer(scalar: &ArrayRef) -> Result> { + if scalar.len() != 1 || scalar.null_count() != 0 { + return Err(ArrowError::InvalidArgumentError( + "Scalar value buffer must be a single non-null value".to_string(), + )); + } + let data = scalar.to_data(); + if data.offset() != 0 { + return Err(ArrowError::InvalidArgumentError( + "Scalar value buffer must have offset=0".to_string(), + )); + } + if !data.child_data().is_empty() { + return Err(ArrowError::InvalidArgumentError( + "Scalar value buffer does not support nested types".to_string(), + )); + } + + // Minimal format (RFC): store the Arrow value buffers for a length-1 array. + // Null bitmap and child data are intentionally not supported here. + // + // | u32 num_buffers | + // | u32 buffer_0_len | ... | u32 buffer_{n-1}_len | + // | buffer_0 bytes | ... | buffer_{n-1} bytes | + let mut out = Vec::with_capacity(128); + let buffers = data.buffers(); + write_u32(&mut out, buffers.len() as u32); + for b in buffers { + write_u32(&mut out, b.len() as u32); + } + for b in buffers { + write_bytes(&mut out, b.as_slice()); + } + Ok(out) +} + +pub fn decode_scalar_from_value_buffer( + data_type: &DataType, + value_buffer: &[u8], +) -> Result { + if matches!( + data_type, + DataType::Struct(_) | DataType::FixedSizeList(_, _) + ) { + return Err(ArrowError::InvalidArgumentError(format!( + "Scalar value buffer does not support nested data type {:?}", + data_type + ))); + } + + let mut offset = 0; + let num_buffers = read_u32(value_buffer, &mut offset)? as usize; + let buffer_lens = (0..num_buffers) + .map(|_| read_u32(value_buffer, &mut offset).map(|l| l as usize)) + .collect::>>()?; + + let mut buffers = Vec::with_capacity(num_buffers); + for len in buffer_lens { + let bytes = read_bytes(value_buffer, &mut offset, len)?; + buffers.push(Buffer::from_vec(bytes.to_vec())); + } + + if offset != value_buffer.len() { + return Err(ArrowError::InvalidArgumentError( + "Invalid scalar value buffer: trailing bytes".to_string(), + )); + } + + let mut builder = ArrayDataBuilder::new(data_type.clone()) + .len(1) + .null_count(0); + for b in buffers { + builder = builder.add_buffer(b); + } + Ok(make_array(builder.build()?)) +} + +pub fn decode_scalar_from_inline_value( + data_type: &DataType, + inline_value: &[u8], +) -> Result { + // I expect our input to be safe here, but I added some debug_assert_eq statements just in case. + // If they are triggered, we may need to change them to return actual errors. + // + // Boolean values are bit-packed in Arrow and therefore are not "fixed-stride" in bytes. + // As a result, `byte_width_opt()` returns `None` for `DataType::Boolean`, even though a + // length-1 scalar can be represented inline using a single byte (matching `try_inline_value`). + if matches!(data_type, DataType::Boolean) { + debug_assert_eq!( + inline_value.len(), + 1, + "Invalid boolean inline scalar length (expected 1 byte, got {})", + inline_value.len() + ); + } else if let Some(byte_width) = data_type.byte_width_opt() { + debug_assert_eq!( + inline_value.len(), + byte_width, + "Inline constant length mismatch for {:?}: expected {} bytes but got {}", + data_type, + byte_width, + inline_value.len() + ); + } + + let data = ArrayDataBuilder::new(data_type.clone()) + .len(1) + .null_count(0) + .add_buffer(Buffer::from_vec(inline_value.to_vec())) + .build()?; + Ok(make_array(data)) +} + +pub fn try_inline_value(scalar: &ArrayRef) -> Option> { + if scalar.null_count() != 0 || scalar.len() != 1 { + return None; + } + let data = scalar.to_data(); + if !data.child_data().is_empty() { + return None; + } + if data.buffers().len() != 1 { + return None; + } + let bytes = data.buffers()[0].as_slice(); + if bytes.len() > INLINE_VALUE_MAX_BYTES { + return None; + } + Some(bytes.to_vec()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{ + BooleanArray, DictionaryArray, FixedSizeBinaryArray, Int8Array, Int32Array, StringArray, + cast::AsArray, types::Int8Type, + }; + + use super::*; + + #[test] + fn test_extract_scalar_value() { + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); + let scalar = extract_scalar_value(&array, 2).unwrap(); + assert_eq!(scalar.len(), 1); + assert_eq!( + scalar + .as_primitive::() + .value(0), + 3 + ); + } + + #[test] + fn test_extract_scalar_value_from_full_dictionary() { + let values = Arc::new(StringArray::from( + (0..=i8::MAX) + .map(|value| format!("value-{value}")) + .collect::>(), + )); + let keys = Int8Array::from((0..=i8::MAX).collect::>()); + let array: ArrayRef = Arc::new(DictionaryArray::::new(keys, values)); + + let scalar = extract_scalar_value(&array, i8::MAX as usize).unwrap(); + + let scalar = scalar.as_dictionary::(); + assert_eq!(scalar.len(), 1); + assert_eq!(scalar.key(0), Some(i8::MAX as usize)); + assert_eq!(scalar.values().len(), i8::MAX as usize + 1); + } + + #[test] + fn test_scalar_value_buffer_utf8_round_trip() { + let scalar: ArrayRef = Arc::new(StringArray::from(vec!["hello"])); + let buf = encode_scalar_value_buffer(&scalar).unwrap(); + let decoded = decode_scalar_from_value_buffer(&DataType::Utf8, &buf).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded.null_count(), 0); + assert_eq!(decoded.as_string::().value(0), "hello"); + } + + #[test] + fn test_scalar_value_buffer_fixed_size_binary_round_trip() { + let val = vec![0xABu8; 33]; + let scalar: ArrayRef = Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + std::iter::once(Some(val.as_slice())), + 33, + ) + .unwrap(), + ); + let buf = encode_scalar_value_buffer(&scalar).unwrap(); + let decoded = + decode_scalar_from_value_buffer(&DataType::FixedSizeBinary(33), &buf).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded.as_fixed_size_binary().value(0), val.as_slice()); + } + + #[test] + fn test_inline_value_boolean_round_trip() { + let scalar: ArrayRef = Arc::new(BooleanArray::from_iter([Some(true)])); + let inline = try_inline_value(&scalar).unwrap(); + let decoded = decode_scalar_from_inline_value(&DataType::Boolean, &inline).unwrap(); + assert_eq!(decoded.len(), 1); + assert_eq!(decoded.null_count(), 0); + assert!(decoded.as_boolean().value(0)); + } + + #[test] + fn test_scalar_value_buffer_rejects_nested_type() { + let field = Arc::new(arrow_schema::Field::new("item", DataType::Int32, false)); + let list: ArrayRef = Arc::new(arrow_array::FixedSizeListArray::new( + field, + 2, + Arc::new(Int32Array::from(vec![1, 2])), + None, + )); + let scalar = list.slice(0, 1); + assert!(encode_scalar_value_buffer(&scalar).is_err()); + } + + #[test] + fn test_decode_scalar_from_value_buffer_rejects_nested_type() { + let buf = Vec::::new(); + let res = + decode_scalar_from_value_buffer(&DataType::Struct(arrow_schema::Fields::empty()), &buf); + assert!(res.is_err()); + } + + #[test] + fn test_decode_scalar_from_value_buffer_trailing_bytes() { + // num_buffers = 0, plus an extra byte + let mut bytes = Vec::new(); + bytes.extend_from_slice(&0u32.to_le_bytes()); + bytes.push(1); + let res = decode_scalar_from_value_buffer(&DataType::Int32, &bytes); + assert!(res.is_err()); + } +} diff --git a/python/vendor/lance-arrow/src/schema.rs b/python/vendor/lance-arrow/src/schema.rs new file mode 100644 index 000000000..0489dd584 --- /dev/null +++ b/python/vendor/lance-arrow/src/schema.rs @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extension to arrow schema + +use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema}; + +use crate::{ARROW_EXT_NAME_KEY, BLOB_META_KEY, BLOB_V2_EXT_NAME}; + +pub enum Indentation { + OneLine, + MultiLine(u8), +} + +impl Indentation { + fn value(&self) -> String { + match self { + Self::OneLine => "".to_string(), + Self::MultiLine(spaces) => " ".repeat(*spaces as usize), + } + } + + fn deepen(&self) -> Self { + match self { + Self::OneLine => Self::OneLine, + Self::MultiLine(spaces) => Self::MultiLine(spaces + 2), + } + } +} + +/// Extends the functionality of [arrow_schema::Field]. +pub trait FieldExt { + /// Create a compact string representation of the field + /// + /// This is intended for display purposes and not for serialization + fn to_compact_string(&self, indent: Indentation) -> String; + + /// Check if the field is marked as a packed struct + fn is_packed_struct(&self) -> bool; + + /// Check if the field is marked as a blob + fn is_blob(&self) -> bool; + + /// Check if the field is marked as a blob + fn is_blob_v2(&self) -> bool; +} + +impl FieldExt for Field { + fn to_compact_string(&self, indent: Indentation) -> String { + let mut result = format!("{}: ", self.name().clone()); + match self.data_type() { + DataType::Struct(fields) => { + result += "{"; + result += &indent.value(); + for (field_idx, field) in fields.iter().enumerate() { + result += field.to_compact_string(indent.deepen()).as_str(); + if field_idx < fields.len() - 1 { + result += ","; + } + result += indent.value().as_str(); + } + result += "}"; + } + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) => { + result += "["; + result += field.to_compact_string(indent.deepen()).as_str(); + result += "]"; + } + DataType::FixedSizeList(child, dimension) => { + result += &format!( + "[{}; {}]", + child.to_compact_string(indent.deepen()), + dimension + ); + } + DataType::Dictionary(key_type, value_type) => { + result += &value_type.to_string(); + result += "@"; + result += &key_type.to_string(); + } + _ => { + result += &self.data_type().to_string(); + } + } + if self.is_nullable() { + result += "?"; + } + result + } + + // Check if field has metadata `packed` set to true, this check is case insensitive. + fn is_packed_struct(&self) -> bool { + let field_metadata = self.metadata(); + const PACKED_KEYS: [&str; 2] = ["packed", "lance-encoding:packed"]; + PACKED_KEYS.iter().any(|key| { + field_metadata + .get(*key) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(false) + }) + } + + fn is_blob(&self) -> bool { + let field_metadata = self.metadata(); + field_metadata.get(BLOB_META_KEY).is_some() + || field_metadata + .get(ARROW_EXT_NAME_KEY) + .map(|value| value == BLOB_V2_EXT_NAME) + .unwrap_or(false) + } + + fn is_blob_v2(&self) -> bool { + let field_metadata = self.metadata(); + field_metadata + .get(ARROW_EXT_NAME_KEY) + .map(|value| value == BLOB_V2_EXT_NAME) + .unwrap_or(false) + } +} + +/// Extends the functionality of [arrow_schema::Schema]. +pub trait SchemaExt { + /// Create a new [`Schema`] with one extra field. + fn try_with_column(&self, field: Field) -> std::result::Result; + + fn try_with_column_at( + &self, + index: usize, + field: Field, + ) -> std::result::Result; + + fn field_names(&self) -> Vec<&String>; + + fn without_column(&self, column_name: &str) -> Schema; + + /// Create a compact string representation of the schema + /// + /// This is intended for display purposes and not for serialization + fn to_compact_string(&self, indent: Indentation) -> String; +} + +impl SchemaExt for Schema { + fn try_with_column(&self, field: Field) -> std::result::Result { + if self.column_with_name(field.name()).is_some() { + return Err(ArrowError::SchemaError(format!( + "Can not append column {} on schema: {:?}", + field.name(), + self + ))); + }; + let mut fields: Vec = self.fields().iter().cloned().collect(); + fields.push(FieldRef::new(field)); + Ok(Self::new_with_metadata(fields, self.metadata.clone())) + } + + fn try_with_column_at( + &self, + index: usize, + field: Field, + ) -> std::result::Result { + if self.column_with_name(field.name()).is_some() { + return Err(ArrowError::SchemaError(format!( + "Failed to modify schema: Inserting column {} would create a duplicate column in schema: {:?}", + field.name(), + self + ))); + }; + let mut fields: Vec = self.fields().iter().cloned().collect(); + fields.insert(index, FieldRef::new(field)); + Ok(Self::new_with_metadata(fields, self.metadata.clone())) + } + + /// Project the schema to remove the given column. + /// + /// This only works on top-level fields right now. If a field does not exist, + /// the schema will be returned as is. + fn without_column(&self, column_name: &str) -> Schema { + let fields: Vec = self + .fields() + .iter() + .filter(|f| f.name() != column_name) + .cloned() + .collect(); + Self::new_with_metadata(fields, self.metadata.clone()) + } + + fn field_names(&self) -> Vec<&String> { + self.fields().iter().map(|f| f.name()).collect() + } + + fn to_compact_string(&self, indent: Indentation) -> String { + let mut result = "{".to_string(); + result += &indent.value(); + for (field_idx, field) in self.fields.iter().enumerate() { + result += field.to_compact_string(indent.deepen()).as_str(); + if field_idx < self.fields.len() - 1 { + result += ","; + } + result += indent.value().as_str(); + } + result += "}"; + result + } +} diff --git a/python/vendor/lance-arrow/src/stream.rs b/python/vendor/lance-arrow/src/stream.rs new file mode 100644 index 000000000..c5e55f093 --- /dev/null +++ b/python/vendor/lance-arrow/src/stream.rs @@ -0,0 +1,554 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Utilities for working with streams of [`RecordBatch`]. + +use arrow_array::RecordBatch; +use arrow_schema::{ArrowError, SchemaRef}; +use futures::stream::{self, Stream, StreamExt}; +use std::pin::Pin; + +use crate::deepcopy::deep_copy_batch_sliced; + +/// Rechunks a stream of [`RecordBatch`] so that each output batch has +/// approximately `target_bytes` of array data. +/// +/// Small input batches are accumulated (by concatenation) until at least +/// `min_bytes` of data has been collected. If the resulting batch exceeds +/// `max_bytes`, it is sliced into roughly equal pieces of ~`max_bytes` +/// (assuming uniform row sizes). +pub fn rechunk_stream_by_size( + input: S, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, +) -> impl Stream> +where + S: Stream>, + E: From, +{ + rechunk_stream_by_size_inner(input, input_schema, min_bytes, max_bytes, false) +} + +/// Like [`rechunk_stream_by_size`] but deep-copies slices so that +/// `get_array_memory_size` reflects the true size of each output batch. +/// +/// After a normal `RecordBatch::slice`, the backing buffers are shared with +/// the original batch, so `get_array_memory_size` still reports the full +/// parent size. This variant deep-copies every slice produced during the +/// splitting phase, which allows the stream to detect and re-split slices +/// that still exceed `max_bytes` (e.g. because a single row is much larger +/// than average). +/// +/// The deep copy is a last resort and potentially expensive for large +/// batches. However, it is only performed when a batch actually needs to be +/// sliced — batches that are already within the target range pass through at +/// zero cost. Use this only when the hard cap on `max_bytes` is a +/// correctness requirement, not merely a performance hint. +pub fn rechunk_stream_by_size_deep_copy( + input: S, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, +) -> impl Stream> +where + S: Stream>, + E: From, +{ + rechunk_stream_by_size_inner(input, input_schema, min_bytes, max_bytes, true) +} + +fn rechunk_stream_by_size_inner( + input: S, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, + deep_copy: bool, +) -> impl Stream> +where + S: Stream>, + E: From, +{ + stream::try_unfold( + RechunkState { + input: Box::pin(input), + accumulated: Vec::new(), + acc_bytes: 0, + done: false, + input_schema, + min_bytes, + max_bytes, + deep_copy, + }, + |mut state| async move { + if state.done && state.accumulated.is_empty() { + return Ok(None); + } + + // Pull batches until we reach the byte target or exhaust input. + // Always pull at least one batch so that min_bytes=0 works. + while !state.done && (state.accumulated.is_empty() || state.acc_bytes < state.min_bytes) + { + match state.input.next().await { + Some(Ok(batch)) => { + state.acc_bytes += batch.get_array_memory_size(); + state.accumulated.push(batch); + } + Some(Err(e)) => return Err(e), + None => { + state.done = true; + } + } + } + + if state.accumulated.is_empty() { + return Ok(None); + } + + // Fast path: if the first accumulated batch already meets the + // byte threshold, deliver it directly instead of concatenating + // everything together (which would just get sliced back apart). + if state.accumulated.len() > 1 + && state.accumulated[0].get_array_memory_size() >= state.min_bytes + { + let b = state.accumulated.remove(0); + state.acc_bytes -= b.get_array_memory_size(); + return Ok(Some((b, state))); + } + + let batch = if state.accumulated.len() == 1 { + state.accumulated.pop().unwrap() + } else { + let b = + arrow_select::concat::concat_batches(&state.input_schema, &state.accumulated) + .map_err(E::from)?; + state.accumulated.clear(); + b + }; + state.acc_bytes = 0; + + // Slice the batch into ~max_bytes pieces assuming uniform row sizes. + let mut slices = + slice_batch(batch, state.max_bytes, state.deep_copy).map_err(E::from)?; + + if slices.len() == 1 { + Ok(Some((slices.pop().unwrap(), state))) + } else { + let first = slices.remove(0); + + // Stash leftover slices for subsequent iterations. + for a in &slices { + state.acc_bytes += a.get_array_memory_size(); + } + state.accumulated = slices; + + Ok(Some((first, state))) + } + }, + ) +} + +/// Slice a batch into pieces of at most `max_bytes`. +/// +/// When `deep_copy` is false, slices share buffers with the original batch +/// and `get_array_memory_size` will still report the parent buffer size. +/// This is fine when the caller only needs approximate sizing. +/// +/// When `deep_copy` is true, each slice is deep-copied so that +/// `get_array_memory_size` reflects the true size. If a deep-copied slice +/// still exceeds `max_bytes` (due to non-uniform row sizes), it is +/// recursively split until every piece is within budget or contains only a +/// single row. +fn slice_batch( + batch: RecordBatch, + max_bytes: usize, + deep_copy: bool, +) -> Result, ArrowError> { + let batch_bytes = batch.get_array_memory_size(); + let num_rows = batch.num_rows(); + + if batch_bytes <= max_bytes { + return Ok(vec![batch]); + } + + if num_rows <= 1 { + // A single row cannot be split further, but the size just measured may + // belong to its source buffer rather than to the row: + // `get_array_memory_size` reports whole buffers, and a slice shares + // them. Returning it uncopied hands the caller that inflated figure -- + // a 100 KB row inside a 200 MB buffer measures as 200 MB, and a caller + // budgeting on it concludes the row cannot fit anywhere. + if deep_copy { + return Ok(vec![deep_copy_batch_sliced(&batch)?]); + } + return Ok(vec![batch]); + } + + let rows_per_chunk = (max_bytes as u64 * num_rows as u64 / batch_bytes as u64).max(1) as usize; + + let mut result = Vec::new(); + let mut offset = 0; + while offset < num_rows { + let len = rows_per_chunk.min(num_rows - offset); + let slice = batch.slice(offset, len); + if deep_copy { + let copied = deep_copy_batch_sliced(&slice)?; + // Recurse: the deep-copied slice has accurate sizes, so if it + // still exceeds max_bytes we can split further. + result.extend(slice_batch(copied, max_bytes, true)?); + } else { + result.push(slice); + } + offset += len; + } + + Ok(result) +} + +/// Internal state for [`rechunk_stream`]. +/// +/// Kept as a named struct so the `try_unfold` closure stays readable. +struct RechunkState { + input: Pin>, + accumulated: Vec, + acc_bytes: usize, + done: bool, + input_schema: SchemaRef, + min_bytes: usize, + max_bytes: usize, + deep_copy: bool, +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow_array::Int32Array; + use arrow_schema::{DataType, Field, Schema}; + use futures::executor::block_on; + + /// A single row that shares a large source buffer must report its own + /// size, not the buffer's. Callers budget on this figure, and an inflated + /// one makes a small row look impossible to place. + #[test] + fn a_single_row_slice_reports_its_own_size() { + use arrow_array::LargeStringArray; + + let row = "x".repeat(100 * 1024); + let values: Vec<&str> = (0..2000).map(|_| row.as_str()).collect(); + let schema = Arc::new(Schema::new(vec![Field::new( + "text", + DataType::LargeUtf8, + false, + )])); + let batch = + RecordBatch::try_new(schema, vec![Arc::new(LargeStringArray::from(values))]).unwrap(); + let whole = batch.get_array_memory_size(); + + let one_row = batch.slice(7, 1); + assert_eq!( + one_row.get_array_memory_size(), + whole, + "a slice shares its source buffer, which is the reason this test exists" + ); + + // Budget far below the row so the split recurses to the single-row floor. + let chunks = slice_batch(one_row, 1024, true).unwrap(); + assert_eq!(chunks.len(), 1); + let measured = chunks[0].get_array_memory_size(); + assert_eq!(chunks[0].num_rows(), 1); + assert!( + measured < whole / 100, + "single row still measured as {measured} bytes against a {whole} byte source" + ); + } + + fn make_batch(num_rows: usize) -> RecordBatch { + let schema = test_schema(); + let values: Vec = (0..num_rows as i32).collect(); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap() + } + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn collect_rechunked( + batches: Vec, + min_bytes: usize, + max_bytes: usize, + ) -> Vec { + let input = stream::iter(batches.into_iter().map(Ok::<_, ArrowError>)); + let rechunked = rechunk_stream_by_size(input, test_schema(), min_bytes, max_bytes); + block_on(rechunked.collect::>()) + .into_iter() + .map(|r| r.unwrap()) + .collect() + } + + fn total_rows(batches: &[RecordBatch]) -> usize { + batches.iter().map(|b| b.num_rows()).sum() + } + + #[test] + fn test_empty_stream() { + let result = collect_rechunked(vec![], 100, 200); + assert!(result.is_empty()); + } + + #[test] + fn test_single_batch_passthrough() { + let batch = make_batch(100); + let bytes = batch.get_array_memory_size(); + // Batch is between min and max — should pass through as-is. + let result = collect_rechunked(vec![batch], bytes / 2, bytes * 2); + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 100); + } + + #[test] + fn test_small_batches_concatenated() { + let one_batch_bytes = make_batch(10).get_array_memory_size(); + let batches: Vec<_> = (0..8).map(|_| make_batch(10)).collect(); + // min = 5 batches worth, max = 10 batches worth. + let result = collect_rechunked(batches, one_batch_bytes * 5, one_batch_bytes * 10); + assert_eq!(total_rows(&result), 80); + // Should have been concatenated into fewer batches than the 8 inputs. + assert!( + result.len() < 8, + "expected fewer output batches, got {}", + result.len() + ); + } + + #[test] + fn test_large_batch_sliced() { + let batch = make_batch(1000); + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], bytes / 8, bytes / 4); + assert_eq!(total_rows(&result), 1000); + assert!( + result.len() >= 4, + "expected at least 4 slices, got {}", + result.len() + ); + } + + #[test] + fn test_sliced_leftovers_are_not_recombined() { + // Key test for the fast-path optimisation. When a large batch is + // sliced, leftover slices should be delivered one-at-a-time without + // being concatenated back together. We verify this by checking that + // every output buffer pointer falls inside the original batch's + // allocation (i.e. they are all zero-copy slices, not fresh copies). + let batch = make_batch(1000); + let bytes = batch.get_array_memory_size(); + let orig_data = batch.column(0).to_data(); + let orig_buf = &orig_data.buffers()[0]; + let orig_start = orig_buf.as_ptr() as usize; + let orig_end = orig_start + orig_buf.len(); + + let result = collect_rechunked(vec![batch], bytes / 8, bytes / 4); + + assert_eq!(total_rows(&result), 1000); + assert!(result.len() >= 4); + + for (i, b) in result.iter().enumerate() { + let ptr = b.column(0).to_data().buffers()[0].as_ptr() as usize; + assert!( + ptr >= orig_start && ptr < orig_end, + "slice {i} buffer at {ptr:#x} is outside the original allocation \ + [{orig_start:#x}, {orig_end:#x}) — it was re-concatenated" + ); + } + } + + #[test] + fn test_flush_remainder_on_stream_end() { + // Data below min_bytes should still be flushed when the stream ends. + let batch = make_batch(10); + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], bytes * 100, bytes * 200); + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 10); + } + + #[test] + fn test_large_then_small_batches() { + // After a large batch is fully drained, subsequent small batches + // should be accumulated normally. + let large = make_batch(1000); + let small_bytes = make_batch(10).get_array_memory_size(); + let batches = vec![ + large, + make_batch(10), + make_batch(10), + make_batch(10), + make_batch(10), + make_batch(10), + ]; + let result = collect_rechunked(batches, small_bytes * 3, small_bytes * 100); + assert_eq!(total_rows(&result), 1050); + // The large batch should appear (possibly sliced) followed by + // concatenated small batches, so we should have fewer output batches + // than the 6 inputs. + assert!(result.len() < 6); + } + + #[test] + fn test_row_preservation_across_slicing() { + // Verify that every input row appears exactly once in the output + // and in the correct order after slicing. + let batch = make_batch(237); // odd count to exercise remainder slice + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], bytes / 8, bytes / 5); + + assert_eq!(total_rows(&result), 237); + + let values: Vec = result + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + let expected: Vec = (0..237).collect(); + assert_eq!(values, expected); + } + + #[test] + fn test_min_bytes_zero_still_yields_all_rows() { + // When min_bytes=0, the stream should still yield every batch. + // This is the "chop only, don't coalesce" use case. + let batches: Vec<_> = (0..5).map(|_| make_batch(100)).collect(); + let batch_bytes = batches[0].get_array_memory_size(); + let result = collect_rechunked(batches, 0, batch_bytes * 2); + assert_eq!(total_rows(&result), 500); + } + + #[test] + fn test_min_bytes_zero_slices_oversized() { + // min_bytes=0 with a small max_bytes should still slice large batches. + let batch = make_batch(1000); + let bytes = batch.get_array_memory_size(); + let result = collect_rechunked(vec![batch], 0, bytes / 4); + assert_eq!(total_rows(&result), 1000); + assert!( + result.len() >= 4, + "expected at least 4 slices, got {}", + result.len() + ); + } + + /// Build a batch with one variable-length string column. + /// Every row is `small_size` bytes except the row at index `big_row_idx` + /// which is `big_size` bytes. + fn make_variable_batch( + num_rows: usize, + small_size: usize, + big_row_idx: usize, + big_size: usize, + ) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let values: Vec = (0..num_rows) + .map(|i| { + if i == big_row_idx { + "X".repeat(big_size) + } else { + "x".repeat(small_size) + } + }) + .collect(); + let array = arrow_array::StringArray::from(values); + RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap() + } + + fn variable_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])) + } + + fn collect_rechunked_variable( + batches: Vec, + min_bytes: usize, + max_bytes: usize, + ) -> Vec { + let input = stream::iter(batches.into_iter().map(Ok::<_, ArrowError>)); + let rechunked = + rechunk_stream_by_size_deep_copy(input, variable_schema(), min_bytes, max_bytes); + block_on(rechunked.collect::>()) + .into_iter() + .map(|r| r.unwrap()) + .collect() + } + + #[test] + fn test_oversized_row_at_end() { + // 100 rows: 99 small (64 bytes each) + 1 large (100KiB) at the end. + let batch = make_variable_batch(100, 64, 99, 100 * 1024); + let max_bytes = 64 * 1024; + let result = collect_rechunked_variable(vec![batch], 0, max_bytes); + assert_eq!(total_rows(&result), 100); + for (i, b) in result.iter().enumerate() { + let size = b.get_array_memory_size(); + assert!( + size <= max_bytes || b.num_rows() == 1, + "batch {i} has {size} bytes (max {max_bytes}) and {} rows", + b.num_rows() + ); + } + } + + #[test] + fn test_oversized_row_at_start() { + // 100 rows: 1 large (100KiB) at the start + 99 small (64 bytes each). + let batch = make_variable_batch(100, 64, 0, 100 * 1024); + let max_bytes = 64 * 1024; + let result = collect_rechunked_variable(vec![batch], 0, max_bytes); + assert_eq!(total_rows(&result), 100); + for (i, b) in result.iter().enumerate() { + let size = b.get_array_memory_size(); + assert!( + size <= max_bytes || b.num_rows() == 1, + "batch {i} has {size} bytes (max {max_bytes}) and {} rows", + b.num_rows() + ); + } + } + + #[test] + fn test_oversized_row_in_middle() { + // 100 rows: 1 large (100KiB) in the middle + 99 small (64 bytes each). + let batch = make_variable_batch(100, 64, 50, 100 * 1024); + let max_bytes = 64 * 1024; + let result = collect_rechunked_variable(vec![batch], 0, max_bytes); + assert_eq!(total_rows(&result), 100); + for (i, b) in result.iter().enumerate() { + let size = b.get_array_memory_size(); + assert!( + size <= max_bytes || b.num_rows() == 1, + "batch {i} has {size} bytes (max {max_bytes}) and {} rows", + b.num_rows() + ); + } + } + + #[test] + fn test_error_propagation() { + let input = stream::iter(vec![ + Ok(make_batch(10)), + Err(ArrowError::ComputeError("boom".into())), + Ok(make_batch(10)), + ]); + let rechunked = rechunk_stream_by_size(input, test_schema(), 1, usize::MAX); + let results: Vec> = block_on(rechunked.collect()); + assert!(results.iter().any(|r| r.is_err())); + } +} diff --git a/python/vendor/lance-arrow/src/struct.rs b/python/vendor/lance-arrow/src/struct.rs new file mode 100644 index 000000000..ed4b6288e --- /dev/null +++ b/python/vendor/lance-arrow/src/struct.rs @@ -0,0 +1,227 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Extension to arrow struct arrays + +use arrow_array::{Array, StructArray, cast::AsArray, make_array}; +use arrow_buffer::NullBuffer; +use arrow_data::{ArrayData, ArrayDataBuilder}; +use arrow_schema::{ArrowError, DataType}; + +pub trait StructArrayExt { + /// Removes the offset / length of the struct array by pushing it into the children + /// + /// In arrow-rs when slice is called it recursively slices the children. + /// In arrow-cpp when slice is called it just sets the offset/length of + /// the struct array and leaves the children as-is + /// + /// Both are legal approaches (╥﹏╥) + /// + /// This method helps reduce complexity by folding into the arrow-rs approach + fn normalize_slicing(&self) -> Result + where + Self: Sized; + + /// Structs are allowed to mask valid items. For example, a struct array might be: + /// + /// [ {"items": [1, 2, 3]}, NULL, {"items": NULL}] + /// + /// However, the underlying items array might be: [[1, 2, 3], [4, 5], NULL] + /// + /// The [4, 5] list is masked out because the struct array is null. + /// + /// The struct validity would be [true, false, true] and the list validity would be [true, true, false] + /// + /// This method pushes nulls down into all children. In the above example the list validity would become + /// [true, false, false]. + /// + /// This method is not recursive. If a child is a struct array it will not push that child's nulls down. + /// + /// This method does not remove garbage lists. It only updates the validity so a future call to + /// [crate::list::ListArrayExt::filter_garbage_nulls] will remove the garbage lists (without + /// this pushdown it would not) + fn pushdown_nulls(&self) -> Result + where + Self: Sized; +} + +fn normalized_struct_array_data(data: ArrayData) -> Result { + let parent_offset = data.offset(); + let parent_len = data.len(); + let modified_children = data + .child_data() + .iter() + .map(|d| { + let d = normalized_struct_array_data(d.clone())?; + let offset = d.offset(); + let len = d.len(); + if len < parent_len + parent_offset { + return Err(ArrowError::InvalidArgumentError(format!( + "Child array {} has length {} which is less than the parent length {} plus the parent offset {}", + d.data_type(), + len, + parent_len, + parent_offset + ))); + } + let new_offset = offset + parent_offset; + d.into_builder().offset(new_offset) + .len(parent_len) + .build() + }) + .collect::, _>>()?; + ArrayDataBuilder::new(data.data_type().clone()) + .len(parent_len) + .offset(0) + .buffers(data.buffers().to_vec()) + .child_data(modified_children) + .build() +} + +impl StructArrayExt for StructArray { + fn normalize_slicing(&self) -> Result + where + Self: Sized, + { + if self.offset() == 0 && self.columns().iter().all(|c| c.len() == self.len()) { + return Ok(self.clone()); + } + + let data = normalized_struct_array_data(self.to_data())?; + Ok(Self::from(data)) + } + + fn pushdown_nulls(&self) -> Result + where + Self: Sized, + { + let Some(validity) = self.nulls() else { + return Ok(self.clone()); + }; + let data = self.to_data(); + let children = data + .child_data() + .iter() + .map(|c| { + if c.data_type() == &DataType::Null { + // A Null array is already null in every slot, and arrow-rs rejects + // a validity buffer on it. + Ok(c.clone()) + } else if let Some(child_validity) = c.nulls() { + let new_validity = child_validity.inner() & validity.inner(); + c.clone() + .into_builder() + .nulls(Some(NullBuffer::from(new_validity))) + .build() + } else { + Ok(c.clone() + .into_builder() + .nulls(Some(validity.clone())) + .build()?) + } + }) + .collect::, _>>()?; + let arr = make_array(data.into_builder().child_data(children).build()?); + Ok(arr.as_struct().clone()) + } +} + +#[cfg(test)] +mod tests { + use arrow_array::{ + Array, Int32Array, NullArray, StructArray, cast::AsArray, make_array, types::Int32Type, + }; + use arrow_buffer::NullBuffer; + use arrow_schema::{DataType, Field, Fields}; + use std::sync::Arc; + + use crate::r#struct::StructArrayExt; + + #[test] + fn test_normalize_slicing_no_offset() { + let x = Int32Array::from(vec![1, 2, 3]); + let y = Int32Array::from(vec![4, 5, 6]); + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]), + vec![Arc::new(x), Arc::new(y)], + None, + ); + + let normalized = struct_array.normalize_slicing().unwrap(); + assert_eq!(normalized, struct_array); + } + + #[test] + fn test_arrow_rs_slicing() { + let x = Int32Array::from(vec![1, 2, 3, 4]); + let y = Int32Array::from(vec![5, 6, 7, 8]); + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]), + vec![Arc::new(x), Arc::new(y)], + None, + ); + + // Slicing with arrow-rs propagates the slicing to the children so there should + // be no change needed to the struct array + let sliced = struct_array.slice(1, 2); + let normalized = sliced.normalize_slicing().unwrap(); + + assert_eq!(normalized, sliced); + } + + #[test] + fn test_arrow_cpp_slicing() { + let x = Int32Array::from(vec![1, 2, 3, 4]); + let y = Int32Array::from(vec![5, 6, 7, 8]); + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]), + vec![Arc::new(x), Arc::new(y)], + None, + ); + + let data = struct_array.to_data(); + let sliced = data.into_builder().offset(1).len(2).build().unwrap(); + let sliced = make_array(sliced); + let normalized = sliced.as_struct().clone().normalize_slicing().unwrap(); + + assert_eq!(normalized, struct_array.slice(1, 2)); + } + + #[test] + fn test_pushdown_nulls_with_null_typed_child() { + // `pa.Table.from_pylist([{"s": {"a": 1, "b": None}}, {"s": None}])` infers a + // struct with a Null-typed child. Arrow forbids a validity buffer on a Null + // array, so the struct's nulls must not be pushed into it. + let struct_array = StructArray::new( + Fields::from(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Null, true), + ]), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(NullArray::new(3)), + ], + Some(NullBuffer::from(vec![true, false, true])), + ); + + let pushed = struct_array.pushdown_nulls().unwrap(); + + assert_eq!(pushed.nulls(), struct_array.nulls()); + assert_eq!( + pushed.column(0).as_primitive::(), + &Int32Array::from(vec![Some(1), None, Some(3)]) + ); + assert_eq!(pushed.column(1).data_type(), &DataType::Null); + assert_eq!(pushed.column(1).logical_null_count(), 3); + } +} diff --git a/rust/lancedb/tests/nullable_struct_merge.rs b/rust/lancedb/tests/nullable_struct_merge.rs new file mode 100644 index 000000000..c245b1c2c --- /dev/null +++ b/rust/lancedb/tests/nullable_struct_merge.rs @@ -0,0 +1,64 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef, BooleanArray, Int32Array, RecordBatch, StructArray}; +use arrow_buffer::NullBuffer; +use arrow_schema::{DataType, Field, Schema}; +use lance_arrow::RecordBatchExt; + +#[test] +fn merge_nullable_struct_with_sliced_boolean_child() { + let left_struct = StructArray::new( + vec![Field::new("id", DataType::Int32, false)].into(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef], + None, + ); + + // BooleanArray retains its offset when sliced. The merged child bitmap must + // cover that offset as well as the four visible rows. + let boolean_child: ArrayRef = Arc::new(BooleanArray::from_iter((0..1030).map(|index| { + if index == 1027 { + None + } else { + Some(index % 2 == 1) + } + }))); + let boolean_child = boolean_child.slice(1025, 4); + assert_eq!(boolean_child.offset(), 1025); + let right_struct = StructArray::new( + vec![Field::new("flag", DataType::Boolean, true)].into(), + vec![boolean_child], + Some(NullBuffer::from(vec![true, false, true, true])), + ); + + let left_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "nested", + left_struct.data_type().clone(), + true, + )])), + vec![Arc::new(left_struct)], + ) + .unwrap(); + let right_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "nested", + right_struct.data_type().clone(), + true, + )])), + vec![Arc::new(right_struct)], + ) + .unwrap(); + + let merged = left_batch.merge(&right_batch).unwrap(); + let nested = merged.column_by_name("nested").unwrap(); + let nested = nested.as_any().downcast_ref::().unwrap(); + let flag = nested.column_by_name("flag").unwrap(); + let flag = flag.as_any().downcast_ref::().unwrap(); + assert_eq!( + flag.iter().collect::>(), + vec![Some(true), None, None, Some(false)] + ); +} From 55897d7c3b03c6b94b461161abc1b28d34903fde Mon Sep 17 00:00:00 2001 From: Yingzi Zhang <164733623+LindseyZ1205@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:26:43 -0400 Subject: [PATCH 17/68] fix(python): preserve Arrow schema when merging reranker scores (#4332) Fixes #4331 With `return_score="all"`, model rerankers reconstruct merged rows without their Arrow schema. This changes fixed-size float32 vectors into variable-length float64 lists, loses metadata, and can overflow valid uint64 row IDs. Unify the original input schemas before inserting placeholder scores, then use the unified schema when reconstructing the merged table. Row merging and reranking behavior stay the same. Capturing the schema first also preserves float64 scores without conflicting with the float32 null placeholders. Validation: - New regression cases: 5 failures and 5 passing controls before the fix; all 10 pass afterward. - Reranker, hybrid-query and TypeSafe reranker tests: 150 passed, 17 skipped. Optional live model/API tests were not run. - The fake-client tests that use random metadata passed again with seeds 7 and 41. - Repository-wide Ruff lint/format checks, configured Pyright check and applicable pre-commit hooks pass. - Tests used the locally built Rust/Python extension and the repository's locked uv environment. No Rust code changed. AI disclosure: I used Codex for investigation, implementation, tests, local checks, and drafting this description. --- python/python/lancedb/rerankers/base.py | 5 +- python/python/tests/test_rerankers.py | 108 ++++++++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/rerankers/base.py b/python/python/lancedb/rerankers/base.py index 4af938a1e..e1bc7f0dd 100644 --- a/python/python/lancedb/rerankers/base.py +++ b/python/python/lancedb/rerankers/base.py @@ -241,6 +241,7 @@ class Reranker(ABC): This op is slower than just keeping relevance score but can be useful for debugging. """ + combined_schema = pa.unify_schemas([vector_results.schema, fts_results.schema]) # add nulls to fts results for _distance if "_distance" not in fts_results.column_names: fts_results = fts_results.append_column( @@ -265,7 +266,9 @@ class Reranker(ABC): else: vector_results_dict[key] = value - combined = pa.Table.from_pylist(list(vector_results_dict.values())) + combined = pa.Table.from_pylist( + list(vector_results_dict.values()), schema=combined_schema + ) return combined def _keep_relevance_score(self, combined_results: pa.Table): diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 850b477a0..d5ee6b6a4 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -785,6 +785,114 @@ def test_empty_hybrid_result_reranker(): assert "_rowid" in result.column_names +@pytest.mark.parametrize("return_score", ["relevance", "all"]) +@pytest.mark.parametrize("first_row_id", [1, 2**63]) +@pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()]) +def test_model_reranker_preserves_arrow_schema( + mocker, return_score, first_row_id, score_type +): + common_schema = pa.schema( + [ + pa.field("_rowid", pa.uint64()), + pa.field("text", pa.string()), + pa.field("vector", pa.list_(pa.float32(), 2)), + ], + metadata={b"source": b"reranker-test"}, + ) + vector_results = pa.Table.from_pylist( + [ + { + "_rowid": first_row_id, + "text": "vector", + "vector": [1, 0], + "_distance": 0.1, + }, + { + "_rowid": first_row_id + 1, + "text": "both", + "vector": [0, 1], + "_distance": 0.2, + }, + ], + schema=common_schema.append(pa.field("_distance", score_type)), + ) + fts_results = pa.Table.from_pylist( + [ + { + "_rowid": first_row_id + 1, + "text": "both", + "vector": [0, 1], + "_score": 0.9, + }, + { + "_rowid": first_row_id + 2, + "text": "fts", + "vector": [1, 1], + "_score": 0.8, + }, + ], + schema=common_schema.append(pa.field("_score", score_type)), + ) + reranker = CohereReranker(return_score=return_score) + reranker._client = mocker.Mock() + reranker._client.rerank.return_value.results = [ + mocker.Mock(index=index, relevance_score=score) + for index, score in [(2, 0.9), (1, 0.8), (0, 0.7)] + ] + + result = reranker.rerank_hybrid("query", vector_results, fts_results) + + assert result.select(common_schema.names).schema.equals( + common_schema, check_metadata=True + ) + assert result["_rowid"].to_pylist() == [ + first_row_id + 2, + first_row_id + 1, + first_row_id, + ] + assert result["text"].to_pylist() == ["fts", "both", "vector"] + if return_score == "all": + assert ( + result["_distance"] + .combine_chunks() + .equals(pa.array([None, 0.2, 0.1], type=score_type)) + ) + assert ( + result["_score"] + .combine_chunks() + .equals(pa.array([0.8, 0.9, None], type=score_type)) + ) + + +@pytest.mark.parametrize("return_score", ["relevance", "all"]) +def test_model_reranker_empty_results_preserve_schema(mocker, return_score): + common_schema = pa.schema( + [ + ("_rowid", pa.uint64()), + ("text", pa.string()), + ("vector", pa.list_(pa.float32(), 2)), + ], + metadata={b"source": b"reranker-test"}, + ) + vector_results = pa.Table.from_pylist( + [], schema=common_schema.append(pa.field("_distance", pa.float32())) + ) + fts_results = pa.Table.from_pylist( + [], schema=common_schema.append(pa.field("_score", pa.float32())) + ) + reranker = CohereReranker(return_score=return_score) + reranker._client = mocker.Mock() + + result = reranker.rerank_hybrid("query", vector_results, fts_results) + + assert len(result) == 0 + assert result.select(common_schema.names).schema.equals( + common_schema, check_metadata=True + ) + assert result.schema.field("_relevance_score").type == pa.float32() + reranker._client.rerank.assert_not_called() + + def test_cross_encoder_reranker_return_all(tmp_path): pytest.importorskip("sentence_transformers") reranker = CrossEncoderReranker(return_score="all") From 7f932b8552c22355902829be1c5615a867d1bd45 Mon Sep 17 00:00:00 2001 From: Prince Jain Date: Thu, 22 Oct 2026 23:15:26 +0530 Subject: [PATCH 18/68] fix(nodejs): reject listNamespaces limit above i32::MAX instead of wrapping (#4584) ## Summary Fixes #4583. `Connection::list_namespaces` narrowed `Option` to `i32` with an unchecked `as` cast, so a JavaScript `limit` between 2^31 and 2^32 - 1 (e.g. `3000000000`) wrapped negative before reaching `ListNamespacesRequest`. The server received a negative page size instead of a clear client-side error. Replaced the cast with a `checked_limit` helper using `i32::try_from`, returning `napi::Error::from_reason` on overflow, matching the existing `try_from` validation used for `limit` elsewhere in the same file. ## Tests Unit tests cover the wrap case (`u32` values above `i32::MAX` reject), the boundary (`i32::MAX` passes), and `None`/zero passthrough. `cargo test -p lancedb-nodejs` passes 8/8 and `pnpm build` compiles the binding and TypeScript definitions cleanly. --- nodejs/src/connection.rs | 32 +++++++++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 58ec9859a..67c4706e4 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -673,7 +673,7 @@ impl Connection { let req = ListNamespacesRequest { id: namespace_path, page_token, - limit: limit.map(|l| l as i32), + limit: checked_limit(limit)?, ..Default::default() }; let resp = self @@ -787,3 +787,33 @@ impl Connection { .default_error() } } + +fn checked_limit(limit: Option) -> napi::Result> { + limit + .map(i32::try_from) + .transpose() + .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) +} + +#[cfg(test)] +mod tests { + use super::checked_limit; + + #[test] + fn checked_limit_rejects_values_above_i32_max() { + // u32 values above i32::MAX wrapped negative through `as i32`, + // sending a negative page size to the server. + assert!(checked_limit(Some(3_000_000_000)).is_err()); + assert!(checked_limit(Some(u32::MAX)).is_err()); + } + + #[test] + fn checked_limit_passes_boundary_and_missing_values() { + assert_eq!( + checked_limit(Some(i32::MAX as u32)).unwrap(), + Some(i32::MAX) + ); + assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); + assert_eq!(checked_limit(None).unwrap(), None); + } +} From 62d83f2ea2b43fcf1d37869bbaf06883fe768882 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 02:19:32 +0800 Subject: [PATCH 19/68] fix(python): expect oidc credential-type header on refreshed tokens (#4632) ## Summary #3971 added `test_zero_lifetime_token_is_refreshed`, but its expected headers omit `x-lancedb-credential-type: oidc`, which `OAuthProvider.get_headers()` has emitted since #2631/#4070. The test fails against current main and breaks the PR CI matrix for every open PR. ## Changes - Align the two zap assertions in `test_zero_lifetime_token_is_refreshed` with the provider contract (see `test_token_refresh` and `test_sync_method`). ## Tests - `pytest python/tests/test_header_provider.py -q` --- python/python/tests/test_header_provider.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/python/tests/test_header_provider.py b/python/python/tests/test_header_provider.py index e50ff3270..9d19e7bef 100644 --- a/python/python/tests/test_header_provider.py +++ b/python/python/tests/test_header_provider.py @@ -128,8 +128,14 @@ class TestOAuthProvider: provider = OAuthProvider(fetcher, refresh_buffer_seconds=0) - assert provider.get_headers() == {"Authorization": "Bearer token1"} - assert provider.get_headers() == {"Authorization": "Bearer token2"} + assert provider.get_headers() == { + "Authorization": "Bearer token1", + "x-lancedb-credential-type": "oidc", + } + assert provider.get_headers() == { + "Authorization": "Bearer token2", + "x-lancedb-credential-type": "oidc", + } assert call_count == 2 def test_missing_access_token(self): From a0946282ca57a58e18b8cea9a50be39195b3323a Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:06 +0100 Subject: [PATCH 20/68] fix(python): return original scores from hybrid search with normalize="rank" (#4343) ## Summary With `rerank(..., normalize="rank")`, a hybrid query returns ranks (1.0, 2.0, 3.0, ...) in the `_distance` and `_score` columns instead of the real vector distances and BM25 scores. `_combine_hybrid_results` converted both result sets to ranks first, and only then captured `original_distances` / `original_scores`. So the step that restores the original values after reranking put the ranks back. With `normalize="score"`, the real values are returned as expected. Repro (RRFReranker(return_score="all"), vector distances 0 / 25 / 100): | normalize | `_distance` returned | `_score` returned | | --- | --- | --- | | `"score"` | 0.0, 25.0, 100.0 | 0.47, 0.67, None | | `"rank"` (before) | 1.0, 2.0, 3.0 | 1.0, 2.0, None | | `"rank"` (after) | 0.0, 25.0, 100.0 | 0.47, 0.67, None | ## Changes - Capture the original `_distance` / `_score` columns before converting them to ranks, then normalize the ranked column. The reranker input is unchanged; only the values restored afterwards are fixed. `AsyncHybridQuery` uses the same function, so both sync and async are covered. ## Test plan - [x] New `test_hybrid_return_all_restores_original_scores` in `test_rerankers.py`, parametrized over `score` / `rank`: the `rank` case fails before the fix and passes after - [x] `pytest python/tests/test_rerankers.py python/tests/test_hybrid_query.py` - [x] `ruff format --check` / `ruff check` I have no local Rust toolchain, so I ran the tests against the released 0.39.0 native module with the changed Python files overlaid. The changed functions are identical between 0.39.0 and main. Co-authored-by: Xuanwo --- python/python/lancedb/query.py | 20 +++++++++----- python/python/tests/test_rerankers.py | 40 +++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 7 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index fe8c33861..9d7610f72 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2575,10 +2575,6 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): with_row_ids: bool, offset: Optional[int] = None, ) -> pa.Table: - if norm == "rank": - vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance") - fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score") - # If both result sets are empty (e.g. after hard filtering), # return early to avoid errors in reranking or score restoration. if vector_results.num_rows == 0 and fts_results.num_rows == 0: @@ -2605,15 +2601,23 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): original_score_row_ids = None # normalize the scores to be between 0 and 1, 0 being most relevant # We check whether the results (vector and FTS) are empty, because when - # they are, they often are missing the _rowid column, which causes an error + # they are, they often are missing the _rowid column, which causes an error. + # The original values are captured before any rank conversion so that + # they can be restored after reranking. if vector_results.num_rows > 0: distance_i = vector_results.column_names.index("_distance") original_distances = vector_results.column(distance_i) original_distance_row_ids = vector_results.column("_rowid") + if norm == "rank": + vector_results = LanceHybridQueryBuilder._rank( + vector_results, "_distance" + ) vector_results = vector_results.set_column( distance_i, vector_results.field(distance_i), - LanceHybridQueryBuilder._normalize_scores(original_distances), + LanceHybridQueryBuilder._normalize_scores( + vector_results.column(distance_i) + ), ) # In fts higher scores represent relevance. Not inverting them here as @@ -2622,10 +2626,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): score_i = fts_results.column_names.index("_score") original_scores = fts_results.column(score_i) original_score_row_ids = fts_results.column("_rowid") + if norm == "rank": + fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score") fts_results = fts_results.set_column( score_i, fts_results.field(score_i), - LanceHybridQueryBuilder._normalize_scores(original_scores), + LanceHybridQueryBuilder._normalize_scores(fts_results.column(score_i)), ) results = reranker.rerank_hybrid(fts_query, vector_results, fts_results) diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index d5ee6b6a4..743669367 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -785,6 +785,46 @@ def test_empty_hybrid_result_reranker(): assert "_rowid" in result.column_names +@pytest.mark.parametrize("norm", ["score", "rank"]) +def test_hybrid_return_all_restores_original_scores(norm): + from lancedb.query import LanceHybridQueryBuilder + + vector_results = pa.table( + { + "text": ["a", "b", "c"], + "_rowid": pa.array([0, 1, 2], type=pa.uint64()), + "_distance": pa.array([0.5, 2.5, 10.0], type=pa.float32()), + } + ) + fts_results = pa.table( + { + "text": ["b", "a"], + "_rowid": pa.array([1, 0], type=pa.uint64()), + "_score": pa.array([3.25, 1.75], type=pa.float32()), + } + ) + + result = LanceHybridQueryBuilder._combine_hybrid_results( + fts_results=fts_results, + vector_results=vector_results, + norm=norm, + fts_query="query", + reranker=RRFReranker(return_score="all"), + limit=10, + with_row_ids=True, + ) + + # The raw distances and scores are returned, not the normalized values + # or ranks used for reranking. + rows = {row["_rowid"]: row for row in result.to_pylist()} + assert rows[0]["_distance"] == 0.5 + assert rows[1]["_distance"] == 2.5 + assert rows[2]["_distance"] == 10.0 + assert rows[0]["_score"] == 1.75 + assert rows[1]["_score"] == 3.25 + assert rows[2]["_score"] is None + + @pytest.mark.parametrize("return_score", ["relevance", "all"]) @pytest.mark.parametrize("first_row_id", [1, 2**63]) @pytest.mark.parametrize("score_type", [pa.float32(), pa.float64()]) From 5f9a8af27c66d8730b9e7cf3c29cac1f2bd237cf Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:21 +0100 Subject: [PATCH 21/68] fix(python): keep embedding functions registered without an alias readable (#4356) `EmbeddingFunctionRegistry.register()` documents that the class name is used when no alias is given, and it does register the class under that name. But it stored the alias itself (`None`) as `__embedding_function_registry_alias__`, which is what `function_to_metadata()` writes into the table schema. A table using such a function gets `"name": null` in its `embedding_functions` metadata, and any later read of the functions fails: ```python @get_registry().register() class MyEmbed(TextEmbeddingFunction): ... f = MyEmbed.create() class Schema(LanceModel): text: str = f.SourceField() vector: Vector(4) = f.VectorField() tbl = db.create_table("t", schema=Schema) tbl.add([{"text": "hi"}]) # KeyError: None tbl.search("hi") # KeyError: None ``` This stores the key the class is registered under instead, so the class name round-trips. For aliased classes the key is the alias, so their metadata is unchanged. The TypeScript registry already falls back to `ctor.name` in the same way. Added `test_embedding_function_registered_without_alias`, which creates a table with such a function, reopens it, and adds a row. It fails with `KeyError` before the change. `test_embeddings.py`, `test_pydantic.py`, `test_table.py` and `test_db.py` pass, and `ruff format --check` / `ruff check` are clean. --- python/python/lancedb/embeddings/registry.py | 4 ++- python/python/tests/test_embeddings.py | 31 ++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/embeddings/registry.py b/python/python/lancedb/embeddings/registry.py index 1dc7ffcb8..3b940fe2e 100644 --- a/python/python/lancedb/embeddings/registry.py +++ b/python/python/lancedb/embeddings/registry.py @@ -73,7 +73,9 @@ class EmbeddingFunctionRegistry: raise KeyError(f"{cls.__name__} was already registered") key = alias or cls.__name__ self._functions[key] = cls - cls.__embedding_function_registry_alias__ = alias + # Store the name the class is registered under, so the table + # metadata written for it can be looked up again when reading. + cls.__embedding_function_registry_alias__ = key return cls return decorator diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index c42172f77..51907f141 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -64,6 +64,37 @@ def test_embedding_function(tmp_path): assert np.allclose(actual, expected) +def test_embedding_function_registered_without_alias(tmp_path): + registry = EmbeddingFunctionRegistry.get_instance() + + @registry.register() + class NoAliasEmbeddingFunction(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [np.array([1.0, 2.0]) for _ in texts] + + assert registry.get("NoAliasEmbeddingFunction") is NoAliasEmbeddingFunction + + func = NoAliasEmbeddingFunction.create() + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(2) = func.VectorField() + + db = lancedb.connect(tmp_path) + db.create_table("test", schema=Schema) + + # The class name is written to the table metadata, so the function can be + # resolved again when the table is reopened. + table = lancedb.connect(tmp_path).open_table("test") + table.add([{"text": "hello"}]) + conf = table.embedding_functions["vector"] + assert isinstance(conf.function, NoAliasEmbeddingFunction) + assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0]] + + def test_instructor_ndims_uses_instruction(): instructor = get_registry().get("instructor").create() model = MagicMock() From 712a8b86995397de03ae1c0389f68824505f80a5 Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi <51020382+MohammadHijjawi97@users.noreply.github.com> Date: Thu, 22 Oct 2026 19:42:36 +0100 Subject: [PATCH 22/68] fix(python): write plain Enum fields of a LanceModel as their values (#4357) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since #3232, `pydantic_to_schema` maps an `Enum` field to the Arrow type of its values, including enums without a `str`/`int` mixin. Writing such models still failed: `model_to_dict()` returns `model.model_dump()`, which keeps the Enum members, and Arrow can only convert members that are also `str`/`int` instances. ```python class Color(Enum): RED = "red" class Item(LanceModel): color: Color tbl = db.create_table("items", schema=Item) tbl.add([Item(color=Color.RED)]) # pyarrow.lib.ArrowTypeError: Expected bytes, got a 'Color' object ``` `create_table(data=[...])` fails the same way, and int-valued enums raise `ArrowInvalid`. `model_to_dict()` now replaces Enum members with their values, including in nested models and lists. Lists of plain scalars (vectors in particular) are returned without being walked, so the per-row cost stays about the same as `model_dump()` alone (21.3µs vs 20.3µs for a 1536-dim model). Reading back is unchanged: Pydantic already validates the stored value into the Enum member. Added `test_plain_enum_values_are_written`, which covers str and int valued enums, `Optional` and `List` fields, and an Enum inside a nested model, and checks that `to_pydantic` round-trips the rows. It fails with the `ArrowTypeError` above before the change. `test_pydantic.py`, `test_embeddings.py`, `test_table.py` and `test_db.py` pass, and `ruff format --check` / `ruff check` are clean. --- python/python/lancedb/pydantic.py | 18 +++++++++- python/python/tests/test_pydantic.py | 49 ++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 313c514c3..8fe3c1c93 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -524,4 +524,20 @@ def get_extras(field_info: FieldInfo, key: str) -> Any: def model_to_dict(model: pydantic.BaseModel) -> dict[str, Any]: """Convert a Pydantic model to a dictionary.""" - return model.model_dump() + return _enum_members_to_values(model.model_dump()) + + +def _enum_members_to_values(value: Any) -> Any: + # model_dump() keeps Enum members, and Arrow cannot convert a member of a + # plain Enum (one without a str or int mixin). Write its value instead, + # which is what the Arrow type of the field is derived from. + if isinstance(value, Enum): + return value.value + if isinstance(value, dict): + return {k: _enum_members_to_values(v) for k, v in value.items()} + if isinstance(value, (list, tuple)) and value: + # Skip lists of plain scalars, such as vectors, without walking them. + if isinstance(value[0], (int, float, str, bytes)): + return value + return [_enum_members_to_values(v) for v in value] + return value diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 81674d735..24fb7ee89 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -733,3 +733,52 @@ def test_enum_types(): assert schema.field("priority").type == pa.int64() assert schema.field("opt_status").type == pa.dictionary(pa.int32(), pa.utf8()) assert schema.field("opt_status").nullable + + +def test_plain_enum_values_are_written(mem_db): + """Members of an Enum without a str/int mixin are written as their values.""" + + class Color(Enum): + RED = "red" + BLUE = "blue" + + class Priority(Enum): + LOW = 1 + HIGH = 2 + + class Tag(LanceModel): + color: Color + + class TestModel(LanceModel): + color: Color + priority: Priority + opt_color: Optional[Color] = None + colors: List[Color] + tag: Tag + vector: Vector(2) + + rows = [ + TestModel( + color=Color.RED, + priority=Priority.HIGH, + opt_color=Color.BLUE, + colors=[Color.BLUE, Color.RED], + tag=Tag(color=Color.BLUE), + vector=[1.0, 2.0], + ), + TestModel( + color=Color.BLUE, + priority=Priority.LOW, + colors=[], + tag=Tag(color=Color.RED), + vector=[3.0, 4.0], + ), + ] + tbl = mem_db.create_table("test", schema=TestModel) + tbl.add(rows) + + assert tbl.to_arrow().select(["color", "priority", "opt_color"]).to_pylist() == [ + {"color": "red", "priority": 2, "opt_color": "blue"}, + {"color": "blue", "priority": 1, "opt_color": None}, + ] + assert tbl.search().to_pydantic(TestModel) == rows From 4bdc8c3c0015c54587445409ad10619481843273 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:42:51 +0800 Subject: [PATCH 23/68] fix: wait for the submitted remote index job (#4428) Addresses the client-side waiting portion of #4402. When create_index returns a job ID, wait_timeout now waits for that submission rather than an already-ready index with the same name. It propagates JobFailed / JobCancelled and bounds job polling by the requested timeout. Older servers without job IDs keep name-based waiting, calls without wait_timeout remain fire-and-forget, and the existing two-hour maximum is preserved. The job wrapper is shared with create_index_async; no synchronous server-side config validation is added. Rebased on main 81d52a98, preserving the new materialized-view SQL constructor. Validation on Linux / Rust 1.97, exact head 61122f92: - All five new native HTTP-mock tests pass (DONE/FAILED/CANCELLED, pending deadline, legacy response, no wait, maximum timeout). On unchanged current-main production the job-tracking and deadline regressions fail by returning early against the old index. - All 274 remote-table tests pass, no failures/skips. - cargo fmt --all -- --check and Clippy --lib --tests -- -D warnings pass. Tests and Clippy used --offline --profile ci --locked -p lancedb --no-default-features --features remote. The affected package was rebuilt from each selected source root; dependency caches were reused. This is Rust mocked-HTTP coverage, not live-cloud, Python/Node binding, full-workspace or performance qualification. --- rust/lancedb/src/index/waiter.rs | 2 +- rust/lancedb/src/remote/table.rs | 194 +++++++++++++++++++++++++++++-- 2 files changed, 185 insertions(+), 11 deletions(-) diff --git a/rust/lancedb/src/index/waiter.rs b/rust/lancedb/src/index/waiter.rs index f201ed6c3..eeef22c98 100644 --- a/rust/lancedb/src/index/waiter.rs +++ b/rust/lancedb/src/index/waiter.rs @@ -9,7 +9,7 @@ use std::time::{Duration, Instant}; use tokio::time::sleep; const DEFAULT_SLEEP_MS: u64 = 1000; -const MAX_WAIT: Duration = Duration::from_secs(2 * 60 * 60); +pub(crate) const MAX_WAIT: Duration = Duration::from_secs(2 * 60 * 60); /// Poll the table using list_indices() and index_stats() until all of the indices have 0 un-indexed rows. /// Will return Error::Timeout if the columns are not fully indexed within the timeout. diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 1c569f6e5..c3b6e7b73 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -16,7 +16,7 @@ use crate::expr::expr_to_sql_string; use crate::index::Index; use crate::index::IndexStatistics; use crate::index::scalar::FtsQuery; -use crate::index::waiter::wait_for_index; +use crate::index::waiter::{MAX_WAIT, wait_for_index}; use crate::job::Job; use crate::materialized_view::{ MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, @@ -650,13 +650,42 @@ impl RemoteTable { .and_then(|body| extract_job_id(&body)); if let Some(wait_timeout) = index.wait_timeout { - let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); - self.wait_for_index(&[&index_name], wait_timeout).await?; + if let Some(job_id) = &job_id { + if wait_timeout > MAX_WAIT { + return Err(Error::InvalidInput { + message: format!("timeout must be less than {:?}", MAX_WAIT), + }); + } + // An older same-named index may still be ready while its + // replacement is building. Wait for this submission instead. + tokio::time::timeout(wait_timeout, self.index_job(job_id.clone()).wait()) + .await + .map_err(|_| Error::Timeout { + message: format!( + "timed out waiting for index job {:?} after {:?}", + job_id, wait_timeout + ), + })??; + } else { + // Older servers do not return a job ID. + let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column)); + self.wait_for_index(&[&index_name], wait_timeout).await?; + } } Ok(job_id) } + fn index_job(&self, job_id: String) -> Job { + Job::new(Box::new(FreshnessJob { + inner: RemoteJob::new(self.client.clone(), job_id), + freshness: self.freshness.clone(), + version: self.version.clone(), + tracked_result: TrackedJobResult::None, + freshness_request: self.snapshot_freshness_headers(), + })) + } + pub(super) fn new_with_sql_client( client: RestfulLanceDbClient, name: String, @@ -3101,13 +3130,7 @@ impl BaseTable for RemoteTable { async fn create_index_async(&self, index: IndexBuilder) -> Result { Ok(match self.submit_create_index(index).await? { - Some(job_id) => Job::new(Box::new(FreshnessJob { - inner: RemoteJob::new(self.client.clone(), job_id), - freshness: self.freshness.clone(), - version: self.version.clone(), - tracked_result: TrackedJobResult::None, - freshness_request: self.snapshot_freshness_headers(), - })), + Some(job_id) => self.index_job(job_id), None => Job::new_done(), }) } @@ -8000,6 +8023,157 @@ mod tests { assert_eq!(err.to_string(), "Job job-err failed: preempted (in commit)"); } + // The name-based endpoint deliberately exposes an already-ready old index. + // A replacement must wait on the new job, not mistake that index for success. + fn index_wait_fixture( + job_id: Option<&'static str>, + terminal_state: &'static str, + ) -> ( + Table, + Arc, + Arc, + ) { + let job_polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let index_polls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let jobs = job_polls.clone(); + let indices = index_polls.clone(); + let table = Table::new_with_handler("my_table", move |request| { + let body = match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + return http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(); + } + "/v1/table/my_table/create_index/" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + // Replacement is the default and omitted from the wire body. + assert!(body.get("replace").is_none()); + match job_id { + Some(id) => json!({"job_id": id}), + None => json!({}), + } + } + "/v1/jobs/describe" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["job_id"], json!(job_id.unwrap())); + let poll = jobs.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + json!({ + "job_id": job_id.unwrap(), + "job_state": if poll == 0 { "IN_PROGRESS" } else { terminal_state }, + "failure": {"phase": "build", "message": "invalid index config", + "retryable": false} + }) + } + "/v1/table/my_table/index/list/" => { + indices.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + json!({"indexes": [{"index_name": "a_idx", "columns": ["a"], + "index_uuid": "3fa85f64-5717-4562-b3fc-2c963f66afa6", + "index_status": "done"}]}) + } + "/v1/table/my_table/index/a_idx/stats/" => json!({ + "num_indexed_rows": 100, "num_unindexed_rows": 0, "index_type": "BTREE" + }), + path => panic!("Unexpected path: {path}"), + }; + http::Response::builder() + .status(200) + .body(body.to_string()) + .unwrap() + }); + (table, job_polls, index_polls) + } + + #[tokio::test] + async fn test_create_index_wait_timeout_tracks_job() { + for state in ["DONE", "FAILED", "CANCELLED"] { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), state); + let result = table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_secs(2)) + .execute() + .await; + match state { + "DONE" => result.unwrap(), + "FAILED" => match result.unwrap_err() { + Error::JobFailed { job_id, failure } => { + assert_eq!(job_id.as_deref(), Some("new-index-job")); + assert_eq!(failure.message.as_deref(), Some("invalid index config")); + assert_eq!(failure.phase.as_deref(), Some("build")); + } + error => panic!("Expected job failure, got {error:?}"), + }, + "CANCELLED" => assert!(matches!(result, Err(Error::JobCancelled { .. }))), + _ => unreachable!(), + } + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 2); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + } + + #[tokio::test] + async fn test_create_index_wait_timeout_bounds_job_polling() { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), "IN_PROGRESS"); + let error = table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_millis(10)) + .execute() + .await + .unwrap_err(); + assert!(matches!(error, Error::Timeout { message } if message.contains("new-index-job"))); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn test_create_index_wait_timeout_legacy_fallback() { + let (table, jobs, indices) = index_wait_fixture(None, "DONE"); + table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_secs(2)) + .execute() + .await + .unwrap(); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_create_index_wait_timeout_absent_does_not_poll() { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), "IN_PROGRESS"); + table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .execute() + .await + .unwrap(); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn test_create_index_wait_timeout_preserves_maximum() { + let (table, jobs, indices) = index_wait_fixture(Some("new-index-job"), "DONE"); + let error = table + .create_index(&["a"], Index::BTree(Default::default())) + .replace(true) + .wait_timeout(Duration::from_secs(2 * 60 * 60 + 1)) + .execute() + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidInput { .. })); + assert_eq!(jobs.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!(indices.load(std::sync::atomic::Ordering::SeqCst), 0); + } + /// Servers that return no job id (e.g. an empty create-index response) /// yield an already-done job. #[tokio::test] From e489d0cd7b272bab25066c25a951fd2bac8f8c15 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:43:07 +0800 Subject: [PATCH 24/68] fix: return errors for unsupported remote alteration types (#4430) RemoteTable::alter_columns currently unwraps the JsonDataType conversion. A valid Arrow input that the remote wire format cannot represent (for example Utf8View, BinaryView, or a struct containing Utf8View) therefore panics instead of returning an error. Propagate conversion/serialization failures as Error::InvalidInput, including the column path. Unsupported inputs are rejected before any HTTP request; supported alteration payloads are unchanged. This does not add server support for view types or change server HTTP-error classification. Three new native Rust cases reproduce the panic on unchanged production and return InvalidInput with the fix. The existing supported-type cases remain passing. Validation at head 88184ea, Linux/Rust 1.97.0, locked offline dependencies: - cargo test --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib remote::table::tests:: (269 passed) - cargo clippy --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib --tests -- -D warnings - cargo fmt --all -- --check Current main merges cleanly; the intervening remote-table change is unrelated nprobes validation. Live Cloud, rebuilt Python/Node bindings and full-workspace integration were not tested. AI-assisted; native regression and before/after results independently checked. --- rust/lancedb/src/remote/table.rs | 37 +++++++++++++++++++++++++++----- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index c3b6e7b73..9e2641c22 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3736,16 +3736,25 @@ impl BaseTable for RemoteTable { value["rename"] = serde_json::Value::String(rename.clone()); } if let Some(data_type) = &alteration.data_type { - let json_data_type = JsonDataType::try_from(data_type).unwrap(); - let json_data_type = serde_json::to_value(&json_data_type).unwrap(); - value["data_type"] = json_data_type; + let json_data_type = + JsonDataType::try_from(data_type).map_err(|err| Error::InvalidInput { + message: format!( + "Cannot serialize data type for column '{}': {err}", + alteration.path + ), + })?; + value["data_type"] = serde_json::to_value(&json_data_type).map_err(|err| { + Error::InvalidInput { + message: format!("Cannot serialize column alteration: {err}"), + } + })?; } if let Some(nullable) = &alteration.nullable { value["nullable"] = serde_json::Value::Bool(*nullable); } - value + Ok(value) }) - .collect::>(); + .collect::>>()?; let mut body = serde_json::json!({ "alterations": body }); self.apply_branch_body(&mut body); let request = self @@ -5067,6 +5076,24 @@ mod tests { } } + #[rstest] + #[case::string_view(DataType::Utf8View)] + #[case::binary_view(DataType::BinaryView)] + #[case::nested_view(DataType::Struct(vec![Field::new("value", DataType::Utf8View, true)].into()))] + #[tokio::test] + async fn test_alter_columns_unsupported_type(#[case] data_type: DataType) { + let table = Table::new_with_handler("my_table", |_| -> http::Response { + panic!("an unserializable alteration must not reach the server") + }); + let err = table + .alter_columns(&[ColumnAlteration::new("value".into()).cast_to(data_type)]) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}"); + assert!(err.to_string().contains("value")); + assert!(err.to_string().contains("Unsupported type")); + } + #[rstest] #[case(true)] #[case(false)] From cc9f01ae2622d2a3726608b4ff875d4478daffd4 Mon Sep 17 00:00:00 2001 From: Preston Zhao Date: Thu, 22 Oct 2026 13:43:21 -0500 Subject: [PATCH 25/68] fix(node): accept exist_ok in createTable (#4432) TypeScript CreateTableOptions.mode omitted exist_ok even though createTable already forwards that mode to the native client. closes: #2900 --------- Co-authored-by: Preston Zhao <195942578+canyang25@users.noreply.github.com> --- docs/src/js/interfaces/CreateTableOptions.md | 5 ++++- nodejs/__test__/connection.test.ts | 14 ++++++++++++++ nodejs/lancedb/connection.ts | 5 ++++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/docs/src/js/interfaces/CreateTableOptions.md b/docs/src/js/interfaces/CreateTableOptions.md index f46cf3722..e6119e74c 100644 --- a/docs/src/js/interfaces/CreateTableOptions.md +++ b/docs/src/js/interfaces/CreateTableOptions.md @@ -65,7 +65,7 @@ then no error will be raised. ### mode ```ts -mode: "overwrite" | "create"; +mode: "overwrite" | "create" | "exist_ok"; ``` The mode to use when creating the table. @@ -76,6 +76,9 @@ happen. Any provided data will be ignored. If this is set to "overwrite" then any existing table will be replaced. +If this is set to "exist_ok" then an existing table is opened and any +provided data is ignored. + *** ### schema? diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 895c2a7c5..1167ae34a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -117,6 +117,20 @@ describe("given a connection", () => { await expect(tbl.countRows()).resolves.toBe(1); }); + it("should keep the existing table when mode is exist_ok", async () => { + const created = await db.createTable("test", [{ id: 1 }, { id: 2 }], { + mode: "exist_ok", + }); + await expect(created.countRows()).resolves.toBe(2); + + const reopened = await db.createTable("test", [{ id: 9 }], { + mode: "exist_ok", + }); + await expect(reopened.countRows()).resolves.toBe(2); + const rows = (await reopened.query().toArray()) as { id: number }[]; + expect(rows.map((row) => row.id).sort()).toEqual([1, 2]); + }); + it("should respect limit and page token when listing tables", async () => { const db = await connect(tmpDir.name); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 11341a8c6..4ed511b86 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -53,8 +53,11 @@ export interface CreateTableOptions { * happen. Any provided data will be ignored. * * If this is set to "overwrite" then any existing table will be replaced. + * + * If this is set to "exist_ok" then an existing table is opened and any + * provided data is ignored. */ - mode: "create" | "overwrite"; + mode: "create" | "overwrite" | "exist_ok"; /** * If this is true and the table already exists and the mode is "create" * then no error will be raised. From 13df747f180da94c3a487067e2a8cec8839a5158 Mon Sep 17 00:00:00 2001 From: Preston Zhao Date: Thu, 22 Oct 2026 13:43:36 -0500 Subject: [PATCH 26/68] fix(python): return the session supplied to LanceDBConnection (#4433) LanceDBConnection.session read AsyncConnection.session, but AsyncConnection has no session attribute, so access raised AttributeError. Store the session passed to the constructor and return that object. Assisted by Cursor (Grok). Co-authored-by: Preston Zhao <195942578+canyang25@users.noreply.github.com> --- python/python/lancedb/db.py | 11 ++++++++++- python/python/tests/test_db.py | 3 +++ python/python/tests/test_session.py | 10 ++++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index e1e5d270a..a74763afb 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1300,6 +1300,10 @@ class LanceDBConnection(DBConnection): self.storage_options = storage_options self._manifest_enabled = manifest_enabled self._namespace_client_properties = namespace_client_properties + # The native connection does not expose the Python Session wrapper. + # Keep the object supplied by the caller so this property stays + # synchronous for debugger inspection. + self._session = session if _inner is not None: self._conn = _inner # Native-derived wrappers resolve this in their async reconstruction @@ -1366,7 +1370,12 @@ class LanceDBConnection(DBConnection): @property def session(self) -> Optional[Session]: - return self._conn.session + """Return the session supplied when this connection was opened. + + None when the caller did not pass a session, including connections + reconstructed with :meth:`from_inner`. + """ + return self._session @property def uri(self) -> str: diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index 7d37910a3..af63987aa 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -86,6 +86,7 @@ def test_sync_debugger_inspection_does_not_use_background_loop(tmp_path, monkeyp assert members["uri"] == str(tmp_path) assert members["read_consistency_interval"] is None + assert members["session"] is None assert repr(db) == f"LanceDBConnection(uri={str(tmp_path)!r})" assert repr(table) == f"LanceTable(name='test', _conn={db!r})" @@ -105,6 +106,8 @@ def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkey assert db.read_consistency_interval == consistency_interval assert db_from_inner.read_consistency_interval == consistency_interval + assert db.session is None + assert db_from_inner.session is None def test_ingest_pd(tmp_path): diff --git a/python/python/tests/test_session.py b/python/python/tests/test_session.py index 8923477c7..d0bc54210 100644 --- a/python/python/tests/test_session.py +++ b/python/python/tests/test_session.py @@ -36,3 +36,13 @@ def test_session_cache_configuration(tmp_path): assert final_cache_size > initial_cache_size # Cache should have grown assert final_cache_items >= initial_cache_items # Items should not decrease assert initial_cache_size < index_cache_size + metadata_cache_size + + +def test_connection_session_returns_supplied_session(tmp_path): + """Sync connections must expose the session they were opened with.""" + db = lancedb.connect(tmp_path) + assert db.session is None + + session = lancedb.Session() + db = lancedb.connect(tmp_path, session=session) + assert db.session is session From 1e313d3c0d828d04f6c7cb3036d896942c48c86c Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:43:50 +0800 Subject: [PATCH 27/68] fix(python): parse decimal exponent literals exactly (#4449) Decimal literals such as Decimal("1E-7") and Decimal("1E+3") currently raise ValueError while building an expression. Parse the coefficient and exponent exactly, and normalize representable negative scales to zero so both typed filters and predicate.to_sql() queries preserve the value. Checked arithmetic bounds the Decimal128 coefficient to 38 digits, preserves existing scale limits, and rejects nonfinite or overflowing values. This does not add Decimal256 literal support. Added 21 cases covering actual Arrow-backed table queries, SQL queries, zero, signed exponents, precision boundaries, and errors. Both Python 3.11 and 3.13 expression suites pass 123 tests; the baseline fails 10 of the new cases. Python binding Rust tests pass 13; official make doctest passes 77 with 6 skips; Ruff, Rustfmt, hooks, Pyright and workspace/tests/all-features Clippy with -D warnings pass. A fresh environment restored from the original dependency freeze passes the full required suite: 1528 passed, 37 skipped, 68 deselected (143.51s); while optional Torch suites have the same baseline remote mock-server ordering failure. Tested against f4544628; observed official main d864e719 contains one Catalog API commit without expr or Cargo.lock changes. AI assistance: gpt-6.1-sol assisted implementation and local validation; all reported checks were actually executed. --- python/python/tests/test_expr.py | 61 ++++++++++++++++++++++++++++++++ python/src/expr.rs | 58 ++++++++++++++++++++++++------ 2 files changed, 109 insertions(+), 10 deletions(-) diff --git a/python/python/tests/test_expr.py b/python/python/tests/test_expr.py index 0f49231f1..6cfc09922 100644 --- a/python/python/tests/test_expr.py +++ b/python/python/tests/test_expr.py @@ -80,6 +80,67 @@ class TestExprConstruction: e = lit(d) assert isinstance(e, Expr) + @pytest.mark.parametrize( + "value, decimal_type", + [ + ("1E-7", pa.decimal128(12, 7)), + ("-1.23E-7", pa.decimal128(12, 9)), + ("1E+3", pa.decimal128(12, 0)), + ("-1.23E+5", pa.decimal128(12, 0)), + ("0E-7", pa.decimal128(12, 7)), + ("-0E+3", pa.decimal128(12, 0)), + ("1E-38", pa.decimal128(38, 38)), + ("1E+37", pa.decimal128(38, 0)), + ("1.234567890123456789012345678901234567E+37", pa.decimal128(38, 0)), + ("0E+128", pa.decimal128(1, 0)), + ], + ) + def test_lit_decimal_exponent_query(self, tmp_path, value, decimal_type): + expected = Decimal(value) + table = lancedb.connect(tmp_path).create_table( + "decimals", + pa.table( + { + "id": [1, 2], + # Arrow's Python ingestion rejects 0E+128 even though its + # numeric value is representable. Store canonical zero to + # exercise the expression literal independently of ingestion. + "price": pa.array( + [Decimal(0) if expected.is_zero() else expected, None], + type=decimal_type, + ), + } + ), + ) + predicate = col("price") == lit(expected) + result = table.search().where(predicate).to_arrow() + assert result["id"].to_pylist() == [1] + assert result["price"].to_pylist() == [expected] + # Exercise the SQL serializer and parser as well as the typed query path. + sql_result = table.search().where(predicate.to_sql()).to_arrow() + assert sql_result["id"].to_pylist() == [1] + assert sql_result["price"].to_pylist() == [expected] + + @pytest.mark.parametrize( + "value", + [ + "NaN", + "sNaN", + "Infinity", + "-Infinity", + "1E-39", + "1E+38", + "1E+129", + "1E+1000000000", + "0E+129", + "12345678901234567890123456789012345678E+1", + "9" * 39, + ], + ) + def test_lit_decimal_unrepresentable_raises(self, value): + with pytest.raises(ValueError): + lit(Decimal(value)) + class TestExprOperators: def test_eq_operator(self): diff --git a/python/src/expr.rs b/python/src/expr.rs index 79b448fdf..bb779bc1d 100644 --- a/python/src/expr.rs +++ b/python/src/expr.rs @@ -240,22 +240,60 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult { fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> { let s = s.trim(); - let dot_pos = s.find('.'); - let scale = if let Some(pos) = dot_pos { - (s.len() - pos - 1) as i8 - } else { - 0 + let (mantissa, exponent) = match s.split_once(['e', 'E']) { + Some((mantissa, exponent)) => ( + mantissa, + exponent.parse::().map_err(|e| { + PyValueError::new_err(format!("failed to parse decimal exponent: {}", e)) + })?, + ), + None => (s, 0), }; + let fractional_digits = mantissa.find('.').map_or(0, |pos| mantissa.len() - pos - 1); + // Keep the coefficient intact instead of expanding an exponent into zeros. + // Arrow Decimal128 has at most 38 digits and an i8 scale (at most 38). + let mut scale = i64::try_from(fractional_digits) + .ok() + .and_then(|fractional_digits| fractional_digits.checked_sub(exponent)) + .and_then(|scale| i8::try_from(scale).ok()) + .filter(|scale| *scale <= 38) + .ok_or_else(|| PyValueError::new_err("decimal scale is outside Decimal128 range"))?; - let digits = s.replace('.', ""); - let val = digits + let digits = mantissa.replace('.', ""); + let mut val = digits .parse::() .map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?; - // Precision is total number of digits - let precision = digits.trim_start_matches('-').len() as u8; + // Positive scale cannot exceed precision, including for tiny values such as 1E-38. + let mut precision = digits + .trim_start_matches('-') + .len() + .max(usize::try_from(scale).unwrap_or(0)); + if scale < 0 { + // Arrow cannot cast SQL text to a decimal with negative scale. Normalize + // positive exponents within Decimal128's bounded coefficient range. + if val != 0 { + let shift = -i16::from(scale) as u32; + precision += shift as usize; + if precision > 38 { + return Err(PyValueError::new_err( + "decimal precision exceeds Decimal128 maximum of 38", + )); + } + val = 10_i128 + .checked_pow(shift) + .and_then(|factor| val.checked_mul(factor)) + .ok_or_else(|| PyValueError::new_err("decimal coefficient exceeds i128 range"))?; + } + scale = 0; + } + if precision > 38 { + return Err(PyValueError::new_err( + "decimal precision exceeds Decimal128 maximum of 38", + )); + } - Ok((val, precision, scale)) + Ok((val, precision as u8, scale)) } /// Call an arbitrary registered SQL function by name. From fcd26c425a8657501b9a438fd8ab35f086bc5943 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:44:05 +0800 Subject: [PATCH 28/68] fix(python): preserve integer Vector and MultiVector values (#4450) When `Vector` or `MultiVector` declares an integer Arrow value type, Pydantic currently validates every element as a float. This silently rounds integers above `2**53`, and converts `uint64`'s maximum value into an out-of-range float before table insertion. Use Pydantic's integer schema for integer Arrow value types in both validators. Floating-point types retain the existing validation. Tests exercise model serialization and real local table creation, insertion, and reads for signed and unsigned integer boundaries and floating-point controls, plus dimensions and optional/nullability behavior. Fractional floats are rejected for integer declarations; integral floats remain accepted. This does not extend ANN index support. Validation on macOS arm64, with the current main native extension built using Rust 1.97.0: - Python 3.11.14 and 3.13.11, Pydantic 2.13.5: full `pytest -m "not slow and not s3_test" -q python/tests` suite: **1,537 passed, 37 skipped, 68 deselected** on each runtime. Comparable unchanged-main baselines: **1,521 passed, 37 skipped, 68 deselected** each. - Added cases on unchanged main: **8 failed, 8 passed**; complete `test_pydantic.py` after the fix: **40 passed** on each runtime. - Documentation tests with the existing embeddings dependencies in an isolated environment, Python 3.11.14: **77 passed, 6 skipped**, matching the unchanged-main baseline. The workflow's Linux/Python 3.13 doctest matrix was not reproduced locally. - Repository Ruff lint/format, Pyright, all applicable pre-commit hooks, compile checks, and `git diff --check` passed. Eight public SDK table round trips passed on each final installed wheel runtime. Slow, credential-dependent S3 integration tests, the remaining OS/runtime matrix, and the minimum Pydantic/PyArrow dependency matrix were not run locally. The native extension was built from unchanged current main; the final local wheel repack contains the verified Python change with that identical native extension. A separate candidate maturin packaging attempt triggered Rust recompilation and was stopped, so an official candidate distribution build was not completed. The project requires Pydantic `>=2.7.4,<3`; Pydantic v1 support was removed in #3990. AI assistance: developed and verified with **gpt-6.1-sol**. --- python/python/lancedb/pydantic.py | 12 +++++- python/python/tests/test_pydantic.py | 61 ++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index 8fe3c1c93..bb70c9a67 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -127,7 +127,11 @@ def Vector( core_schema.list_schema( min_length=dim, max_length=dim, - items_schema=core_schema.float_schema(), + items_schema=( + core_schema.int_schema() + if pa.types.is_integer(value_type) + else core_schema.float_schema() + ), ), ) @@ -208,7 +212,11 @@ def MultiVector( items_schema=core_schema.list_schema( min_length=dim, max_length=dim, - items_schema=core_schema.float_schema(), + items_schema=( + core_schema.int_schema() + if pa.types.is_integer(value_type) + else core_schema.float_schema() + ), ), ), ) diff --git a/python/python/tests/test_pydantic.py b/python/python/tests/test_pydantic.py index 24fb7ee89..6717bbc32 100644 --- a/python/python/tests/test_pydantic.py +++ b/python/python/tests/test_pydantic.py @@ -415,6 +415,67 @@ def test_nullable_vector(): assert schema == pa.schema([pa.field("vec", pa.list_(pa.float32(), 16), True)]) +@pytest.mark.parametrize("vector_type", [Vector, MultiVector]) +@pytest.mark.parametrize( + "value_type, values", + [ + (pa.int64(), [-(2**53 + 1), 2**53 + 3]), + (pa.uint64(), [2**53 + 1, 2**64 - 1]), + (pa.int8(), [-128, 127]), + (pa.uint8(), [0, 255]), + (pa.float32(), [1.25, 2.5]), + (pa.float64(), [1.25, 2.5]), + ], +) +def test_typed_vector_round_trip(tmp_path, vector_type, value_type, values): + import lancedb + + annotation = vector_type(2, value_type=value_type) + expected = [values, values] if vector_type is MultiVector else values + + class Row(LanceModel): + id: int + payload: annotation + + row = Row(id=1, payload=expected) + table = lancedb.connect(tmp_path).create_table("typed_vectors", [row], schema=Row) + table.add([Row(id=2, payload=expected)]) + assert table.schema == Row.to_arrow_schema() + assert table.search().to_arrow().to_pylist() == [ + {"id": 1, "payload": expected}, + {"id": 2, "payload": expected}, + ] + assert json.loads(row.model_dump_json())["payload"] == expected + + +@pytest.mark.parametrize("vector_type", [Vector, MultiVector]) +@pytest.mark.parametrize("nullable", [False, True]) +def test_integer_vector_validation(vector_type, nullable): + annotation = vector_type(2, value_type=pa.int64(), nullable=nullable) + + class Row(LanceModel): + payload: annotation + + def shape(values): + return [values] if vector_type is MultiVector else values + + assert Row.to_arrow_schema().field("payload").nullable is nullable + assert Row(payload=shape([1.0, -2.0])).model_dump()["payload"] == shape([1, -2]) + with pytest.raises(pydantic.ValidationError, match="valid integer"): + Row(payload=shape([1.5, 2])) + with pytest.raises(pydantic.ValidationError): + Row(payload=shape([1])) + with pytest.raises(pydantic.ValidationError): + Row(payload=shape([1, 2, 3])) + with pytest.raises(pydantic.ValidationError): + Row(payload=None) + + class OptionalRow(LanceModel): + payload: Optional[annotation] = None + + assert OptionalRow(payload=None).payload is None + + def test_bare_vector_raises_clear_error(): namespace = { "__name__": "test_model_without_pyarrow", From 92ada619d2c5aac50c987534a4ebaafd295828c2 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:44:20 +0800 Subject: [PATCH 29/68] fix(python): preserve metadata during vector schema inference (#4451) Creating a table from an Arrow `list` or `large_list` vector field infers a fixed-size vector type, but currently drops that field's metadata. Descriptions and tags such as `lancedb:description` and `lancedb:tag:model` disappear even though the input field carries them. Preserve the original field metadata while replacing its inferred type. Tests use public synchronous and asynchronous table creation, reopening, and reads for floating-point and integer inference, checking the complete schema and row values. Existing dtype inference and nullability remain unchanged. Validation on macOS arm64 with an official native extension built from main `d864e719` using Rust 1.97.0: - Python 3.11.14 and 3.13.11: full `pytest -m "not slow and not s3_test" -q -rs python/tests`: **1,533 passed, 37 skipped, 68 deselected** each, wall times 331.464 s and 416.308 s. - Added public SDK cases on unchanged main: **12 failed**; after the fix: **12 passed** on each runtime. Final installed distributions passed **8 public SDK round trips** each, including fixed-size fields and fields without metadata. - Documentation tests on Python 3.11.14: **77 passed, 6 skipped**. Ruff lint/format, Pyright (**0 errors**), applicable pre-commit hooks (**4 passed, 2 without matching files**), compile checks and `git diff --check` passed. - Comparable unchanged-main full baselines from the same day, native extension and dependency versions: **1,521 passed, 37 skipped, 68 deselected** on each runtime. These were reused rather than rerun. The final local wheel is an explicitly verified repack of the official main wheel, replacing only `table.py` and `RECORD`; all production Python files match the tested source and the native extension is unchanged. No candidate Rust rebuild was performed. Slow/S3 integration tests, the remaining OS/runtime and minimum-dependency matrices, and the workflow's Linux/Python 3.13 doctest job were not run locally. The existing Rust migration draft #2970 and issues #3216/#3222 do not fix this metadata loss in the current Python path; a future migration should preserve the same metadata. During final verification main advanced to `81d52a98` through #4373 and the beta.13 version bump. Neither modified the two files here; the local native and full-suite validation above remain explicitly based on `d864e719`, and the updated native changes have not been tested locally. AI assistance: developed and verified with **gpt-6.1-sol**. --- python/python/lancedb/table.py | 1 + python/python/tests/test_table.py | 52 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 35ec7a466..9b96a1b2a 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5404,6 +5404,7 @@ def _infer_target_schema( field.name, # preserve original field name target_type, nullable=field.nullable, + metadata=field.metadata, ) schema = schema.set(i, new_field) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b241adb0a..b2d499933 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4507,6 +4507,58 @@ async def test_alter_columns_rejects_invalid_async( assert await table.schema() == initial_schema +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("list_type", [pa.list_, pa.large_list]) +@pytest.mark.parametrize( + "value_type, values, inferred_type", + [ + (pa.float64(), [1.25, 2.5], pa.float32()), + (pa.int32(), [0, 255], pa.uint8()), + (pa.int32(), [-1, 256], pa.float32()), + ], +) +async def test_inferred_vector_field_metadata( + tmp_db: DBConnection, + tmp_db_async: AsyncConnection, + asynchronous, + list_type, + value_type, + values, + inferred_type, +): + metadata = { + "lancedb:description": "Embedding values", + "lancedb:tag:model": "test-model", + } + schema = pa.schema( + [ + pa.field( + "vector", list_type(value_type), nullable=False, metadata=metadata + ), + pa.field("text", pa.string(), metadata={"description": "Source text"}), + ], + metadata={"description": "Source dataset"}, + ) + rows = [{"vector": values, "text": "first"}, {"vector": values, "text": "second"}] + data = pa.Table.from_pylist(rows, schema=schema) + if asynchronous: + await tmp_db_async.create_table("vector_metadata", data=data) + table = await tmp_db_async.open_table("vector_metadata") + actual_schema = await table.schema() + actual_data = await table.to_arrow() + else: + tmp_db.create_table("vector_metadata", data=data) + table = tmp_db.open_table("vector_metadata") + actual_schema = table.schema + actual_data = table.to_arrow() + expected_schema = schema.set( + 0, schema.field("vector").with_type(pa.list_(inferred_type, 2)) + ) + assert actual_schema.equals(expected_schema, check_metadata=True) + assert actual_data.to_pylist() == rows + + def test_update_field_metadata(mem_db: DBConnection): data = pa.table({"id": [0, 1], "category": ["a", "b"]}) table = mem_db.create_table("my_table", data=data) From b9ef048abbf62840befb3727efb6f0e2ed5932ca Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:44:34 +0800 Subject: [PATCH 30/68] fix: validate Python read consistency intervals (#4454) Negative `read_consistency_interval` values currently reach Rust's panicking `Duration::from_secs_f64` conversion. For example: ```python from datetime import timedelta import lancedb lancedb.connect("./db", read_consistency_interval=timedelta(seconds=-1)) ``` Validate the interval with `Duration::try_from_secs_f64` at the three native connection entry points and return `ValueError` before constructing the connection or async future. `None`, zero and positive intervals retain their behavior. The raw native binding also rejects NaN, infinity and overflow; the public Python API still accepts `timedelta`, not floats. Tests cover public sync/async connections, native namespace connections, real Python namespace clients, the three raw binding entry points and valid-value controls. Validation on Python 3.12 / Rust 1.97: - Same 26 regression cases against separately built native baseline and candidate: baseline 22 failed / 4 passed; candidate 26 passed. - Existing basic connection and background-loop regression tests pass. - `cargo fmt --all -- --check`, Ruff check/format and `git diff --check` pass. The native builds used the unchanged repository `ci` profile and Cargo.lock with `--no-default-features --features remote,lancedb/metrics-otel,pyo3/extension-module`. Python runtime and focused test dependency versions came from `python/uv.lock`. Source-tree tests used a minimal distribution-version metadata fixture derived from the unchanged Cargo/pyproject manifests; both imported module origins and native versions were checked in fresh processes. This is not a complete Maturin packaging, cloud backend or full Python suite result. No GPU was used. AI-assisted diagnosis, implementation and validation. Signed-off-by: Xucheng Zhou --- .../tests/test_read_consistency_interval.py | 83 +++++++++++++++++++ python/src/connection.rs | 18 +++- 2 files changed, 98 insertions(+), 3 deletions(-) create mode 100644 python/python/tests/test_read_consistency_interval.py diff --git a/python/python/tests/test_read_consistency_interval.py b/python/python/tests/test_read_consistency_interval.py new file mode 100644 index 000000000..bc69775ea --- /dev/null +++ b/python/python/tests/test_read_consistency_interval.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import timedelta + +import lancedb +import pytest +from lance_namespace import connect as connect_namespace_client +from lancedb import _lancedb +from lancedb.namespace import LanceNamespaceDBConnection + + +@pytest.mark.parametrize("seconds", [-1.0, -0.000001]) +@pytest.mark.parametrize("kind", ["sync", "namespace", "namespace_async", "client"]) +def test_negative_read_consistency_interval(tmp_path, seconds, kind): + interval = timedelta(seconds=seconds) + with pytest.raises(ValueError, match="read_consistency_interval"): + if kind == "sync": + lancedb.connect(tmp_path, read_consistency_interval=interval) + elif kind == "client": + client = connect_namespace_client("dir", {"root": str(tmp_path)}) + LanceNamespaceDBConnection(client, read_consistency_interval=interval) + else: + connect = ( + lancedb.connect_namespace + if kind == "namespace" + else lancedb.connect_namespace_async + ) + connect("dir", {"root": str(tmp_path)}, read_consistency_interval=interval) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [-1.0, -0.000001]) +async def test_negative_async_read_consistency_interval(tmp_path, seconds): + with pytest.raises(ValueError, match="read_consistency_interval"): + await lancedb.connect_async( + tmp_path, read_consistency_interval=timedelta(seconds=seconds) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("seconds", [float("nan"), float("inf"), -float("inf"), 1e30]) +@pytest.mark.parametrize("kind", ["connect", "namespace", "client"]) +async def test_binding_rejects_invalid_read_consistency_interval( + tmp_path, seconds, kind +): + with pytest.raises(ValueError, match="read_consistency_interval"): + if kind == "connect": + await _lancedb.connect(str(tmp_path), read_consistency_interval=seconds) + elif kind == "namespace": + _lancedb.connect_namespace( + "dir", {"root": str(tmp_path)}, read_consistency_interval=seconds + ) + else: + client = connect_namespace_client("dir", {"root": str(tmp_path)}) + _lancedb.connect_namespace_client(client, read_consistency_interval=seconds) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "interval", [None, timedelta(0), timedelta(microseconds=1), timedelta(seconds=5)] +) +async def test_valid_read_consistency_intervals(tmp_path, interval): + db = lancedb.connect(tmp_path / "sync", read_consistency_interval=interval) + assert db.read_consistency_interval == interval + db.close() + + db = await lancedb.connect_async( + tmp_path / "async", read_consistency_interval=interval + ) + assert await db.get_read_consistency_interval() == interval + db.close() + + db = lancedb.connect_namespace( + "dir", {"root": str(tmp_path / "namespace")}, read_consistency_interval=interval + ) + assert db.read_consistency_interval == interval + db.close() + + client = connect_namespace_client("dir", {"root": str(tmp_path / "client")}) + db = LanceNamespaceDBConnection(client, read_consistency_interval=interval) + assert db.read_consistency_interval == interval + db.close() diff --git a/python/src/connection.rs b/python/src/connection.rs index 7e20f5f99..baeb29780 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -49,6 +49,18 @@ impl Connection { } } +fn parse_read_consistency_interval(interval: Option) -> PyResult> { + interval + .map(|seconds| { + Duration::try_from_secs_f64(seconds).map_err(|_| { + PyValueError::new_err( + "read_consistency_interval must be finite, non-negative, and within the supported range", + ) + }) + }) + .transpose() +} + fn parse_namespace_client_pushdown_operations( operations: Option>, ) -> PyResult> { @@ -1077,6 +1089,7 @@ pub fn connect( namespace_client_properties: Option>, oauth_config: Option, ) -> PyResult> { + let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?; future_into_py(py, async move { let mut builder = lancedb::connect(&uri); if let Some(api_key) = api_key { @@ -1095,7 +1108,6 @@ pub fn connect( #[cfg(not(feature = "remote"))] let _ = sql_host_override; if let Some(read_consistency_interval) = read_consistency_interval { - let read_consistency_interval = Duration::from_secs_f64(read_consistency_interval); builder = builder.read_consistency_interval(read_consistency_interval); } if let Some(storage_options) = storage_options { @@ -1144,7 +1156,7 @@ pub fn connect_namespace_client( namespace_client_impl: Option, namespace_client_properties: Option>, ) -> PyResult { - let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64); + let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?; let namespace_client_pushdown_operations = parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?; let ns_properties = namespace_client_properties.unwrap_or_default(); @@ -1201,7 +1213,7 @@ pub fn connect_namespace( session: Option, namespace_client_pushdown_operations: Option>, ) -> PyResult { - let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64); + let read_consistency_interval = parse_read_consistency_interval(read_consistency_interval)?; let namespace_client_pushdown_operations = parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?; From a12aef2a8ce39b7d85c28033bd7eff0d1e60b1da Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:44:48 +0800 Subject: [PATCH 31/68] fix: return a Python error when querying a closed table (#4455) `AsyncTable.query()` unwraps the native closed-table error, so querying a closed table raises `PanicException` instead of the documented `RuntimeError`. `AsyncTable.head()` reaches the same path. Propagate `inner_ref()`'s error through `PyResult` instead. Add regressions for closed `query()` and `head()`, with `count_rows()`, repeated close, and normal open-table queries as controls. Validation against separately built native bindings: - Before: 2 regression failures (`PanicException`), 2 controls passed. - After: all 4 regressions passed. - Existing query tests plus the basic database test: 102 passed, 1 skipped on each arm (Torch is not installed). - `cargo fmt --all -- --check`, Ruff, and `git diff --check` passed. The native build used the locked Cargo dependencies with `remote,lancedb/metrics-otel,pyo3/extension-module`. This is source-tree native testing, not a complete wheel-packaging or cloud-service test. Signed-off-by: Xucheng Zhou --- .../python/tests/test_closed_table_query.py | 32 +++++++++++++++++++ python/src/table.rs | 4 +-- 2 files changed, 34 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_closed_table_query.py diff --git a/python/python/tests/test_closed_table_query.py b/python/python/tests/test_closed_table_query.py new file mode 100644 index 000000000..2ff3a8e98 --- /dev/null +++ b/python/python/tests/test_closed_table_query.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import pytest +from lancedb.db import AsyncConnection + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["count_rows", "query", "head"]) +async def test_closed_table_queries_raise_runtime_error( + mem_db_async: AsyncConnection, operation +): + table = await mem_db_async.create_table("closed_query", data=[{"id": 0}]) + assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}] + table.close() + table.close() + assert not table.is_open() + + with pytest.raises(RuntimeError, match="Table closed_query is closed"): + if operation == "query": + table.query() + else: + await getattr(table, operation)() + + +@pytest.mark.asyncio +async def test_open_table_queries(mem_db_async: AsyncConnection): + table = await mem_db_async.create_table("open_query", data=[{"id": 0}]) + assert await table.count_rows() == 1 + assert (await table.query().to_arrow()).to_pylist() == [{"id": 0}] + assert (await table.head()).to_pylist() == [{"id": 0}] + table.close() diff --git a/python/src/table.rs b/python/src/table.rs index c1d23d151..27ed8be8b 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1428,8 +1428,8 @@ impl Table { }) } - pub fn query(&self) -> Query { - Query::new(self.inner_ref().unwrap().query()) + pub fn query(&self) -> PyResult { + Ok(Query::new(self.inner_ref()?.query())) } #[getter] From 7a3709b3768f3d1a496e1baf5832493de25e6c87 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:45:02 +0800 Subject: [PATCH 32/68] fix(remote): preserve request IDs in multipart response errors (#4456) Keep the request ID when a multipart create/complete response has a missing or incorrectly typed `upload_id`/`version`. These errors currently use an empty ID, unlike JSON parse errors, so callers lose the ID needed to correlate the failed request. Add six cases through `Table.add(...).write_parallelism(2).execute()`, capturing the actual outgoing `x-request-id`: four missing/wrong-type cases fail on unchanged main, while two invalid-JSON controls already pass. After the fix all six pass, and all 15 multipart tests pass. Native Linux validation (locked dependencies, `ci` profile): - `cargo test --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib remote::table::tests::test_multipart_write`: 15 passed. - `cargo fmt --all -- --check`: passed. - `cargo clippy --offline --profile ci --locked -p lancedb --no-default-features --features remote --lib --tests -- -D warnings`: passed. The committed source bytes match the tested source. HTTP endpoints are mocked; this does not claim a live Cloud or full-repository test run. --- rust/lancedb/src/remote/table.rs | 69 ++++++++++++++++++++++++++++++-- 1 file changed, 65 insertions(+), 4 deletions(-) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 9e2641c22..84c1b3fef 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1277,7 +1277,7 @@ impl RemoteTable { let body = response.text().await.err_to_http(request_id.clone())?; let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { source: format!("Failed to parse multipart create response: {}", e).into(), - request_id, + request_id: request_id.clone(), status_code: None, })?; parsed["upload_id"] @@ -1285,7 +1285,7 @@ impl RemoteTable { .map(|s| s.to_string()) .ok_or_else(|| Error::Http { source: "Missing upload_id in multipart create response".into(), - request_id: String::new(), + request_id, status_code: None, }) } @@ -1304,12 +1304,12 @@ impl RemoteTable { let body = response.text().await.err_to_http(request_id.clone())?; let parsed: serde_json::Value = serde_json::from_str(&body).map_err(|e| Error::Http { source: format!("Failed to parse multipart complete response: {}", e).into(), - request_id, + request_id: request_id.clone(), status_code: None, })?; let version = parsed["version"].as_u64().ok_or_else(|| Error::Http { source: "Missing version in multipart complete response".into(), - request_id: String::new(), + request_id, status_code: None, })?; Ok(AddResult { version }) @@ -12596,6 +12596,67 @@ mod tests { ); } + #[rstest] + #[case("create", "{}")] + #[case("create", r#"{"upload_id": 1}"#)] + #[case("create", "not json")] + #[case("complete", "{}")] + #[case("complete", r#"{"version": "invalid"}"#)] + #[case("complete", "not json")] + #[tokio::test] + async fn test_multipart_write_error_preserves_request_id( + #[case] failed_stage: &str, + #[case] body: &str, + ) { + let failed_stage = failed_stage.to_string(); + let body = body.to_string(); + let request_id = Arc::new(std::sync::Mutex::new(String::new())); + let observed_request_id = request_id.clone(); + let table = Table::new_with_handler_version( + "my_table", + semver::Version::new(0, 4, 0), + move |request| { + let path = request.url().path(); + if path == "/v1/table/my_table/describe/" { + return simple_describe_response(); + } + if path == format!("/v1/table/my_table/multipart_write/{failed_stage}") { + *observed_request_id.lock().unwrap() = request.headers()["x-request-id"] + .to_str() + .unwrap() + .to_string(); + return http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap(); + } + let response = match path { + "/v1/table/my_table/multipart_write/create" => r#"{"upload_id":"upload-1"}"#, + "/v1/table/my_table/insert/" => r#"{"version":1}"#, + "/v1/table/my_table/multipart_write/abort" => "", + _ => panic!("Unexpected request path: {path}"), + }; + http::Response::builder() + .status(200) + .body(response.to_string()) + .unwrap() + }, + ); + let batch = record_batch!(("id", Int32, [1, 2, 3])).unwrap(); + let error = table + .add(vec![batch]) + .write_parallelism(2) + .execute() + .await + .unwrap_err(); + let expected = request_id.lock().unwrap(); + assert!(!expected.is_empty()); + match error { + Error::Http { request_id, .. } => assert_eq!(request_id, *expected), + other => panic!("Expected Http error, got {other:?}"), + } + } + #[tokio::test] async fn test_multipart_write_progress() { let callback_count = Arc::new(AtomicUsize::new(0)); From 9c161869e1b90f3e7733ff851d113f04935c309f Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 02:45:19 +0800 Subject: [PATCH 33/68] fix(node): respect transformers tokenizer option (#4457) ## Summary `tokenizer` is a documented option, but initialization always passed the model name to `AutoTokenizer.from_pretrained`. Preserve the resolved tokenizer option, fall back to the model name when it is omitted, and identify the selected tokenizer in loading errors. Add five initialization regressions covering default model/tokenizer selection, an explicit tokenizer, registry variables, and tokenizer loading errors. They run in native Node child processes so the real ESM import works without changing Jest/CI flags. Model loaders are mocked and local/remote model access is disabled; these tests do not download models or test inference. ## Validation - Before: three regressions fail and two default-selection controls pass. After: all five pass. - New tests plus the existing registry and embedding-registry suites: 48 passed. - TypeScript compilation, changed-file Biome checks, and `pnpm run docs` pass; generated tracked docs are unchanged. - Checked with Node 24.14.1 / pnpm 11.1.1 on Linux, both with default Jest flags and the Linux CI ESM flag. The native binding and declarations came from the official `0.40.0-beta.13` packages, with registry integrity verified. Rebuilding Rust from source timed out during Cargo metadata resolution, before compilation; no Rust code is changed. --- .../fixtures/transformers_tokenizer.cjs | 66 +++++++++++++++++++ nodejs/__test__/transformers.test.ts | 26 ++++++++ nodejs/lancedb/embedding/transformers.ts | 6 +- 3 files changed, 96 insertions(+), 2 deletions(-) create mode 100644 nodejs/__test__/fixtures/transformers_tokenizer.cjs create mode 100644 nodejs/__test__/transformers.test.ts diff --git a/nodejs/__test__/fixtures/transformers_tokenizer.cjs b/nodejs/__test__/fixtures/transformers_tokenizer.cjs new file mode 100644 index 000000000..bf30cc229 --- /dev/null +++ b/nodejs/__test__/fixtures/transformers_tokenizer.cjs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +const assert = require("node:assert/strict"); +const { + TransformersEmbeddingFunction, +} = require("../../dist/embedding/transformers"); +const { getRegistry } = require("../../dist/embedding/registry"); + +async function main() { + const transformers = await import("@huggingface/transformers"); + // Never fetch a real model, including if a loader mock stops intercepting. + transformers.env.allowRemoteModels = false; + transformers.env.allowLocalModels = false; + const modelCalls = []; + const tokenizerCalls = []; + const testCase = Number(process.argv[2]); + const defaults = "Xenova/all-MiniLM-L6-v2"; + const cases = [ + [{}, defaults, defaults], + [{ model: "test/model" }, "test/model", "test/model"], + [ + { model: "test/model", tokenizer: "test/tokenizer" }, + "test/model", + "test/tokenizer", + ], + [ + { model: "test/model", tokenizer: "$var:tokenizer" }, + "test/model", + "test/tokenizer", + ], + [ + { model: "test/model", tokenizer: "test/tokenizer" }, + "test/model", + "test/tokenizer", + ], + ]; + assert.ok( + Number.isInteger(testCase) && testCase >= 0 && testCase < cases.length, + ); + transformers.AutoModel.from_pretrained = async (...args) => { + modelCalls.push(args); + return {}; + }; + transformers.AutoTokenizer.from_pretrained = async (...args) => { + tokenizerCalls.push(args); + if (testCase === 4) throw new Error("missing tokenizer"); + return {}; + }; + getRegistry().setVar("tokenizer", "test/tokenizer"); + const [options, model, tokenizer] = cases[testCase]; + const fn = new TransformersEmbeddingFunction(options); + if (testCase === 4) { + await assert.rejects(fn.init(), /error loading tokenizer test\/tokenizer/); + } else { + await fn.init(); + } + assert.deepEqual(modelCalls, [[model, { dtype: "fp32" }]]); + assert.deepEqual(tokenizerCalls, [[tokenizer]]); + assert.deepEqual(fn.toJSON(), options); +} + +main().catch((error) => { + console.error(error); + process.exitCode = 1; +}); diff --git a/nodejs/__test__/transformers.test.ts b/nodejs/__test__/transformers.test.ts new file mode 100644 index 000000000..ac66c1499 --- /dev/null +++ b/nodejs/__test__/transformers.test.ts @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import { execFileSync } from "node:child_process"; +import { resolve } from "node:path"; + +describe("Transformers embedding tokenizer selection", () => { + // Run init's native ESM import outside Jest's CommonJS sandbox, like a user. + // The fixture mocks model loading and forbids local/remote model access. + test.each([ + ["default model", 0], + ["default tokenizer", 1], + ["custom tokenizer", 2], + ["tokenizer registry variable", 3], + ["tokenizer loading error", 4], + ])("handles %s", (_name, testCase) => { + execFileSync( + process.execPath, + [ + resolve(__dirname, "fixtures", "transformers_tokenizer.cjs"), + String(testCase), + ], + { stdio: "pipe", timeout: 15_000 }, + ); + }); +}); diff --git a/nodejs/lancedb/embedding/transformers.ts b/nodejs/lancedb/embedding/transformers.ts index 06043ea7c..7e57bc6bc 100644 --- a/nodejs/lancedb/embedding/transformers.ts +++ b/nodejs/lancedb/embedding/transformers.ts @@ -38,6 +38,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< #model?: import("@huggingface/transformers").PreTrainedModel; #tokenizer?: import("@huggingface/transformers").PreTrainedTokenizer; #modelName: XenovaTransformerOptions["model"]; + #tokenizerName: string; #initialized = false; #tokenizerOptions: XenovaTransformerOptions["tokenizerOptions"]; #ndims?: number; @@ -58,6 +59,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< this.#ndims = options.ndims; this.#modelName = modelName; + this.#tokenizerName = options.tokenizer ?? modelName; } async init() { @@ -86,11 +88,11 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction< } try { this.#tokenizer = await transformers.AutoTokenizer.from_pretrained( - this.#modelName, + this.#tokenizerName, ); } catch (e) { throw new Error( - `error loading tokenizer for ${this.#modelName}. Make sure you are using a wasm compatible model:\nReason: ${e}`, + `error loading tokenizer ${this.#tokenizerName}. Make sure you are using a wasm compatible model:\nReason: ${e}`, ); } this.#initialized = true; From e1e87e1cae05a39ccb736c516405649c1ce92380 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:45:36 +0800 Subject: [PATCH 34/68] fix(python): preserve logical rows from Hugging Face datasets (#4461) Hugging Face `select`, `filter`, and `shuffle` can retain a physical Arrow table while storing logical row selection/order in an indices table. Ingesting the physical table then silently writes excluded rows or the wrong order and can lose duplicate selections. Read indexed datasets through bounded Arrow-format iteration, and keep the existing direct batch path for datasets without indices. Apply the same conversion to sync/async table creation and additions, including DatasetDict splits. The conversion does not mutate the user's formatting settings. Regression coverage includes reordered/duplicate selections, filtering, shuffling, empty selections, projected numpy/pandas/Arrow formatting, DatasetDicts, rescanning, and bounded batches. Validation on base d864e7198fe55a4e5fbe747fee574737fa102aab with its byte-identical native/config: - Python 3.11 full CI test selection: 1564 passed, 36 skipped, 68 deselected; 162.51s pytest /171.68s wall. - Python 3.13 full CI test selection: 1521 passed, 37 skipped, 68 deselected; 160.32s pytest /163.74s wall. Optional HF tests were skipped in this environment. - Python 3.13 complete doctest command: 71 passed, 6 skipped; 8.42s pytest /12.14s wall. - Python 3.11 HF regression module from a candidate wheel: 43 passed; the same regression module on the unchanged baseline gives 24 failed /19 passed. - Ruff check/format, repository pre-commit hooks, pyright with the source import path, and `git diff --check` passed. This is a Python-only change. Native binaries were reused from the verified baseline; no new native build or external-service integration tests were run. The separate schema-metadata preservation change in #4451 addresses a different conversion defect. The target branch currently includes newer native changes than the verified base. Latest-main native compatibility, Linux CI, and remote-service tests remain for CI/maintainer validation. Prepared with AI assistance (gpt-6.1-sol). --- python/python/lancedb/scannable.py | 20 ++++- python/python/lancedb/table.py | 6 +- python/python/tests/test_huggingface.py | 103 ++++++++++++++++++++++++ 3 files changed, 124 insertions(+), 5 deletions(-) diff --git a/python/python/lancedb/scannable.py b/python/python/lancedb/scannable.py index baf8e0c2b..9a6f71a5d 100644 --- a/python/python/lancedb/scannable.py +++ b/python/python/lancedb/scannable.py @@ -276,6 +276,14 @@ def _from_iterable(data: Iterator) -> Scannable: _registered_modules: set[str] = set() +def _hf_batches(data): + if data._indices is None: + yield from data.data.to_batches() + else: + for table in data.with_format("arrow").iter(batch_size=1000): + yield from table.to_batches() + + def _register_optional_converters(): """Register converters for optional dependencies that are already imported.""" @@ -323,7 +331,15 @@ def _register_optional_converters(): def _from_hf_dataset(data: HFDataset) -> Scannable: table = data.data.table # Access underlying Arrow table return Scannable( - schema=table.schema, num_rows=len(data), reader=table.to_reader + schema=table.schema, + num_rows=len(data), + reader=( + table.to_reader + if data._indices is None + else lambda: pa.RecordBatchReader.from_batches( + table.schema, _hf_batches(data) + ) + ), ) @to_scannable.register(HFDatasetDict) @@ -335,7 +351,7 @@ def _register_optional_converters(): def gen(): for split_name, dataset in data.items(): - for batch in dataset.data.to_batches(): + for batch in _hf_batches(dataset): split_arr = pa.array( [split_name] * len(batch), type=pa.string() ) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 9b96a1b2a..f9cfac6e1 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -28,7 +28,7 @@ from typing import ( ) from urllib.parse import urlparse -from lancedb.scannable import _register_optional_converters, to_scannable +from lancedb.scannable import _hf_batches, _register_optional_converters, to_scannable from . import __version__ from ._blob import ( @@ -400,7 +400,7 @@ def _into_pyarrow_reader( if _check_for_hugging_face(data): if isinstance(data, datasets.Dataset): schema = data.features.arrow_schema - return pa.RecordBatchReader.from_batches(schema, data.data.to_batches()) + return pa.RecordBatchReader.from_batches(schema, _hf_batches(data)) elif isinstance(data, datasets.dataset_dict.DatasetDict): schema = _schema_from_hf(data, None) if "split" not in schema.names: @@ -1195,7 +1195,7 @@ def _to_batches_with_split(data): with an extra `split` column """ for key, dataset in data.items(): - for batch in dataset.data.to_batches(): + for batch in _hf_batches(dataset): table = pa.Table.from_batches([batch]) if "split" not in table.column_names: table = table.append_column( diff --git a/python/python/tests/test_huggingface.py b/python/python/tests/test_huggingface.py index 9b597bbc9..91421b79c 100644 --- a/python/python/tests/test_huggingface.py +++ b/python/python/tests/test_huggingface.py @@ -129,3 +129,106 @@ def test_generator(tmp_path: Path): assert len(tbl) == 2 assert tbl.schema == ds.features.arrow_schema + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["create", "add"]) +@pytest.mark.parametrize( + "selection", + ["normal", "contiguous", "selected", "duplicates", "filtered", "shuffled", "empty"], +) +async def test_write_selected_hf_dataset(tmp_path, asynchronous, operation, selection): + source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")}) + data = { + "normal": source, + "contiguous": source.select(range(1, 4)), + "selected": source.select([4, 1]), + "duplicates": source.select([2, 2, 0]), + "filtered": source.filter(lambda row: row["id"] % 2 == 1), + "shuffled": source.shuffle(seed=42), + "empty": source.select([]), + }[selection] + expected = data.to_list() + if asynchronous: + db = await lancedb.connect_async(tmp_path) + if operation == "create": + table = await db.create_table("selected", data) + else: + table = await db.create_table("selected", schema=data.features.arrow_schema) + await table.add(data) + actual = await table.to_arrow() + else: + db = lancedb.connect(tmp_path) + if operation == "create": + table = db.create_table("selected", data) + else: + table = db.create_table("selected", schema=data.features.arrow_schema) + table.add(data) + actual = table.to_arrow() + assert actual.to_pylist() == expected + assert actual.schema.equals(data.features.arrow_schema, check_metadata=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["create", "add"]) +async def test_write_selected_hf_splits(tmp_path, asynchronous, operation): + source = datasets.Dataset.from_dict({"id": list(range(5)), "text": list("abcde")}) + data = datasets.DatasetDict( + { + "train": source.select([4, 1]), + "test": source.select([3, 0, 2]), + "empty": source.select([]), + } + ) + expected = [ + dict(row, split=split) + for split, dataset in data.items() + for row in dataset.to_list() + ] + schema = source.features.arrow_schema.append(pa.field("split", pa.string())) + if asynchronous: + db = await lancedb.connect_async(tmp_path) + if operation == "create": + table = await db.create_table("splits", data) + else: + table = await db.create_table("splits", schema=schema) + await table.add(data) + actual = await table.to_arrow() + else: + db = lancedb.connect(tmp_path) + if operation == "create": + table = db.create_table("splits", data) + else: + table = db.create_table("splits", schema=schema) + table.add(data) + actual = table.to_arrow() + assert actual.to_pylist() == expected + assert actual.schema.equals(schema, check_metadata=True) + + +@pytest.mark.parametrize("format_type", [None, "numpy", "pandas", "arrow"]) +@pytest.mark.parametrize("empty", [False, True]) +def test_selected_hf_reader_is_rescannable(format_type, empty): + from lancedb.scannable import _register_optional_converters, to_scannable + + source = datasets.Dataset.from_dict( + {"id": list(range(2105)), "text": [str(i) for i in range(2105)]} + ) + indices = [] if empty else list(reversed(range(2105))) + [1, 1] + data = source.select(indices).with_format(format_type, columns=["id"]) + original_format = data.format + expected = source.select(indices).to_list() + _register_optional_converters() + scannable = to_scannable(data) + assert scannable.rescannable + assert scannable.num_rows == len(expected) + for _ in range(2): + batches = list(scannable.reader()) + if not empty: + assert all(batch.num_rows <= 1000 for batch in batches) + actual = pa.Table.from_batches(batches, schema=scannable.schema) + assert actual.to_pylist() == expected + assert actual.schema.equals(data.features.arrow_schema, check_metadata=True) + assert data.format == original_format From 8552fd2f5b8213ce6be5fdd9d3cdfe7978f10a1d Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:45:51 +0800 Subject: [PATCH 35/68] fix(python): prevent embedding aliases from replacing registered functions (#4462) Registering a different embedding class under an existing alias silently replaced the original class. A deterministic real local SDK table reproduction changed the top result for the same query from alpha to beta after reopening the persisted table. Check the resolved key before inserting a class. Preserve repeated registration of the same class under the same nonempty explicit alias as idempotent, as existing SDK tests rely on it. Repeated implicit class-name registration, including an empty alias, still raises KeyError. Failed registration preserves the mapping and class metadata. Existing embedding tests receive isolated registry mappings; test_util.py remains unchanged. Eight regression cases cover different-class alias collisions, alias/class-name collisions, implicit and empty-alias duplicates, equal class names under distinct aliases, explicit-alias idempotence, and real local creation/reopening/search. This is independent of #4356, which fixes no-alias serialization. Validation on 52a1ccee5a5b8dbfff8568abafb57099c55383b9: - Original registry plus eight new cases: 4 failed, 4 passed; Python 3.11.14; exit 1; 2.902604 s. Candidate plus full utility module: 83 passed, 1 skipped; exit 0; 3.144414 s. - Python 3.11.14 full official non-slow/non-S3 suite: 1532 passed, 36 skipped, 68 deselected; exit 0; 152.631543 s. - Python 3.13.11 full official non-slow/non-S3 suite: 1529 passed, 37 skipped, 68 deselected; exit 0; 163.970571 s. - Official Python 3.13 doctest with existing compatible embeddings/torch/mlx dependencies: 71 passed, 6 skipped; exit 0; 9.635034 s. - Pyright with matching interpreter/source path: 0 errors, 0 warnings; exit 0; 4.299743 s. - Repository Ruff 0.9.9 check and format: exit 0; 157 files already formatted; 0.047418 s and 0.016815 s. Pre-commit changed files: four hooks passed, two inapplicable skipped; exit 0; 0.406136 s. git diff --check passed. Base is official d864e7198fe55a4e5fbe747fee574737fa102aab with its verified matching native extension reused read-only. All 174 native/config paths are unchanged from that base. Latest main bbc4aebe55b69471fa1bd7655599b8d2e3f9352c has byte-identical registry production code; its complete native build was not run. Verification was macOS arm64, with optional dependency/service skips above; Linux/Windows CI remains remote verification. Exact commands, cwd, identities, raw logs and intermediate failures are in the local manifest. The earlier strict same-class check caused 13 utility failures and was corrected in production; it is not described as a baseline failure. Reproduction commands after preparing the repository's declared environment (from repository root unless noted): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/embeddings/registry.py python/python/tests/test_embedding_registry.py python/python/tests/test_embeddings.py git diff --check ``` These relative commands describe reproducible source/cwd equivalents. The recorded local runs used absolute paths to existing Python 3.11.14/3.13.11 environments, a verified matching reused native extension, and read-only compatible optional-dependency cache paths for the 3.13 doctest, rather than a newly installed `.venv`. Full tests used separate task-private working directories and `--basetemp` directories to avoid cross-run database collisions. Implemented with assistance from gpt-6.1-sol; the final diff was reviewed locally by the collaborating agents. --- python/python/lancedb/embeddings/registry.py | 6 +- .../python/tests/test_embedding_registry.py | 88 +++++++++++++++++++ python/python/tests/test_embeddings.py | 6 ++ 3 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_embedding_registry.py diff --git a/python/python/lancedb/embeddings/registry.py b/python/python/lancedb/embeddings/registry.py index 3b940fe2e..728ff36da 100644 --- a/python/python/lancedb/embeddings/registry.py +++ b/python/python/lancedb/embeddings/registry.py @@ -69,9 +69,11 @@ class EmbeddingFunctionRegistry: def decorator(cls): if not issubclass(cls, EmbeddingFunction): raise TypeError("Must be a subclass of EmbeddingFunction") - if cls.__name__ in self._functions: - raise KeyError(f"{cls.__name__} was already registered") key = alias or cls.__name__ + if key in self._functions and ( + not alias or self._functions[key] is not cls + ): + raise KeyError(f"{key} was already registered") self._functions[key] = cls # Store the name the class is registered under, so the table # metadata written for it can be looked up again when reading. diff --git a/python/python/tests/test_embedding_registry.py b/python/python/tests/test_embedding_registry.py new file mode 100644 index 000000000..07c10bd7e --- /dev/null +++ b/python/python/tests/test_embedding_registry.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import lancedb +import pytest +from lancedb.embeddings.base import TextEmbeddingFunction +from lancedb.embeddings.registry import EmbeddingFunctionRegistry, get_registry +from lancedb.pydantic import LanceModel, Vector + + +class FirstEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts] + + +class SecondEmbedding(FirstEmbedding): + def generate_embeddings(self, texts): + return [[0.0, 1.0] if text == "alpha" else [1.0, 0.0] for text in texts] + + +@pytest.mark.parametrize( + "first_alias,second_alias,second_class", + [ + ("shared", "shared", SecondEmbedding), + (None, "FirstEmbedding", SecondEmbedding), + ("SecondEmbedding", None, SecondEmbedding), + (None, None, FirstEmbedding), + ("", "", FirstEmbedding), + ], +) +def test_duplicate_registration_preserves_original( + first_alias, second_alias, second_class +): + registry = EmbeddingFunctionRegistry() + registry.register(first_alias)(FirstEmbedding) + key = first_alias or "FirstEmbedding" + original_alias = FirstEmbedding.__embedding_function_registry_alias__ + second_attributes = dict(second_class.__dict__) + with pytest.raises(KeyError, match=f"{key} was already registered"): + registry.register(second_alias)(second_class) + assert registry.get(key) is FirstEmbedding + assert FirstEmbedding.__embedding_function_registry_alias__ == original_alias + assert dict(second_class.__dict__) == second_attributes + + +def test_class_name_does_not_collide_with_distinct_alias(): + registry = EmbeddingFunctionRegistry() + registry.register()(FirstEmbedding) + same_name = type("FirstEmbedding", (SecondEmbedding,), {}) + registry.register("separate")(same_name) + assert registry.get("FirstEmbedding") is FirstEmbedding + assert registry.get("separate") is same_name + + +def test_same_class_explicit_alias_registration_is_idempotent(): + registry = EmbeddingFunctionRegistry() + assert registry.register("shared")(FirstEmbedding) is FirstEmbedding + attributes = dict(FirstEmbedding.__dict__) + assert registry.register("shared")(FirstEmbedding) is FirstEmbedding + assert registry.get("shared") is FirstEmbedding + assert dict(FirstEmbedding.__dict__) == attributes + + +def test_alias_collision_does_not_change_persisted_table(tmp_path, monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + alias = "test-alias-collision-persisted-table" + registry.register(alias)(FirstEmbedding) + function = FirstEmbedding.create(max_retries=0) + + class Row(LanceModel): + text: str = function.SourceField() + vector: Vector(2) = function.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table( + "rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}] + ) + assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + assert registry.register(alias)(FirstEmbedding) is FirstEmbedding + with pytest.raises(KeyError, match=f"{alias} was already registered"): + registry.register(alias)(SecondEmbedding) + reopened = lancedb.connect(tmp_path).open_table("rows") + assert isinstance(reopened.embedding_functions["vector"].function, FirstEmbedding) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 51907f141..d6d128845 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -23,6 +23,12 @@ from lancedb.embeddings.utils import retry from lancedb.pydantic import LanceModel, Vector +@pytest.fixture(autouse=True) +def isolated_embedding_registry(monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + + def mock_embed_func(input_data): return [np.random.randn(128).tolist() for _ in range(len(input_data))] From 554dc01f4d653f95bf044059e1f88ad7e3e71457 Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:06 +0800 Subject: [PATCH 36/68] fix(python): preserve custom embedding field metadata (#4463) `SourceField(json_schema_extra={...})` and `VectorField(json_schema_extra={...})` currently raise `TypeError` because the helper supplies the same `Field` keyword twice. Explicit `None` also fails. This prevents users from adding ordinary Pydantic JSON schema metadata to embedding fields. Copy and merge the supplied dictionary before attaching the helper's internal embedding identity. Caller dictionaries remain unchanged, and reserved embedding keys retain the real function. Other `Field` arguments keep their existing behavior. This change covers dictionaries and `None`; callable `json_schema_extra` remains unsupported by the existing embedding metadata reader and is outside this fix. The 11 regression cases exercise both helpers, empty/custom/None metadata, reserved-key collisions, other Field arguments, and a deterministic local embedding model through JSON schema generation, table creation, reopening, and search. On the original source, 9 fail and 2 controls pass; the candidate passes all 11 on Python 3.11 and 3.13. Validation uses read-only reused environments and the native extension built for d864e719. The production helper is byte-identical to current main bbc4aebe before this fix; no current-main native rebuild was performed. The branch contains only this helper fix and its regression test, independently based on d864e719. Reproducible checkout commands below assume the documented environment is installed as `.venv`; the actual local runs used existing interpreters through absolute paths and isolated generated-data working directories: ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) git diff --check .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/embeddings/base.py python/python/tests/test_embedding_field_metadata.py ``` Doctests require the repository's tests/dev/embeddings dependencies, including mlx on this platform. All checks below ran against frozen head 3fa6820. Implemented with assistance from gpt-6.1-sol; the final diff was reviewed locally. Frozen-head validation (all exit 0): - Python 3.11 full: 1535 passed, 36 skipped, 68 deselected; 190.272415 seconds wall time. - Python 3.13 full: 1532 passed, 37 skipped, 68 deselected; 249.289778 seconds wall time. - Python 3.13 doctests: 71 passed, 6 skipped; 10.391002 seconds wall time. - Pyright: 0 errors, 0 warnings; 4.884641 seconds wall time. - pre-commit: 4 passed, 2 skipped; 0.747764 seconds wall time. - Repository Ruff check: passed; 0.078177 seconds wall time. - Repository Ruff format: 157 files already formatted; 0.034561 seconds wall time. --- python/python/lancedb/embeddings/base.py | 8 +- .../tests/test_embedding_field_metadata.py | 84 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_embedding_field_metadata.py diff --git a/python/python/lancedb/embeddings/base.py b/python/python/lancedb/embeddings/base.py index f26a940df..96ca6556c 100644 --- a/python/python/lancedb/embeddings/base.py +++ b/python/python/lancedb/embeddings/base.py @@ -239,7 +239,9 @@ class EmbeddingFunction(BaseModel, ABC): Creates a pydantic Field that can automatically annotate the source column for this embedding function """ - return Field(json_schema_extra={"source_column_for": self}, **kwargs) + json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {}) + json_schema_extra["source_column_for"] = self + return Field(json_schema_extra=json_schema_extra, **kwargs) def VectorField(self, **kwargs): """ @@ -250,7 +252,9 @@ class EmbeddingFunction(BaseModel, ABC): static type checkers. LanceDB will infer the fixed vector dimension from this embedding function. """ - return Field(json_schema_extra={"vector_column_for": self}, **kwargs) + json_schema_extra = dict(kwargs.pop("json_schema_extra", None) or {}) + json_schema_extra["vector_column_for"] = self + return Field(json_schema_extra=json_schema_extra, **kwargs) def __eq__(self, __value: object) -> bool: if not hasattr(__value, "__dict__"): diff --git a/python/python/tests/test_embedding_field_metadata.py b/python/python/tests/test_embedding_field_metadata.py new file mode 100644 index 000000000..68754ee35 --- /dev/null +++ b/python/python/tests/test_embedding_field_metadata.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import copy + +import lancedb +import pytest +from lancedb.embeddings.base import TextEmbeddingFunction +from lancedb.embeddings.registry import get_registry +from lancedb.pydantic import LanceModel, Vector + + +class LocalEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + return [[1.0, 0.0] if text == "alpha" else [0.0, 1.0] for text in texts] + + +@pytest.mark.parametrize( + "method,tag", + [("SourceField", "source_column_for"), ("VectorField", "vector_column_for")], +) +@pytest.mark.parametrize("extra", [None, {}, {"x-purpose": "retrieval"}]) +def test_embedding_field_preserves_extra_metadata(method, tag, extra): + function = LocalEmbedding.create(max_retries=0) + original = copy.deepcopy(extra) + field = getattr(function, method)(json_schema_extra=extra) + assert extra == original + assert field.json_schema_extra[tag] is function + for key, value in (extra or {}).items(): + assert field.json_schema_extra[key] == value + + +@pytest.mark.parametrize( + "method,tag", + [("SourceField", "source_column_for"), ("VectorField", "vector_column_for")], +) +def test_extra_metadata_cannot_replace_embedding_identity(method, tag): + function = LocalEmbedding.create(max_retries=0) + supplied = {tag: "caller value", "x-purpose": "retrieval"} + field = getattr(function, method)(json_schema_extra=supplied) + assert supplied[tag] == "caller value" + assert field.json_schema_extra[tag] is function + assert field.json_schema_extra["x-purpose"] == "retrieval" + + +@pytest.mark.parametrize("method", ["SourceField", "VectorField"]) +def test_embedding_field_preserves_other_field_arguments(method): + function = LocalEmbedding.create(max_retries=0) + field = getattr(function, method)(title="Custom title", description="Custom field") + assert field.title == "Custom title" + assert field.description == "Custom field" + + +def test_extra_metadata_preserves_schema_and_local_embedding(tmp_path, monkeypatch): + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + registry.register("test-embedding-field-extra")(LocalEmbedding) + function = LocalEmbedding.create(max_retries=0) + + class Row(LanceModel): + text: str = function.SourceField(json_schema_extra={"x-purpose": "source"}) + vector: Vector(2) = function.VectorField( + json_schema_extra={"x-purpose": "embedding"} + ) + + json_schema = Row.model_json_schema() + assert json_schema["properties"]["text"]["x-purpose"] == "source" + assert json_schema["properties"]["vector"]["x-purpose"] == "embedding" + configs = Row.parse_embedding_functions() + assert len(configs) == 1 + assert configs[0].function is function + assert configs[0].source_column == "text" + assert configs[0].vector_column == "vector" + db = lancedb.connect(tmp_path) + table = db.create_table( + "rows", schema=Row, data=[{"text": "alpha"}, {"text": "beta"}] + ) + assert table.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + reopened = lancedb.connect(tmp_path).open_table("rows") + assert isinstance(reopened.embedding_functions["vector"].function, LocalEmbedding) + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" From 694942a1e0dcc85927ecb6c6362fb4003af8a50a Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:21 +0800 Subject: [PATCH 37/68] fix(python): convert tuple and NumPy datetime update values (#4464) `Table.update` and `AsyncTable.update` convert Python values into SQL literals. Tuple arrays and `numpy.datetime64` values currently raise `NotImplementedError`, even when the same values can be ingested into the corresponding local Arrow columns. Register tuples with the existing recursive list converter and format NumPy timestamps as complete nanosecond ISO strings with an explicit UTC offset. `NaT` becomes SQL NULL; subnanosecond units raise an explicit error rather than silently losing precision. Regression tests exercise local synchronous and asynchronous create/update/Arrow reads, date and timestamp units, nanosecond precision, NULL, and existing list/ndarray/Python datetime behavior. Nested and empty tuple arrays are compared with equivalent lists on the actual SDK: successful results must match, and the matching d864 native engine's existing list limitations must produce the same stable error without changing either table. This does not change SQL map handling. Implemented with assistance from gpt-6.1-sol; final diff reviewed locally. Local verification uses Python 3.11.14 and 3.13.11 in existing read-only development environments with the native extension matching base d864e7198fe55a4e5fbe747fee574737fa102aab. The current main `util.py` is byte-identical to that base; current-main native code has not been rebuilt. Slow tests, S3 tests, remote production services and other operating systems are not validated here. Reproducible commands after preparing the repository's documented development environment (these relative `.venv` examples describe equivalent commands; the local runs use absolute read-only environment paths and fresh task-private working directories): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests/test_util_update_values.py -q -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/util.py python/python/tests/test_util_update_values.py git diff --check d864e7198fe55a4e5fbe747fee574737fa102aab HEAD ``` Frozen commit `33e74d0ed3d90a20419816a3f1fc15ff038db452` passed all local gates: Python 3.11.14 full suite 1,590 passed / 36 skipped / 68 deselected; Python 3.13.11 full suite 1,587 passed / 37 skipped / 68 deselected. All 66 added SDK cases passed on both runtimes. Python 3.13 doctests: 71 passed / 6 skipped. Pyright: 0 errors / 0 warnings. Ruff, formatting, applicable pre-commit hooks, two-file syntax/import checks and the PR merge-base diff check exited 0. Exact local argv/runtime/wall metadata and raw logs are retained in the acceptance manifest. The follow-up addresses the non-UTC timestamp issue reported in review: an update previously changed the epoch ingested by Arrow (New York +5 hours, Tokyo -9 hours). On the published original head, 4 non-UTC sync/async SDK regressions failed and 12 naive/UTC/NaT controls passed (exit 1; wall 3.713 seconds). Formatting with `timezone="UTC"` emits a Z suffix and preserves the ingested nanosecond instant independently of the destination timezone. This is consistent with the Arrow-aligned epoch convention in merged #3775. The original 50 cases also remain green. Full-suite wall times were 158.619 seconds (3.11) and 163.631 seconds (3.13), both exit 0. --- python/python/lancedb/util.py | 12 +- .../python/tests/test_util_update_values.py | 215 ++++++++++++++++++ 2 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 python/python/tests/test_util_update_values.py diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index f1a652924..a043f0c8d 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -426,7 +426,8 @@ def _(value: date): @value_to_sql.register(list) -def _(value: list): +@value_to_sql.register(tuple) +def _(value: Union[list, tuple]): return "[" + ", ".join(map(value_to_sql, value)) + "]" @@ -471,6 +472,15 @@ def _(value: np.floating): return value_to_sql(float(value)) +@value_to_sql.register(np.datetime64) +def _(value: np.datetime64): + if np.isnat(value): + return value_to_sql(None) + if np.datetime_data(value.dtype)[0] in ("ps", "fs", "as"): + raise ValueError("SQL timestamps cannot preserve sub-nanosecond precision") + return value_to_sql(str(np.datetime_as_string(value, unit="ns", timezone="UTC"))) + + def deprecated(func): """This is a decorator which can be used to mark functions as deprecated. It will result in a warning being emitted diff --git a/python/python/tests/test_util_update_values.py b/python/python/tests/test_util_update_values.py new file mode 100644 index 000000000..fdef867da --- /dev/null +++ b/python/python/tests/test_util_update_values.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +from datetime import date, datetime + +import lancedb +import numpy as np +import pyarrow as pa +import pytest + + +ARRAY_VALUES = [ + ((1.0, 2.0), [1.0, 2.0]), + ([1.0, 2.0], [1.0, 2.0]), + (np.array([1.0, 2.0]), [1.0, 2.0]), + (None, None), +] + + +@pytest.mark.parametrize("value,expected", ARRAY_VALUES) +def test_update_array_values(tmp_path, value, expected): + dtype = pa.list_(pa.float64()) + db = lancedb.connect(tmp_path) + table = db.create_table( + "values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)}) + ) + table.update(values={"array": value}) + assert table.to_arrow()["array"].to_pylist() == [expected] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value,expected", ARRAY_VALUES) +async def test_update_array_values_async(tmp_path, value, expected): + dtype = pa.list_(pa.float64()) + db = await lancedb.connect_async(tmp_path) + table = await db.create_table( + "values", pa.table({"array": pa.array([[0.0, 0.0]], type=dtype)}) + ) + await table.update({"array": value}) + assert (await table.to_arrow())["array"].to_pylist() == [expected] + + +TIME_VALUES = [ + np.datetime64("2025", "Y"), + np.datetime64("2025-01", "M"), + np.datetime64("2025-01-02", "D"), + np.datetime64("2025-01-02T03", "h"), + np.datetime64("2025-01-02T03:04", "m"), + np.datetime64("2025-01-02T03:04:05", "s"), + np.datetime64("2025-01-02T03:04:05.123", "ms"), + np.datetime64("2025-01-02T03:04:05.123456", "us"), + np.datetime64("2025-01-02T03:04:05.123456789", "ns"), + np.datetime64("NaT", "ns"), + datetime(2025, 1, 2, 3, 4, 5, 123456), + None, +] + + +def expected_time(value): + if value is None or isinstance(value, np.datetime64) and np.isnat(value): + return None + return int(np.datetime64(value, "ns").astype(np.int64)) + + +@pytest.mark.parametrize("value", TIME_VALUES) +def test_update_timestamp_values(tmp_path, value): + db = lancedb.connect(tmp_path) + data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))}) + table = db.create_table("times", data) + table.update(values={"time": value}) + assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [ + expected_time(value) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", TIME_VALUES) +async def test_update_timestamp_values_async(tmp_path, value): + db = await lancedb.connect_async(tmp_path) + data = pa.table({"time": pa.array([datetime(2025, 1, 1)], type=pa.timestamp("ns"))}) + table = await db.create_table("times", data) + await table.update({"time": value}) + assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [ + expected_time(value) + ] + + +@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []]) +def test_nested_array_tuple_matches_list(tmp_path, values): + db = lancedb.connect(tmp_path) + data = pa.table( + {"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))} + ) + control = db.create_table("list_values", data) + candidate = db.create_table("tuple_values", data) + tuples = tuple(tuple(value) for value in values) + try: + control.update(values={"array": values}) + except ValueError as error: + reason = ( + "Expected a literal value in array" if values else "concat requires input" + ) + assert reason in str(error) + with pytest.raises(ValueError, match=reason): + candidate.update(values={"array": tuples}) + assert control.to_arrow().equals(data) + assert candidate.to_arrow().equals(data) + else: + candidate.update(values={"array": tuples}) + assert candidate.to_arrow().equals(control.to_arrow()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("values", [[[1, 2], [3, 4]], []]) +async def test_nested_array_tuple_matches_list_async(tmp_path, values): + db = await lancedb.connect_async(tmp_path) + data = pa.table( + {"array": pa.array([[[0, 0]]], type=pa.list_(pa.list_(pa.int64())))} + ) + control = await db.create_table("list_values", data) + candidate = await db.create_table("tuple_values", data) + tuples = tuple(tuple(value) for value in values) + try: + await control.update({"array": values}) + except ValueError as error: + reason = ( + "Expected a literal value in array" if values else "concat requires input" + ) + assert reason in str(error) + with pytest.raises(ValueError, match=reason): + await candidate.update({"array": tuples}) + assert (await control.to_arrow()).equals(data) + assert (await candidate.to_arrow()).equals(data) + else: + await candidate.update({"array": tuples}) + assert (await candidate.to_arrow()).equals(await control.to_arrow()) + + +@pytest.mark.parametrize("unit", ["ps", "fs", "as"]) +def test_subnanosecond_update_does_not_round(tmp_path, unit): + db = lancedb.connect(tmp_path) + data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))}) + table = db.create_table("precise", data) + with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"): + table.update(values={"time": np.datetime64(1, unit)}) + assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == [0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unit", ["ps", "fs", "as"]) +async def test_subnanosecond_update_does_not_round_async(tmp_path, unit): + db = await lancedb.connect_async(tmp_path) + data = pa.table({"time": pa.array([datetime(1970, 1, 1)], type=pa.timestamp("ns"))}) + table = await db.create_table("precise", data) + with pytest.raises(ValueError, match="cannot preserve sub-nanosecond precision"): + await table.update({"time": np.datetime64(1, unit)}) + assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == [0] + + +@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"]) +def test_update_numpy_date_column(tmp_path, unit): + value = np.datetime64("2025-01-02", unit) + db = lancedb.connect(tmp_path) + table = db.create_table( + "dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())}) + ) + table.update(values={"day": value}) + assert table.to_arrow()["day"].to_pylist() == [value.astype("datetime64[D]").item()] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unit", ["Y", "M", "W", "D"]) +async def test_update_numpy_date_column_async(tmp_path, unit): + value = np.datetime64("2025-01-02", unit) + db = await lancedb.connect_async(tmp_path) + table = await db.create_table( + "dates", pa.table({"day": pa.array([date(1970, 1, 1)], type=pa.date32())}) + ) + await table.update({"day": value}) + assert (await table.to_arrow())["day"].to_pylist() == [ + value.astype("datetime64[D]").item() + ] + + +@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"]) +@pytest.mark.parametrize( + "value", + [np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")], +) +def test_update_numpy_timestamp_matches_arrow_ingestion(tmp_path, timezone, value): + dtype = pa.timestamp("ns", tz=timezone) + ingested = pa.array(np.array([value]), type=dtype) + db = lancedb.connect(tmp_path) + table = db.create_table("timestamps", pa.table({"time": ingested})) + expected = table.to_arrow()["time"].cast(pa.int64()).to_pylist() + table.update(values={"time": value}) + assert table.to_arrow()["time"].cast(pa.int64()).to_pylist() == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York", "Asia/Tokyo"]) +@pytest.mark.parametrize( + "value", + [np.datetime64("2025-01-02T03:04:05.123456789", "ns"), np.datetime64("NaT", "ns")], +) +async def test_update_numpy_timestamp_matches_arrow_ingestion_async( + tmp_path, timezone, value +): + dtype = pa.timestamp("ns", tz=timezone) + ingested = pa.array(np.array([value]), type=dtype) + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("timestamps", pa.table({"time": ingested})) + expected = (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() + await table.update({"time": value}) + assert (await table.to_arrow())["time"].cast(pa.int64()).to_pylist() == expected From 57a6cce6819913adfca2852bb795a1ec0dd4641c Mon Sep 17 00:00:00 2001 From: LeonSGP <154585401+LeonSGP43@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:46:34 +0800 Subject: [PATCH 38/68] fix(python): omit excluded model fields from persistence (#4465) LanceDB serializes Pydantic rows with `model_dump()`, which omits `Field(exclude=True)`, but currently constructs an Arrow schema containing those same fields. Creating a local table therefore fails when an excluded field is non-nullable, including defaulted transient fields and nested models. Filter unconditionally excluded fields when converting model schemas, including nested structs. Skip excluded source/vector fields when deriving embedding configuration so persistent metadata does not refer to absent columns. Keep `field_names()` as model-field reflection. Tests exercise implicit and explicit schemas, synchronous and asynchronous create/add/reopen/readback, nested fields, defaults and default factories, retained fields, and local deterministic embeddings. Excluded values are not persisted; `to_pydantic` restores declared defaults. Conditional exclusion and readback of required excluded fields without defaults are outside this change. Implemented with assistance from gpt-6.1-sol; final diff reviewed locally. Local validation reuses the native extension matching official base d864e7198fe55a4e5fbe747fee574737fa102aab. The current-main Pydantic module is byte-identical to that base, but current-main native code has not been rebuilt. Slow/S3 tests, remote services, and other operating systems were not validated. Reproducible commands after preparing the documented development environment (relative `.venv` examples are equivalent commands; actual local runs use absolute read-only environments and fresh task-private working directories): ```sh PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests/test_pydantic_excluded_fields.py -q -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest python/python/tests -m 'not slow and not s3_test' -q -rs -p no:cacheprovider PYTHONPATH=python/python .venv/bin/python -m pytest --doctest-modules python/python/lancedb -q --tb=short (cd python && PYTHONPATH=python ../.venv/bin/pyright --pythonpath ../.venv/bin/python) .venv/bin/ruff check . .venv/bin/ruff format --check . .venv/bin/python -m pre_commit run --files python/python/lancedb/pydantic.py python/python/tests/test_pydantic_excluded_fields.py git diff --check d864e7198fe55a4e5fbe747fee574737fa102aab HEAD ``` Final frozen-source results: - EXCLUDE_TARGET_COMPLETE_ORIGINAL_RED: Python 3.11.14, exit 1, 1.900076000019908 seconds, {"failed"=>7}. - EXCLUDE_TARGET_EMBEDDING: Python 3.11.14, exit 0, 1.3686820000002626 seconds, {"passed"=>7}. - EXCLUDE_TARGET313: Python 3.13.11, exit 0, 1.7214069999754429 seconds, {"passed"=>7}. - EXCLUDE_FULL311: Python 3.11.14, exit 0, 217.56785399999353 seconds, {"passed"=>1531, "skipped"=>36, "deselected"=>68, "warnings"=>251}. - EXCLUDE_FULL313: Python 3.13.11, exit 0, 197.93847399999504 seconds, {"passed"=>1528, "skipped"=>37, "deselected"=>68, "warnings"=>251}. - EXCLUDE_DOCTEST313: Python 3.13.11, exit 0, 11.8951390000002 seconds, {"passed"=>71, "skipped"=>6, "warnings"=>2}. - EXCLUDE_PYRIGHT: Python 3.13.11, exit 0, 5.808195000019623 seconds, {"warnings"=>0, "errors"=>0}. - EXCLUDE_PRECOMMIT: Python 3.11.14, exit 0, 0.5624079999979585 seconds, {"passed"=>4, "skipped"=>2}. - EXCLUDE_RUFF_FINAL: Ruff 0.9.9, exit 0, 0.06722099997568876 seconds, {}. - EXCLUDE_FORMAT_FINAL: Ruff 0.9.9, exit 0, 0.033322000002954155 seconds, {"files_formatted"=>157}. - EXCLUDE_SYNTAX: Python 3.11.14, exit 0, 0.031629999983124435 seconds, {"files_compiled"=>2}. - EXCLUDE_DIFF_CHECK: Git, exit 0, 0.024460999993607402 seconds, {}. --- python/python/lancedb/pydantic.py | 9 +- .../tests/test_pydantic_excluded_fields.py | 128 ++++++++++++++++++ 2 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 python/python/tests/test_pydantic_excluded_fields.py diff --git a/python/python/lancedb/pydantic.py b/python/python/lancedb/pydantic.py index bb70c9a67..db8723a99 100644 --- a/python/python/lancedb/pydantic.py +++ b/python/python/lancedb/pydantic.py @@ -268,7 +268,9 @@ def _py_type_to_arrow_type(py_type: Type[Any], field: FieldInfo) -> pa.DataType: def _pydantic_model_to_fields(model: pydantic.BaseModel) -> List[pa.Field]: return [ - _pydantic_to_field(name, field) for name, field in model.model_fields.items() + _pydantic_to_field(name, field) + for name, field in model.model_fields.items() + if field.exclude is not True ] @@ -501,6 +503,9 @@ class LanceModel(pydantic.BaseModel): vec_and_function = [] for name, field_info in cls.safe_get_fields().items(): + if field_info.exclude is True: + continue + func = get_extras(field_info, "vector_column_for") if func is not None: vec_and_function.append([name, func]) @@ -508,6 +513,8 @@ class LanceModel(pydantic.BaseModel): configs = [] for vec, func in vec_and_function: for source, field_info in cls.safe_get_fields().items(): + if field_info.exclude is True: + continue src_func = get_extras(field_info, "source_column_for") if src_func is func: # note we can't use == here since the function is a pydantic diff --git a/python/python/tests/test_pydantic_excluded_fields.py b/python/python/tests/test_pydantic_excluded_fields.py new file mode 100644 index 000000000..5f1fd1adb --- /dev/null +++ b/python/python/tests/test_pydantic_excluded_fields.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import lancedb +import pytest +from pydantic import BaseModel, Field +from lancedb.pydantic import LanceModel + + +class Details(BaseModel): + label: str + transient: str = Field(default="nested default", exclude=True) + + +class Row(LanceModel): + id: int + details: Details + transient: str = Field(default="local default", exclude=True) + generated: list[str] = Field(default_factory=lambda: ["local"], exclude=True) + retained: str = Field(default="persisted", exclude=False) + + +def rows(): + return [ + Row( + id=i, + details=Details(label=str(i), transient="private"), + transient="private", + generated=["private"], + ) + for i in (1, 2) + ] + + +def assert_schema(schema): + assert schema.names == ["id", "details", "retained"] + assert [field.name for field in schema.field("details").type] == ["label"] + assert Row.field_names() == ["id", "details", "transient", "generated", "retained"] + + +def assert_restored(restored): + assert sorted(row.id for row in restored) == [1, 2] + for row in restored: + assert row.transient == "local default" + assert row.generated == ["local"] + assert row.details.transient == "nested default" + assert row.details.label == str(row.id) + assert row.retained == "persisted" + + +@pytest.mark.parametrize("explicit_schema", [False, True]) +def test_excluded_fields_round_trip(tmp_path, explicit_schema): + data = rows() + kwargs = {"schema": Row} if explicit_schema else {} + db = lancedb.connect(tmp_path) + table = db.create_table("rows", data=data[:1], **kwargs) + table.add(data[1:]) + assert_schema(table.schema) + reopened = lancedb.connect(tmp_path).open_table("rows") + assert_restored(reopened.search().to_pydantic(Row)) + assert data[0].transient == "private" + assert data[0].generated == ["private"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("explicit_schema", [False, True]) +async def test_excluded_fields_round_trip_async(tmp_path, explicit_schema): + data = rows() + kwargs = {"schema": Row} if explicit_schema else {} + db = await lancedb.connect_async(tmp_path) + table = await db.create_table("rows", data=data[:1], **kwargs) + await table.add(data[1:]) + assert_schema(await table.schema()) + reopened_db = await lancedb.connect_async(tmp_path) + reopened = await reopened_db.open_table("rows") + assert_restored(await reopened.query().to_pydantic(Row)) + assert data[0].details.transient == "private" + + +@pytest.mark.parametrize("excluded", ["transient", "source", "vector"]) +def test_excluded_fields_keep_embedding_metadata_consistent( + tmp_path, monkeypatch, excluded +): + from lancedb.embeddings.base import TextEmbeddingFunction + from lancedb.embeddings.registry import get_registry + from lancedb.pydantic import Vector + + calls = [] + + class LocalEmbedding(TextEmbeddingFunction): + def ndims(self): + return 2 + + def generate_embeddings(self, texts): + calls.extend(texts) + return [[1.0, 0.0] for _ in texts] + + registry = get_registry() + monkeypatch.setattr(registry, "_functions", registry._functions.copy()) + registry.register("test-excluded-fields")(LocalEmbedding) + function = LocalEmbedding.create(max_retries=0) + + class EmbeddedRow(LanceModel): + text: str = function.SourceField(default="alpha", exclude=excluded == "source") + vector: Vector(2) = function.VectorField( + default_factory=lambda: [1.0, 0.0], exclude=excluded == "vector" + ) + transient: str = Field(default="local", exclude=True) + + configs = EmbeddedRow.parse_embedding_functions() + assert len(configs) == (1 if excluded == "transient" else 0) + db = lancedb.connect(tmp_path) + data = [{"text": "alpha"}] if excluded == "transient" else [EmbeddedRow()] + table = db.create_table("embeddings", schema=EmbeddedRow, data=data) + assert "transient" not in table.schema.names + assert ("text" in table.schema.names) == (excluded != "source") + assert ("vector" in table.schema.names) == (excluded != "vector") + reopened = lancedb.connect(tmp_path).open_table("embeddings") + assert len(reopened.embedding_functions) == len(configs) + restored = reopened.search().to_pydantic(EmbeddedRow) + assert restored[0].text == "alpha" + assert restored[0].vector == [1.0, 0.0] + assert restored[0].transient == "local" + if excluded == "transient": + assert calls == ["alpha"] + assert reopened.search("alpha").limit(1).to_list()[0]["text"] == "alpha" + else: + assert calls == [] From 036f4ec53fac660b7703fe707dbcd7c5b405e020 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:13:31 +0800 Subject: [PATCH 39/68] fix(node): preserve options when namespace is omitted (#4479) The namespace overloads allow `undefined` to select the root namespace, but overload dispatch treats it as the legacy options argument and discards the trailing options. For example, `createTable(name, data, undefined, { mode: "overwrite" })` raises an already-exists error instead of overwriting. Preserve the trailing options for `createTable`, `createEmptyTable`, `tableNames`, and `listTables` when the namespace is omitted. Legacy options-object calls and explicit namespace arrays keep their existing behavior. Adds native-database regressions for overwrite, existOk, and both listing APIs, parameterized over `undefined` and `[]`. Validation on Windows / Node 24.19.0: - Before the fix: all five omitted-namespace cases failed; all five explicit-root controls passed. - Current TypeScript source with official, integrity-verified 0.41.0-beta.0 native binaries: the entire connection test file passed, 47/47 (transpile-only). - The same fix replayed onto the matching 0.41.0-beta.0 source: TypeScript compilation, strict ts-jest connection tests (47/47), and TypeDoc passed. - Biome check/format on both changed files and diff whitespace checks passed. No Rust changes. A complete current-main native rebuild remains for CI: released declarations predate the unrelated lazy-iterator APIs, so current-main typechecking against those declarations reports six pre-existing iterator signature errors. Signed-off-by: Xucheng Zhou Co-authored-by: Xuanwo --- nodejs/__test__/connection.test.ts | 60 ++++++++++++++++++++++++++++++ nodejs/lancedb/connection.ts | 8 ++-- 2 files changed, 64 insertions(+), 4 deletions(-) diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 1167ae34a..355b86ccb 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -117,6 +117,66 @@ describe("given a connection", () => { await expect(tbl.countRows()).resolves.toBe(1); }); + describe.each([undefined, []])("root namespace %p", (namespacePath) => { + it("should preserve createTable options after the namespace", async () => { + await db.createTable("test", [{ id: 1 }, { id: 2 }]); + + const table = await db.createTable("test", [{ id: 3 }], namespacePath, { + mode: "overwrite", + }); + + const rows = await table.query().toArray(); + expect(rows.map((row) => row.id)).toEqual([3]); + }); + + it("should preserve createEmptyTable overwrite after the namespace", async () => { + await db.createTable("test", [{ id: 1 }]); + const schema = new Schema([new Field("id", new Float64(), true)]); + + const table = await db.createEmptyTable("test", schema, namespacePath, { + mode: "overwrite", + }); + + await expect(table.countRows()).resolves.toBe(0); + }); + + it("should preserve createEmptyTable existOk after the namespace", async () => { + await db.createTable("test", [{ id: 1 }]); + const schema = new Schema([new Field("id", new Float64(), true)]); + + const table = await db.createEmptyTable("test", schema, namespacePath, { + existOk: true, + }); + + await expect(table.countRows()).resolves.toBe(1); + }); + + it("should preserve tableNames options after the namespace", async () => { + for (const name of ["a", "b", "c"]) { + await db.createTable(name, [{ id: 1 }]); + } + + await expect( + db.tableNames(namespacePath, { limit: 1, startAfter: "a" }), + ).resolves.toEqual(["b"]); + }); + + it("should preserve listTables options after the namespace", async () => { + for (const name of ["a", "b", "c"]) { + await db.createTable(name, [{ id: 1 }]); + } + + const first = await db.listTables(namespacePath, { limit: 1 }); + expect(first.tables).toEqual(["a"]); + expect(first.pageToken).toBeDefined(); + + const second = await db.listTables(namespacePath, { + limit: 1, + pageToken: first.pageToken, + }); + expect(second.tables).toEqual(["b"]); + }); + }); it("should keep the existing table when mode is exist_ok", async () => { const created = await db.createTable("test", [{ id: 1 }, { id: 2 }], { mode: "exist_ok", diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 4ed511b86..f866d1ecd 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -720,7 +720,7 @@ export class LocalConnection extends Connection { } else { // First argument is options object (backwards compatibility) namespacePath = undefined; - tableNamesOptions = namespacePathOrOptions; + tableNamesOptions = namespacePathOrOptions ?? options; } return this.inner.tableNames( @@ -831,7 +831,7 @@ export class LocalConnection extends Connection { : undefined; const listTablesOptions = Array.isArray(namespacePathOrOptions) ? options - : namespacePathOrOptions; + : (namespacePathOrOptions ?? options); return this.inner.listTables( namespacePath ?? [], @@ -942,7 +942,7 @@ export class LocalConnection extends Connection { } else { // Third argument is options object (backwards compatibility) namespacePath = undefined; - createOptions = namespacePathOrOptions; + createOptions = namespacePathOrOptions ?? options; } return this._createTableImpl(name, data, namespacePath, createOptions); @@ -989,7 +989,7 @@ export class LocalConnection extends Connection { } else { // Third argument is options object (backwards compatibility) namespacePath = undefined; - createOptions = namespacePathOrOptions; + createOptions = namespacePathOrOptions ?? options; } let mode: string = createOptions?.mode ?? "create"; From 13d54db59411029f6c80755b7453090e17beb417 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:13:46 +0800 Subject: [PATCH 40/68] fix(node): preserve caller-owned storage options (#4480) `getStorageOptions` writes typed creation options back into the caller's `storageOptions` object. This makes frozen options throw and causes overrides to leak into later calls that reuse the storage map. For example, using `{ storageOptions: shared, enableV2ManifestPaths: true }` changes `shared.newTableEnableV2ManifestPaths` to `"true"`. A subsequent creation with just `{ storageOptions: shared }` then uses V2 manifest paths even if `shared` originally requested `"false"`. Merge the typed overrides into a local map instead. Their existing precedence, key normalization, and undefined/default behavior remain unchanged. This affects both `createTable` and `createEmptyTable`; no native/Rust changes. Adds six native-database regressions covering frozen creation options, frozen nested storage options with conflicting defaults, and cross-call reuse for both APIs. Validation (Linux x64 / Node 24.19.0 / official 0.41.0-beta.0 native library): - Before: all six new regressions fail; an existing namespace storage-options control passes. - After: the full connection test file passes, 43/43 (current TypeScript source, transpile-only). - Same fix on the matching 0.41.0-beta.0 source: TypeScript compilation, strict ts-jest connection tests (43/43), and TypeDoc with warnings-as-errors pass. - Biome on changed files and whitespace checks pass. A current-main native rebuild was not run. Released native declarations predate six unrelated lazy-iterator signatures; current-main typechecking reports the same six errors before and after this fix (only downstream line coordinates shift). ### Connection entry-point follow-up The same ownership problem occurs in `connect(uri, options)` and `connectNamespace`: normalization assigns a new storage map into the caller's options. Copy the options before normalization instead. The `connect(options)` overload already copies its input and is retained as a control. Nine added native-database cases cover frozen outer options (with/without storage options) and repeated reuse, including table creation and row-count checks. Before this entry-point fix, the URI/namespace cases fail (6), while the object-overload controls pass (3). Afterward the complete connection file passes 52/52 on current TypeScript with the official 0.41.0-beta.0 native library. Applying the entry-point fix to matching-release source also passes strict TypeScript, strict ts-jest (52/52), and TypeDoc. Biome passes on both newly changed files. The same six current-main/released-native lazy-iterator declaration errors remain unchanged. No Rust/native rebuild or GPU claim is added. --------- Signed-off-by: Xucheng Zhou --- nodejs/__test__/connection.test.ts | 144 +++++++++++++++++++++++++++++ nodejs/lancedb/connection.ts | 23 ++--- nodejs/lancedb/index.ts | 17 ++-- 3 files changed, 164 insertions(+), 20 deletions(-) diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index 355b86ccb..e32e7e626 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -36,6 +36,87 @@ describe("when connecting", () => { }); }); +describe.each([ + [ + "URI overload", + (uri: string, options: { storageOptions?: Record }) => + connect(uri, options), + ], + [ + "object overload", + async ( + uri: string, + options: { storageOptions?: Record }, + ) => { + const input = { uri, ...options }; + if (Object.isFrozen(options)) { + Object.freeze(input); + } + const db = await connect(input); + expect(input).toEqual({ uri, ...options }); + expect(input.storageOptions).toBe(options.storageOptions); + return db; + }, + ], + [ + "namespace", + (uri: string, options: { storageOptions?: Record }) => + connectNamespace("dir", { root: uri }, options), + ], +])("%s connection options", (_name, open) => { + let tmpDir: tmp.DirResult; + beforeEach(() => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + }); + afterEach(() => tmpDir.removeCallback()); + + it("should accept frozen options without storage options", async () => { + const options = Object.freeze({}); + const db = await open(tmpDir.name, options); + try { + await db.createTable("frozen", [{ id: 1 }]); + await expect(db.tableNames()).resolves.toContain("frozen"); + expect(options).toEqual({}); + } finally { + await db.close(); + } + }); + + it("should accept frozen options with storage options", async () => { + const storageOptions = Object.freeze({ + newTableDataStorageVersion: "stable", + }); + const options = Object.freeze({ storageOptions }); + const db = await open(tmpDir.name, options); + try { + const table = await db.createTable("frozen", [{ id: 1 }]); + await expect(table.countRows()).resolves.toBe(1); + expect(options.storageOptions).toBe(storageOptions); + expect(storageOptions).toEqual({ newTableDataStorageVersion: "stable" }); + } finally { + await db.close(); + } + }); + + it("should preserve reused options and their storage options", async () => { + const storageOptions = { newTableDataStorageVersion: "stable" }; + const options = { storageOptions }; + for (let i = 0; i < 2; i++) { + const db = await open(tmpDir.name, options); + try { + const table = await db.createTable(`reused_${i}`, [{ id: i }]); + await expect(table.countRows()).resolves.toBe(1); + expect(options.storageOptions).toBe(storageOptions); + expect(storageOptions).toEqual({ + newTableDataStorageVersion: "stable", + }); + } finally { + await db.close(); + } + } + }); +}); + describe("given a connection", () => { let tmpDir: tmp.DirResult; let db: Connection; @@ -367,6 +448,69 @@ describe("given a connection", () => { }); }); + describe.each(["createTable", "createEmptyTable"] as const)( + "%s storage options", + (method) => { + const schema = new Schema([new Field("id", new Float64(), true)]); + + it("should accept frozen creation options without storage options", async () => { + const options = Object.freeze({ + dataStorageVersion: "stable", + enableV2ManifestPaths: true, + }); + const table = + method === "createTable" + ? await db.createTable("frozen_options", [{ id: 1 }], options) + : await db.createEmptyTable("frozen_options", schema, options); + expect(await (table as LocalTable).usesV2ManifestPaths()).toBe(true); + expect(options).not.toHaveProperty("storageOptions"); + }); + + it("should accept frozen storage options and preserve override precedence", async () => { + const storageOptions = Object.freeze({ + newTableDataStorageVersion: "legacy", + newTableEnableV2ManifestPaths: "false", + }); + const options = Object.freeze({ + storageOptions, + dataStorageVersion: "stable", + enableV2ManifestPaths: true, + }); + const table = + method === "createTable" + ? await db.createTable("frozen_storage", [{ id: 1 }], options) + : await db.createEmptyTable("frozen_storage", schema, options); + expect(await (table as LocalTable).usesV2ManifestPaths()).toBe(true); + expect(options.storageOptions).toBe(storageOptions); + expect(storageOptions).toEqual({ + newTableDataStorageVersion: "legacy", + newTableEnableV2ManifestPaths: "false", + }); + }); + + it("should not retain a previous override when reusing storage options", async () => { + const storageOptions = { newTableEnableV2ManifestPaths: "false" }; + const create = async (name: string, enableV2ManifestPaths?: boolean) => + method === "createTable" + ? db.createTable(name, [{ id: 1 }], { + storageOptions, + enableV2ManifestPaths, + }) + : db.createEmptyTable(name, schema, { + storageOptions, + enableV2ManifestPaths, + }); + const first = await create("override", true); + expect(await (first as LocalTable).usesV2ManifestPaths()).toBe(true); + const second = await create("without_override"); + expect(await (second as LocalTable).usesV2ManifestPaths()).toBe(false); + expect(storageOptions).toEqual({ + newTableEnableV2ManifestPaths: "false", + }); + }); + }, + ); + it("should be able to migrate tables to the V2 manifest paths", async () => { const db = await connect(tmpDir.name); const table = (await db.createEmptyTable( diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index f866d1ecd..44fde96f1 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -893,23 +893,24 @@ export class LocalConnection extends Connection { private getStorageOptions( options?: Partial, ): Record | undefined { + let storageOptions = options?.storageOptions; if (options?.dataStorageVersion !== undefined) { - if (options.storageOptions === undefined) { - options.storageOptions = {}; - } - options.storageOptions["newTableDataStorageVersion"] = - options.dataStorageVersion; + storageOptions = { + ...storageOptions, + newTableDataStorageVersion: options.dataStorageVersion, + }; } if (options?.enableV2ManifestPaths !== undefined) { - if (options.storageOptions === undefined) { - options.storageOptions = {}; - } - options.storageOptions["newTableEnableV2ManifestPaths"] = - options.enableV2ManifestPaths ? "true" : "false"; + storageOptions = { + ...storageOptions, + newTableEnableV2ManifestPaths: options.enableV2ManifestPaths + ? "true" + : "false", + }; } - return cleanseStorageOptions(options?.storageOptions); + return cleanseStorageOptions(storageOptions); } async createTable( diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 1057e2d53..4ad8b2a93 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -420,10 +420,10 @@ export async function connect( throw new Error("uri is required"); } - finalOptions = (finalOptions as ConnectionOptions) ?? {}; - (finalOptions).storageOptions = cleanseStorageOptions( - (finalOptions).storageOptions, - ); + finalOptions = { + ...finalOptions, + storageOptions: cleanseStorageOptions(finalOptions.storageOptions), + }; // Create native header provider if one was provided let nativeProvider: NativeJsHeaderProvider | undefined; @@ -630,11 +630,10 @@ export async function connectNamespace( properties = configOrProperties as Record; } - const finalOptions: ConnectNamespaceOptions = (options ?? - {}) as ConnectNamespaceOptions; - finalOptions.storageOptions = cleanseStorageOptions( - finalOptions.storageOptions, - ); + const finalOptions: ConnectNamespaceOptions = { + ...options, + storageOptions: cleanseStorageOptions(options?.storageOptions), + }; const nativeConn = await LanceDbConnection.newWithNamespace( implName, From e256b3c2cfc49f45c74d4ac5e31a2e60b6a40efa Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:03 +0800 Subject: [PATCH 41/68] fix(python): honor namespace table storage options (#4483) Follow-up on current main `7b3ba4c4` (merge commit `0b0475c1`): resolved the test import conflict, preserving both upstream imports and the original `warnings` import. The final PR diff is still limited to the same storage-option helper/call sites and regression class. AST comparisons preserve every upstream test definition and the original regression class/helper; both changed files compile and pass Ruff 0.15.1 check/format. Fresh locked-environment runtime validation was attempted with `uv run --locked --extra tests --extra dev pytest python/tests/test_namespace.py -q`, but dependency setup failed downloading the locked `pylance==14.0.0b12` wheel due a network connection error. No current-main native test result is claimed. Repository-wide Ruff reports 49 findings and 11 formatting differences in unchanged upstream docs/web-tests files; these were left outside this PR. CI approval and a human review are still pending. --- Original implementation and historical validation (before the main merge): Namespace-backed `create_table` accepts `data_storage_version` and `enable_v2_manifest_paths`, but neither sync nor async creation forwards them. The requested format/manifest settings are silently ignored. Normalize these deprecated arguments into the existing `new_table_*` storage options, matching ordinary connections. Explicit arguments take precedence; the caller's dictionary is not mutated. No Rust or native interface changes. Regression tests inspect persisted tables, including sync/async creation, populated/empty tables, stable/legacy formats, both manifest settings and explicit-option precedence. Validation on Linux/Python 3.12, based on current main `d3390a9`: - Rebuilt the exact-main native extension and reused that identical extension in separate baseline/candidate environments; only the candidate Python namespace wrapper differs. Dependencies came from the unchanged `uv.lock`. - Same complete `test_namespace.py`: baseline 16 regression failures and 39 passing controls; candidate **55 passed**, no skips/errors. - Ruff check/format and diff whitespace checks passed. This validates the local directory-namespace path, not a live Cloud service or the complete repository CI suite. Signed-off-by: Xucheng Zhou --- python/python/lancedb/namespace.py | 32 +++++++ python/python/tests/test_namespace.py | 120 ++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index f8675ebb5..a5d88eb1f 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -11,6 +11,7 @@ through a namespace abstraction. from __future__ import annotations import sys +import warnings from typing import Any, Dict, Iterable, List, Optional, Union from uuid import UUID @@ -82,6 +83,31 @@ from ._lancedb import Session _MAX_QUERY_K = 2**31 - 1 +def _create_table_storage_options( + storage_options: Optional[Dict[str, str]], + data_storage_version: Optional[str], + enable_v2_manifest_paths: Optional[bool], +) -> Optional[Dict[str, str]]: + if data_storage_version is None and enable_v2_manifest_paths is None: + return storage_options + options = dict(storage_options or {}) + for name, value in ( + ("data_storage_version", data_storage_version), + ("enable_v2_manifest_paths", enable_v2_manifest_paths), + ): + if value is not None: + warnings.warn( + f"setting {name} directly on create_table is deprecated. " + "Use storage_options instead.", + DeprecationWarning, + stacklevel=3, + ) + options[f"new_table_{name}"] = ( + str(value).lower() if isinstance(value, bool) else value + ) + return options + + def _query_to_namespace_request( table_id: List[str], query: "Query", @@ -587,6 +613,9 @@ class LanceNamespaceDBConnection(DBConnection): if mode.lower() not in ["create", "overwrite"]: raise ValueError("mode must be either 'create' or 'overwrite'") validate_table_name(name) + storage_options = _create_table_storage_options( + storage_options, data_storage_version, enable_v2_manifest_paths + ) async_table = LOOP.run( self._inner.create_table( name, @@ -1230,6 +1259,9 @@ class AsyncLanceNamespaceDBConnection: if mode.lower() not in ["create", "overwrite"]: raise ValueError("mode must be either 'create' or 'overwrite'") validate_table_name(name) + storage_options = _create_table_storage_options( + storage_options, data_storage_version, enable_v2_manifest_paths + ) table = await self._inner.create_table( name, data, diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index 520a6992d..b55da4d57 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -6,6 +6,7 @@ import tempfile import shutil import importlib +import warnings import http.server import json import threading @@ -25,6 +26,125 @@ PUSHDOWN_DATA = pa.table( ) +class TestNamespaceCreateTableOptions: + @pytest.mark.parametrize("version", ["legacy", "stable"]) + @pytest.mark.parametrize("empty", [False, True]) + def test_deprecated_storage_version(self, tmp_path, version, empty): + db = lancedb.connect_namespace("dir", {"root": str(tmp_path / "namespace")}) + data = None if empty else [{"id": 1}] + schema = pa.schema([("id", pa.int64())]) + other_version = "stable" if version == "legacy" else "legacy" + options = {"new_table_data_storage_version": other_version} + control = lancedb.connect(tmp_path / "control").create_table( + "test", + data=data, + schema=schema, + storage_options={"new_table_data_storage_version": version}, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = db.create_table( + "test", + data=data, + schema=schema, + storage_options=options, + data_storage_version=version, + ) + + assert ( + table.to_lance().data_storage_version + == control.to_lance().data_storage_version + ) + assert table.to_arrow().to_pylist() == ([] if empty else data) + assert options == {"new_table_data_storage_version": other_version} + assert any("data_storage_version" in str(w.message) for w in caught) + + @pytest.mark.asyncio + @pytest.mark.parametrize("version", ["legacy", "stable"]) + @pytest.mark.parametrize("empty", [False, True]) + async def test_deprecated_storage_version_async(self, tmp_path, version, empty): + root = str(tmp_path / "namespace") + db = lancedb.connect_namespace_async("dir", {"root": root}) + data = None if empty else [{"id": 1}] + schema = pa.schema([("id", pa.int64())]) + other_version = "stable" if version == "legacy" else "legacy" + options = {"new_table_data_storage_version": other_version} + control = lancedb.connect(tmp_path / "control").create_table( + "test", + data=data, + schema=schema, + storage_options={"new_table_data_storage_version": version}, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = await db.create_table( + "test", + data=data, + schema=schema, + storage_options=options, + data_storage_version=version, + ) + + reopened = lancedb.connect_namespace("dir", {"root": root}).open_table("test") + assert ( + reopened.to_lance().data_storage_version + == control.to_lance().data_storage_version + ) + assert (await table.to_arrow()).to_pylist() == ([] if empty else data) + assert options == {"new_table_data_storage_version": other_version} + assert any("data_storage_version" in str(w.message) for w in caught) + + @pytest.mark.parametrize("enabled", [False, True]) + @pytest.mark.parametrize("empty", [False, True]) + def test_deprecated_manifest_paths(self, tmp_path, enabled, empty): + db = lancedb.connect_namespace("dir", {"root": str(tmp_path)}) + data = None if empty else [{"id": 1}] + options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = db.create_table( + "test", + data=data, + schema=pa.schema([("id", pa.int64())]), + storage_options=options, + enable_v2_manifest_paths=enabled, + ) + + assert table.uses_v2_manifest_paths() == enabled + assert table.to_arrow().to_pylist() == ([] if empty else data) + assert db.open_table("test").uses_v2_manifest_paths() == enabled + assert options == { + "new_table_enable_v2_manifest_paths": str(not enabled).lower() + } + assert any("enable_v2_manifest_paths" in str(w.message) for w in caught) + + @pytest.mark.asyncio + @pytest.mark.parametrize("enabled", [False, True]) + @pytest.mark.parametrize("empty", [False, True]) + async def test_deprecated_manifest_paths_async(self, tmp_path, enabled, empty): + db = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)}) + data = None if empty else [{"id": 1}] + options = {"new_table_enable_v2_manifest_paths": str(not enabled).lower()} + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + table = await db.create_table( + "test", + data=data, + schema=pa.schema([("id", pa.int64())]), + storage_options=options, + enable_v2_manifest_paths=enabled, + ) + + assert await table.uses_v2_manifest_paths() == enabled + assert (await table.to_arrow()).to_pylist() == ([] if empty else data) + reopened = await db.open_table("test") + assert await reopened.uses_v2_manifest_paths() == enabled + assert options == { + "new_table_enable_v2_manifest_paths": str(not enabled).lower() + } + assert any("enable_v2_manifest_paths" in str(w.message) for w in caught) + + def _ipc_file(table: pa.Table = PUSHDOWN_DATA) -> bytes: sink = pa.BufferOutputStream() with pa.ipc.new_file(sink, table.schema) as writer: From b205e7c628c327a10f778a8d6db217c63da96e97 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:17 +0800 Subject: [PATCH 42/68] fix(node): prevent table listing limits from wrapping (#4492) ## Summary Prevent large JavaScript table-listing limits from wrapping during the native `u32` conversion. With two tables, `limit: 2 ** 32` currently returns no tables and `2 ** 32 + 1` returns only one. Validate nonnegative safe integers using the existing helper and clamp valid limits above `u32::MAX` before calling the native binding. Apply the same normalization to `tableNames` and `listTables`, including namespace overloads. Undefined, zero, small limits and `u32::MAX` keep their existing behavior. ## Validation Built the real Node native extension from main `0aa842b17f13d1f8ffd2851bc57c222bc5a376ba` with the repository's `pnpm build`, using locked dependencies and reusable build caches. No mocked/released native extension or disabled TypeScript diagnostics. With identical regression tests and that native library in fresh Jest processes: - Baseline: 40 passed, 8 failed; candidate: **48 passed**, no failures. - Both arms: `pnpm run tsc` and `pnpm run docs` passed. - Changed-file Biome check passed. Tests cover wrapping limits, invalid numbers, zero/small/u32 limits, defaults and both overload forms. Validation is the focused connection suite, not the complete repository suite. No performance claim. AI-assisted implementation and verification; the patch and exact tested source bytes were reviewed before publication. Signed-off-by: aden Co-authored-by: aden --- nodejs/__test__/connection.test.ts | 50 ++++++++++++++++++++++++++++++ nodejs/lancedb/connection.ts | 11 +++++-- 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index e32e7e626..bbd0c3c52 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -336,6 +336,56 @@ describe("given a connection", () => { expect(seen).toEqual(created); }); + it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( + "should not wrap a large table listing limit %s", + async (limit) => { + await db.createTable("a", [{ id: 1 }]); + await db.createTable("b", [{ id: 2 }]); + for (const names of [ + db.tableNames({ limit }), + db.tableNames([], { limit }), + ]) { + await expect(names).resolves.toEqual(["a", "b"]); + } + for (const page of [ + db.listTables({ limit }), + db.listTables([], { limit }), + ]) { + await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + } + }, + ); + + it.each([ + -1, + 1.5, + Number.NaN, + Number.POSITIVE_INFINITY, + Number.NEGATIVE_INFINITY, + Number.MAX_SAFE_INTEGER + 1, + ])("should reject an invalid table listing limit %s", async (limit) => { + await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); + await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); + await expect(db.listTables({ limit })).rejects.toThrow(/limit/); + await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + }); + + it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { + await db.createTable("a", [{ id: 1 }]); + await db.createTable("b", [{ id: 2 }]); + for (const [limit, names] of [ + [0, []], + [1, ["a"]], + [2 ** 32 - 1, ["a", "b"]], + ] as const) { + await expect(db.tableNames({ limit })).resolves.toEqual(names); + const page = await db.listTables({ limit }); + expect(page.tables).toEqual(names); + } + await expect(db.tableNames()).resolves.toEqual(["a", "b"]); + await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + }); + it("should list all names and use a default page size of 100", async () => { const names = Array.from( { length: 130 }, diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 44fde96f1..c93414fab 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -726,7 +726,7 @@ export class LocalConnection extends Connection { return this.inner.tableNames( namespacePath ?? [], tableNamesOptions?.startAfter, - tableNamesOptions?.limit, + normalizeTableListingLimit(tableNamesOptions?.limit), ); } @@ -836,7 +836,7 @@ export class LocalConnection extends Connection { return this.inner.listTables( namespacePath ?? [], listTablesOptions?.pageToken, - listTablesOptions?.limit, + normalizeTableListingLimit(listTablesOptions?.limit), ); } @@ -1107,6 +1107,13 @@ export class LocalConnection extends Connection { } } +function normalizeTableListingLimit(limit?: number): number | undefined { + validateNonNegativeInteger(limit, "limit"); + // The native binding takes u32: larger JS numbers otherwise wrap to zero or + // a smaller page size. A listing limit is an upper bound, so saturate it. + return limit === undefined ? undefined : Math.min(limit, 0xffffffff); +} + /** * Takes storage options and makes all the keys snake case. */ From 3d883af674924a7e7266da19d9ce1d774298f74c Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:31 +0800 Subject: [PATCH 43/68] fix(python): propagate invalid index wait_timeout errors (#4495) Invalid `wait_timeout` values currently reach `Duration::extract(...).unwrap()` in both native index-creation bindings. A normal Python argument error becomes `pyo3_runtime.PanicException`, which is not caught by `except Exception`. Propagate the extraction error with `map(...).transpose()?` before constructing the index operation. Valid `None` and `timedelta` behavior is unchanged. Tests cover both `create_index` and `create_index_async`, three invalid types, valid `None`/`timedelta`, and no index created after rejection. Validation against main `0aa842b`, Python 3.12, uv-locked dependencies, freshly compiled PyO3 extension for each arm: - Unmodified source: 6 new regressions fail with `PanicException`; 6 controls pass. - Fixed source: 12/12 pass; broader index-creation regression selection: 18/18 pass. - `cargo fmt --all --check`, Ruff check/format, and diff check pass. The broader selection includes the 12 cases above. No cloud, full-suite, or performance claim. Signed-off-by: Xucheng Zhou --- python/python/tests/test_table.py | 19 +++++++++++++++++++ python/src/table.rs | 8 ++++++-- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index b2d499933..238ee5330 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1633,6 +1633,25 @@ def test_create_index_async_returns_done_job(mem_db: DBConnection): job.cancel() +@pytest.mark.parametrize("method", ["create_index", "create_index_async"]) +@pytest.mark.parametrize("wait_timeout", ["invalid", 1, []]) +def test_create_index_invalid_wait_timeout(mem_db: DBConnection, method, wait_timeout): + table = mem_db.create_table("invalid_timeout", [{"id": 1}]) + with pytest.raises(TypeError): + getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout) + assert table.list_indices() == [] + + +@pytest.mark.parametrize("method", ["create_index", "create_index_async"]) +@pytest.mark.parametrize("wait_timeout", [None, timedelta(seconds=1)]) +def test_create_index_valid_wait_timeout(mem_db: DBConnection, method, wait_timeout): + table = mem_db.create_table("valid_timeout", [{"id": 1}]) + result = getattr(table, method)("id", config=BTree(), wait_timeout=wait_timeout) + if result is not None: + result.wait() + assert len(table.list_indices()) == 1 + + @patch("lancedb.table.AsyncTable.create_index") def test_create_index_method(mock_create_index, mem_db: DBConnection): table = mem_db.create_table( diff --git a/python/src/table.rs b/python/src/table.rs index 27ed8be8b..e77c58578 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1134,7 +1134,9 @@ impl Table { train: Option, ) -> PyResult> { let index = extract_index_params(&index)?; - let timeout = wait_timeout.map(|t| t.extract::().unwrap()); + let timeout = wait_timeout + .map(|t| t.extract::()) + .transpose()?; let mut op = self_ .inner_ref()? .create_index_with_timeout(&[column], index, timeout); @@ -1165,7 +1167,9 @@ impl Table { train: Option, ) -> PyResult> { let index = extract_index_params(&index)?; - let timeout = wait_timeout.map(|t| t.extract::().unwrap()); + let timeout = wait_timeout + .map(|t| t.extract::()) + .transpose()?; let mut op = self_ .inner_ref()? .create_index_with_timeout(&[column], index, timeout); From 93f054bcc5b02927720cdcd52fd7b1cc13c9a731 Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:14:45 +0800 Subject: [PATCH 44/68] fix(rust): finish query streams after a zero-duration timeout (#4506) A positive query timeout emits one error and then ends, but `Duration::ZERO` leaves `TimeoutStream` in `NotStarted`, returning the same error on every subsequent poll. Transition the zero-duration path to `Completed` before returning its error. The initial error is unchanged. Extend both timeout tests to check that subsequent polls return `None`; the normal-completion test is unchanged. This affects callers that continue polling after an error, not `try_collect`, which already stops at the first error. Validation: - `cargo fmt --all` and `cargo fmt --all --check`: pass. - Mechanically extracted the unchanged implementation and three upstream timeout tests into a native Rust test binary using real cached Arrow 58.4.0, DataFusion 54.1.0, futures 0.3.32 and Tokio 1.53.1 libraries (matching the current lock). With the new assertions, baseline: 2 pass / 1 fail; fix: 3 pass. - Full-crate `cargo test --locked --offline --features remote -p lancedb --lib test_timeout_stream` could not start because this environment lacks the newly locked Lance `v14.0.0-beta.10` checkout. The component comparison is not a claim that the full Cargo suite passed; CI still needs to validate integration. Signed-off-by: Xucheng Zhou --- rust/lancedb/src/utils/mod.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/rust/lancedb/src/utils/mod.rs b/rust/lancedb/src/utils/mod.rs index 20f209f2e..ebf3f697b 100644 --- a/rust/lancedb/src/utils/mod.rs +++ b/rust/lancedb/src/utils/mod.rs @@ -686,7 +686,9 @@ impl Stream for TimeoutStream { match &mut self.state { TimeoutState::NotStarted { timeout } => { if timeout.is_zero() { - return std::task::Poll::Ready(Some(Err(Self::timeout_error(timeout)))); + let err = Self::timeout_error(timeout); + self.state = TimeoutState::Completed; + return std::task::Poll::Ready(Some(Err(err))); } let deadline = Box::pin(tokio::time::sleep(*timeout)); self.state = TimeoutState::Started { @@ -1159,6 +1161,8 @@ mod tests { .to_string() .contains("Query timeout") ); + assert!(timeout_stream.next().await.is_none()); + assert!(timeout_stream.next().await.is_none()); } #[tokio::test] @@ -1178,6 +1182,8 @@ mod tests { let result = timeout_stream.next().await.unwrap(); assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Query timeout")); + assert!(timeout_stream.next().await.is_none()); + assert!(timeout_stream.next().await.is_none()); } #[tokio::test] From d826b6115a0ca30998d9aa1bfb83d01e8f65bf4b Mon Sep 17 00:00:00 2001 From: Aviral Saxena Date: Fri, 23 Oct 2026 00:45:01 +0530 Subject: [PATCH 45/68] fix(rust): explain client-side model files when remote tokenize fails (#4519) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses option 2 of #4440. tokenize() builds the FTS analyzer on the client. For jieba/* and lindera/* that needs model files under LANCE_LANGUAGE_MODEL_HOME, and when they’re missing you only got lance’s raw Invalid directory path: ... error, even though server-side search works. Now, if a jieba/* or lindera/* tokenizer can’t be built, the error keeps the original cause and adds that tokenize() runs on the client, which directory the files need to be in, and that for remote tables they must match the server’s. I also updated the tokenize docs (Rust, Python, TS, generated JS docs) to say a mismatch can’t be detected, so tokens may differ from the server’s index. Behavior when the files exist is unchanged. Detecting a mismatch needs a server-side tokenize endpoint (option 1), so I left it out. The hint isn’t limited to remote tables, since Table::as_native is slated for removal and the wording fits local tables too. Tests: Rust test with mocked jieba/default index_details and an empty model dir, one checking other tokenizer errors don’t get the hint, and Python sync/async tests. --- docs/src/js/classes/Table.md | 6 +- nodejs/lancedb/table.ts | 6 +- python/python/lancedb/remote/table.py | 8 +- python/python/lancedb/table.py | 18 +++- python/python/tests/test_remote_db.py | 73 +++++++++++++++ rust/lancedb/src/remote/table.rs | 125 ++++++++++++++++++++++++++ rust/lancedb/src/table.rs | 34 +++++-- 7 files changed, 257 insertions(+), 13 deletions(-) diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index b41ade302..7cf87c405 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -1298,7 +1298,11 @@ Specify exactly one of `column` or `indexName`. Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in the client process from index metadata. For remote tables, this means the -same tokenizer model files must also exist locally. +same tokenizer model files must also exist locally (under +`LANCE_LANGUAGE_MODEL_HOME`) and must match the files the server uses. A +mismatch cannot be detected: the returned tokens may differ from what the +server's index produces. If the model files are missing, an error naming +the tokenizer is raised. #### Parameters diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 0c79e6a7e..4ae1aefb8 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -989,7 +989,11 @@ export abstract class Table { * * Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in * the client process from index metadata. For remote tables, this means the - * same tokenizer model files must also exist locally. + * same tokenizer model files must also exist locally (under + * `LANCE_LANGUAGE_MODEL_HOME`) and must match the files the server uses. A + * mismatch cannot be detected: the returned tokens may differ from what the + * server's index produces. If the model files are missing, an error naming + * the tokenizer is raised. */ abstract tokenize( query: string, diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 96b988911..86e8fb8a0 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -279,8 +279,12 @@ class RemoteTable(Table): """Tokenize a query using the tokenizer configured on an FTS index. Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are - rebuilt in the client process from index metadata, so the same tokenizer - model files must exist locally. + rebuilt in the client process from index metadata. For remote tables, + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ return LOOP.run( self._table.tokenize(query, column=column, index_name=index_name) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index f9cfac6e1..124fa34b5 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2558,7 +2558,11 @@ class Table(ABC): Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are rebuilt in the client process from index metadata. For remote tables, - this means the same tokenizer model files must also exist locally. + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ @abstractmethod @@ -4813,7 +4817,11 @@ class LanceTable(Table): Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are rebuilt in the client process from index metadata. For remote tables, - this means the same tokenizer model files must also exist locally. + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ return LOOP.run( self._table.tokenize(query, column=column, index_name=index_name) @@ -7505,7 +7513,11 @@ class AsyncTable: Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are rebuilt in the client process from index metadata. For remote tables, - this means the same tokenizer model files must also exist locally. + this means the same tokenizer model files must also exist locally + (under ``LANCE_LANGUAGE_MODEL_HOME``) and must match the files the + server uses. A mismatch cannot be detected: the returned tokens may + differ from what the server's index produces. If the model files are + missing, an error naming the tokenizer is raised. """ return await self._inner.tokenize(query, column=column, index_name=index_name) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 80bd10289..a45e1aaef 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1793,6 +1793,79 @@ def test_table_wait_for_index_timeout(): table.wait_for_index(["id_idx"], timedelta(seconds=1)) +def _jieba_fts_handler(request): + """Serve a table with a ``jieba/default`` FTS index and no other state.""" + if request.path == "/v1/table/test/describe/": + body = { + "version": 1, + "schema": { + "fields": [ + {"name": "text", "type": {"type": "string"}, "nullable": True} + ] + }, + } + elif request.path == "/v1/table/test/index/list/": + index_details = { + "base_tokenizer": "jieba/default", + "language": "English", + "with_position": False, + "max_token_length": 40, + "lower_case": True, + "stem": False, + "remove_stop_words": False, + "ascii_folding": False, + } + body = { + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": json.dumps(index_details), + } + ] + } + else: + request.send_response(404) + request.end_headers() + return + request.send_response(200) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(body).encode()) + + +_MODEL_FILES_ERROR = ( + r"(?s)Invalid directory path.*" + r"full text search index 'text_idx'.*" + r"tokenize\(\) runs on the client.*'jieba/default'.*" + r"LANCE_LANGUAGE_MODEL_HOME.*match the files the server uses" +) + + +def test_remote_tokenize_model_backed_without_model_files(tmp_path, monkeypatch): + monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path)) + with mock_lancedb_connection(_jieba_fts_handler) as db: + table = db.open_table("test") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + table.tokenize("我们是中国人", index_name="text_idx") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + table.tokenize("我们是中国人", column="text") + + +@pytest.mark.asyncio +async def test_remote_tokenize_model_backed_without_model_files_async( + tmp_path, monkeypatch +): + monkeypatch.setenv("LANCE_LANGUAGE_MODEL_HOME", str(tmp_path)) + async with mock_lancedb_connection_async(_jieba_fts_handler) as db: + table = await db.open_table("test") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + await table.tokenize("我们是中国人", index_name="text_idx") + with pytest.raises(ValueError, match=_MODEL_FILES_ERROR): + await table.tokenize("我们是中国人", column="text") + + def test_stats(): stats = { "total_bytes": 38, diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 84c1b3fef..b722dc4a1 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -8871,6 +8871,131 @@ mod tests { ); } + #[tokio::test] + #[serial_test::serial(language_model_env)] + async fn test_tokenize_model_backed_tokenizer_without_model_files() { + let model_home = tempfile::tempdir().unwrap(); + // SAFETY: This is only called in tests, serialized on the env key. + unsafe { + std::env::set_var("LANCE_LANGUAGE_MODEL_HOME", model_home.path()); + } + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "base_tokenizer": "jieba/default", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": false, + "ascii_folding": false, + }) + .to_string(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": index_details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let result = table.tokenize("我们是中国人", "text_idx").await; + // SAFETY: This is only called in tests, serialized on the env key. + unsafe { + std::env::remove_var("LANCE_LANGUAGE_MODEL_HOME"); + } + + let err = result.unwrap_err(); + let Error::InvalidInput { message } = err else { + panic!("expected InvalidInput, got {err:?}"); + }; + // The original cause and the index context are preserved... + assert!(message.contains("Invalid directory path"), "{message}"); + assert!( + message.contains("full text search index 'text_idx'"), + "{message}" + ); + // ...and the message explains that this is a client-side limitation. + assert!( + message.contains("tokenize() runs on the client"), + "{message}" + ); + assert!(message.contains("'jieba/default'"), "{message}"); + assert!(message.contains("LANCE_LANGUAGE_MODEL_HOME"), "{message}"); + assert!( + message.contains("match the files the server uses"), + "{message}" + ); + } + + #[tokio::test] + async fn test_tokenize_unknown_tokenizer_error_has_no_model_files_hint() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "base_tokenizer": "no_such_tokenizer", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": false, + "ascii_folding": false, + }) + .to_string(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": index_details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let err = table.tokenize("hello", "text_idx").await.unwrap_err(); + let message = err.to_string(); + assert!( + !message.contains("tokenize() runs on the client"), + "{message}" + ); + } + #[tokio::test] async fn test_tokenize_requires_existing_index_name() { let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 39e6f1911..3123eb776 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -2264,7 +2264,11 @@ impl Table { /// /// Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in /// the client process from index metadata. For remote tables, this means the - /// same tokenizer model files must also exist locally. + /// same tokenizer model files must also exist locally, and must match the + /// files the server uses: the client cannot detect a mismatch, so it may + /// return tokens the server's index never produces. If the model files are + /// missing, this returns an error naming the tokenizer and + /// `LANCE_LANGUAGE_MODEL_HOME`. pub async fn tokenize(&self, query: &str, index_name: &str) -> Result> { let indices = self.inner.list_indices().await?; let matches = indices @@ -2298,7 +2302,8 @@ impl Table { /// The column must have exactly one FTS index. Model-backed tokenizers such /// as `jieba/*` and `lindera/*` are rebuilt in the client process from /// index metadata. For remote tables, this means the same tokenizer model - /// files must also exist locally. + /// files must also exist locally, and must match the files the server uses + /// (see [`Self::tokenize`]). pub async fn tokenize_with_column(&self, query: &str, column: &str) -> Result> { let schema = self.inner.schema().await?; let (column, _) = resolve_arrow_field_path(schema.as_ref(), column)?; @@ -2354,13 +2359,30 @@ impl Table { ), } })?; + let model_backed_tokenizer = serde_json::from_str::(details) + .ok() + .and_then(|details| { + let tokenizer = details.get("base_tokenizer")?.as_str()?.to_string(); + (tokenizer.starts_with("jieba/") || tokenizer.starts_with("lindera/")) + .then_some(tokenizer) + }); tokenize(query, ¶ms).map_err(|err| match err { - Error::InvalidInput { message } => Error::InvalidInput { - message: format!( + Error::InvalidInput { message } => { + let mut message = format!( "{} for full text search index '{}' for {}", message, index.name, selector_description - ), - }, + ); + if let Some(tokenizer) = model_backed_tokenizer { + message.push_str(&format!( + ". tokenize() runs on the client: the '{}' tokenizer is rebuilt \ + in this process from the index metadata, so its model files must \ + be available under LANCE_LANGUAGE_MODEL_HOME and, for remote \ + tables, match the files the server uses", + tokenizer + )); + } + Error::InvalidInput { message } + } err => err, }) } From b6d15558ec8df6c15f25b1012e330802b4d1f830 Mon Sep 17 00:00:00 2001 From: Christopher Krah Date: Thu, 22 Oct 2026 21:18:52 +0200 Subject: [PATCH 46/68] fix(python): stop retrying a missing API key or provider package (#4521) `retry_with_exponential_backoff` retries every exception except an auth error (`python/python/lancedb/embeddings/utils.py:298-318`). For Cohere, Gemini, Jina and VoyageAI a missing API key raises from `api_key_not_found_help`, and a missing provider package raises `ImportError` everywhere. Neither can succeed on a retry, yet both sleep through all seven rounds and then surface wrapped in `Exception("Maximum number of retries...")`. Before, on `142d032`, a VoyageAI function with no key set (script below, `time.sleep` recorded, `random.seed(0)`): ``` $ python lance_retry_deterministic.py python/python raised: Exception('Maximum number of retries (7) exceeded.', ValueError('Please set the VOYAGEAI_API_KEY environment v attempts: 8 sleeps: [4, 13, 37, 93, 280, 788, 2811] slept 4026 s before a missing-API-key error surfaced ``` After: ``` $ python lance_retry_deterministic.py python/python raised: MissingApiKeyError('Please set the VOYAGEAI_API_KEY environment variable.') attempts: 1 sleeps: [] ```
lance_retry_deterministic.py ```python # loads the pure-Python embeddings package from the tree, native extension stubbed import os, random, sys, time, types TREE = sys.argv[1] pkg = types.ModuleType("lancedb"); pkg.__path__ = [TREE + "/lancedb"]; sys.modules["lancedb"] = pkg native = types.ModuleType("lancedb._lancedb"); native.__getattr__ = lambda name: None sys.modules["lancedb._lancedb"] = native import pyarrow as pa from lancedb.embeddings import get_registry random.seed(0) os.environ.pop("VOYAGE_API_KEY", None) slept = [] time.sleep = slept.append func = get_registry().get("voyageai").create(name="voyage-3") try: func.compute_source_embeddings_with_retry(pa.array(["hello"])) except Exception as e: print("raised:", repr(e)[:110]) print("attempts:", len(slept) + 1, "sleeps:", [round(s) for s in slept]) print("slept %.0f s before a missing-API-key error surfaced" % sum(slept)) ```
The message names the wrong variable: `voyageai.py:501` reads `VOYAGE_API_KEY`. I left that for a separate fix. `api_key_not_found_help` now raises `MissingApiKeyError`, a `ValueError` subclass, so existing `except ValueError` callers still catch it. That error and `ImportError` join the 401 exemption from #2995. A plain `ValueError` is still retried, because `bedrock.py:171` wraps boto3 throttling in one. No issue exists for this. OpenAI is not covered: with no key, `openai.OpenAI()` raises `OpenAIError`, which has no status code, so it still retries. I can follow up on that. ### Verification ``` # lance_stub: an empty stand-in for lancedb._lancedb, loaded from the source tree $ pytest -p lance_stub --noconftest python/tests/test_embeddings.py -k "retry or 401 or deterministic" -q 5 passed, 23 deselected in 0.59s # utils.py from 142d032, tests kept: 2 failed, 3 passed, 23 deselected in 0.74s $ pytest -p lance_stub --noconftest python/tests/test_voyageai_embeddings.py -q 18 passed in 0.18s # not run: the full suite or pyright; the native extension did not build here ``` Baseline on `142d032` has 2 of these tests. The new ones cover a missing key, a missing package, and a wrapped throttling error that must still retry. ruff 0.16.10 check and format are clean. @wjones127 you merged the 401 change. Is a `ValueError` subclass the way you want to mark this? Co-authored-by: Xuanwo --- python/python/lancedb/embeddings/utils.py | 13 ++++++- python/python/tests/test_embeddings.py | 45 +++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/embeddings/utils.py b/python/python/lancedb/embeddings/utils.py index f76c9b1a5..fb2a94ecb 100644 --- a/python/python/lancedb/embeddings/utils.py +++ b/python/python/lancedb/embeddings/utils.py @@ -320,6 +320,11 @@ def _is_non_retryable_error(error: Exception) -> bool: if error.status_code in (401, 403): return True + # A missing API key or a missing optional dependency fails the same way on + # every attempt + if isinstance(error, (MissingApiKeyError, ImportError)): + return True + return False @@ -337,9 +342,15 @@ def url_retrieve(url: str): raise ConnectionError("could not download {} due to {}".format(url, err)) +class MissingApiKeyError(ValueError): + """Raised when an embedding provider's API key is not set.""" + + def api_key_not_found_help(provider): logging.error("Could not find API key for %s", provider) - raise ValueError(f"Please set the {provider.upper()}_API_KEY environment variable.") + raise MissingApiKeyError( + f"Please set the {provider.upper()}_API_KEY environment variable." + ) def is_flash_attn_2_available(): diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index d6d128845..2438bfe80 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -633,6 +633,51 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +@patch("time.sleep") +@pytest.mark.parametrize( + "error", + [ + pytest.param("missing_key", id="missing-api-key"), + pytest.param(ImportError("No module named 'voyageai'"), id="missing-package"), + ], +) +def test_no_retry_on_deterministic_errors(mock_sleep, error): + from lancedb.embeddings.utils import ( + api_key_not_found_help, + retry_with_exponential_backoff, + ) + + def func(): + if error == "missing_key": + api_key_not_found_help("voyageai") + raise error + + mock_func = MagicMock(side_effect=func) + wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3) + + with pytest.raises((ValueError, ImportError)): + wrapped_func() + + assert mock_func.call_count == 1 + assert mock_sleep.call_count == 0 + + +@patch("time.sleep") +def test_retry_on_wrapped_provider_error(mock_sleep): + """Providers such as Bedrock wrap throttling in a plain ValueError.""" + from lancedb.embeddings.utils import retry_with_exponential_backoff + + mock_func = MagicMock( + side_effect=[ + ValueError("Error raised by boto3 client: ThrottlingException"), + [0.1], + ] + ) + wrapped_func = retry_with_exponential_backoff(mock_func, max_retries=3) + + assert wrapped_func() == [0.1] + assert mock_func.call_count == 2 + assert mock_sleep.call_count == 1 def test_openai_bad_request_partial_batch_recovery(): # Regression test for issue #1677 From 9f4cd1618bddf1041c8ebdaa0b0ccb33656dc57c Mon Sep 17 00:00:00 2001 From: adenzhou1350 Date: Fri, 23 Oct 2026 03:19:06 +0800 Subject: [PATCH 47/68] fix(python): handle empty batch iterables with a schema (#4538) ## Summary `create_table(data=iter(()), schema=...)` currently raises `RuntimeError: coroutine raised StopIteration`, although an empty list with the same schema works. A tuple of batches also fails because the helper calls `next()` before obtaining an iterator. Normalize the iterable once and use the supplied schema for an empty reader. Without a schema, raise a clear `ValueError`. Nonempty schema inference and batch casting are unchanged. Tests cover synchronous/asynchronous creation, empty/nonempty iterators and tuples, and schema-aware `add()` sanitization. The default schema-less `add()` rejection is deliberately unchanged. ## Validation - Built the current Rust/PyO3 extension from base `9f37f392` with the repository's `ci` profile; both Python-only arms use the same verified extension and locked Python dependency versions. - Complete `python/tests/test_table.py` + `python/tests/test_scannable.py`: **268 passed, 0 skipped**. The unchanged production code with identical tests has **24 failures / 244 passes**. - Both arms used four-CPU affinity. An initial unrestricted-affinity run exhausted thread creation under the private memory limit; that failed attempt is not counted as a pass. - Package-wide Ruff check/format pass. Repository-wide checks still report pre-existing findings in unchanged `docs/web-tests`; those files are not modified here. AI-assisted change, reproduced against unchanged production code and validated with a native source build. This does not claim the full repository/platform matrix or a performance improvement. Signed-off-by: Xucheng Zhou --- python/python/lancedb/table.py | 15 +++++-- python/python/tests/test_table.py | 75 +++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 124fa34b5..c7de4161d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -469,7 +469,7 @@ def _into_pyarrow_reader( ): return data.collect().to_arrow().to_reader() elif isinstance(data, Iterable): - return _iterator_to_reader(data) + return _iterator_to_reader(data, schema) else: raise TypeError( f"Unknown data type {type(data)}. " @@ -479,10 +479,19 @@ def _into_pyarrow_reader( ) -def _iterator_to_reader(data: Iterable) -> pa.RecordBatchReader: +def _iterator_to_reader( + data: Iterable, schema: Optional[pa.Schema] = None +) -> pa.RecordBatchReader: # Each batch is treated as it's own reader, mainly so we can # re-use the _into_pyarrow_reader logic. - first = _into_pyarrow_reader(next(data)) + data = iter(data) + try: + first_batch = next(data) + except StopIteration: + if schema is None: + raise ValueError("Cannot create table from empty iterable without a schema") + return pa.RecordBatchReader.from_batches(schema, []) + first = _into_pyarrow_reader(first_batch) schema = first.schema def gen(): diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 238ee5330..e23b478f5 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4847,6 +4847,81 @@ def test_stats(mem_db: DBConnection): assert stats_with_index["fragment_stats"]["num_fragments"] == 1 +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.parametrize("with_schema", [False, True]) +def test_create_table_recordbatch_iterable( + mem_db: DBConnection, iterable_type, empty, with_schema +): + schema = pa.schema([pa.field("id", pa.int64())]) + batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)] + data = iterable_type(batches) + if empty and not with_schema: + with pytest.raises(ValueError, match="empty iterable without a schema"): + mem_db.create_table("test", data=data) + return + + table = mem_db.create_table( + "test", data=data, schema=schema if with_schema else None + ) + assert table.schema == schema + assert table.to_arrow().to_pylist() == ([] if empty else [{"id": 1}]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.parametrize("with_schema", [False, True]) +async def test_create_table_recordbatch_iterable_async( + mem_db_async: AsyncConnection, iterable_type, empty, with_schema +): + schema = pa.schema([pa.field("id", pa.int64())]) + batches = [] if empty else [pa.RecordBatch.from_pydict({"id": [1]}, schema)] + data = iterable_type(batches) + if empty and not with_schema: + with pytest.raises(ValueError, match="empty iterable without a schema"): + await mem_db_async.create_table("test", data=data) + return + + table = await mem_db_async.create_table( + "test", data=data, schema=schema if with_schema else None + ) + assert await table.schema() == schema + assert (await table.to_arrow()).to_pylist() == ([] if empty else [{"id": 1}]) + + +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_add_empty_recordbatch_iterable( + mem_db: DBConnection, iterable_type, on_bad_vectors +): + table = mem_db.create_table("test", [{"id": 1}]) + if on_bad_vectors == "error": + # The default scannable path still requires an input schema. This + # fix does not change its validation policy for schema-less input. + with pytest.raises(ValueError, match="empty iterator"): + table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + else: + # Sanitization has the existing table schema available. + table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + assert table.to_arrow().to_pylist() == [{"id": 1}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterable_type", [iter, tuple]) +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +async def test_add_empty_recordbatch_iterable_async( + mem_db_async: AsyncConnection, iterable_type, on_bad_vectors +): + table = await mem_db_async.create_table("test", [{"id": 1}]) + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="empty iterator"): + await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + else: + await table.add(iterable_type([]), on_bad_vectors=on_bad_vectors) + assert (await table.to_arrow()).to_pylist() == [{"id": 1}] + + def test_create_table_empty_list_with_schema(mem_db: DBConnection): """Test creating table with empty list data and schema From 2d54630a3c8cfc1faaec67289c864bba66578d2a Mon Sep 17 00:00:00 2001 From: Avanish Garg Date: Fri, 23 Oct 2026 00:52:24 +0530 Subject: [PATCH 48/68] fix(python): stop retrying permanent embedding errors and cap the backoff delay (#4542) Permanent embedding errors (HTTP 400/404/422 and other 4xx) were retried the full `max_retries` times with an uncapped, multiplicatively jittered backoff, so a typo'd model name looked like a hang of about an hour (median 50 min of sleeping with a 404; see the issue for the repro and numbers). Changes (`python/python/lancedb/embeddings/utils.py`, `embeddings/base.py`): - An HTTP status in 4xx other than 408/409/429 is raised immediately. The status is read from `status_code`, `http_status`, `response.status_code`, an integer `code` (google-api-core, urllib) or a botocore `response["ResponseMetadata"]["HTTPStatusCode"]`. AWS throttling reported as 400 stays retryable (matched by error code). `AuthenticationError` by name is unchanged. - Still retried: 408, 409, 429, >=500, and anything without a recognisable status (connection errors, timeouts, custom exceptions), so existing custom embedding functions behave as before. - New keyword `max_delay=60` (last in the signature, default keeps call sites working). `delay = min(delay * base, max_delay)` with equal jitter `delay/2 + random() * delay/2`. With defaults the waits are [1,2], [2,4], [4,8], [8,16], [16,32], [30,60], [30,60] s, worst case about 182 s instead of up to 2.5 h. - The "Maximum number of retries" exception now has `from e`. The `max_retries` comment in `base.py` describes the rules. Not changed: `embeddings/openai.py` (it swallows `BadRequestError`; open #4154/#4005/#3152 cover it), Retry-After handling (the code never used it), and the "C Data interface error" wrapping from the native layer (the original message is now included in it). Behaviour changes to be aware of: every 4xx except 408/409/429 fails fast, including 401/403 coming through `response.status_code` (requests/httpx), which used to be retried; a provider that signals transient failures with a 4xx and no AWS-style code is no longer retried; the first retry waits 1-2 s instead of 2-4 s. Tests (`python/python/tests/test_embeddings.py`): 48 new cases (23 test functions): permanent errors are not retried, transient errors are, botocore-style errors, the delay is capped, exact exponential growth without jitter, jitter bounds, no float overflow for a large `max_retries`, the original error is attached, and end to end through `table.add` and `table.search` with a stub function (one call and no sleep). On `main` 23 of them fail with the change stashed; the rest are guards that must not change. `test_embeddings.py`: 61 passed, 4 skipped. `ruff format`/`ruff check` are clean on the touched files. Verification caveat: `main`'s Python layer needs a newer native module than the 0.40.0 wheel, so the tests were run with the wheel's native module preloaded through a small shim (the same shim ran the baseline on `main`); a wider run of embeddings/pydantic/query/table tests has the same 36 failures before and after (that mismatch). Not run: real provider APIs (a fake OpenAI-style server was used), Node or Rust code. Fixes #4541 Co-authored-by: Xuanwo --- python/python/lancedb/embeddings/base.py | 8 +- python/python/lancedb/embeddings/utils.py | 103 +++++++-- python/python/tests/test_embeddings.py | 263 ++++++++++++++++++++++ 3 files changed, 354 insertions(+), 20 deletions(-) diff --git a/python/python/lancedb/embeddings/base.py b/python/python/lancedb/embeddings/base.py index 96ca6556c..c11914fab 100644 --- a/python/python/lancedb/embeddings/base.py +++ b/python/python/lancedb/embeddings/base.py @@ -38,9 +38,11 @@ class EmbeddingFunction(BaseModel, ABC): 3. ndims() which returns the number of dimensions of the vector column """ - max_retries: int = ( - 7 # Setting 0 disables retires. Maybe this should not be enabled by default, - ) + # Maximum number of retries for transient errors (rate limits, 5xx responses, + # connection errors, ...), with exponential backoff capped at 60 seconds between + # attempts. Permanent errors (HTTP 4xx other than 408/409/429) are never retried. + # Setting 0 disables retries. + max_retries: int = 7 _ndims: int = PrivateAttr() _original_args: dict = PrivateAttr() diff --git a/python/python/lancedb/embeddings/utils.py b/python/python/lancedb/embeddings/utils.py index fb2a94ecb..ff8fc6a18 100644 --- a/python/python/lancedb/embeddings/utils.py +++ b/python/python/lancedb/embeddings/utils.py @@ -241,15 +241,30 @@ def retry_with_exponential_backoff( exponential_base: float = 2, jitter: bool = True, max_retries: int = 7, + max_delay: float = 60, ): """Retry a function with exponential backoff. + Only errors that may succeed when repeated are retried. Errors that carry an + HTTP status code are retried when the status is 408, 409, 429 or >= 500; any + other 4xx status (for example 400, 401, 403, 404 or 422) is a permanent failure + that is raised immediately. Errors without a recognisable status code (for + example connection errors, timeouts, or exceptions raised by a custom embedding + function) are retried. + + The delay before the n-th retry is ``initial_delay * exponential_base ** n``, + never more than ``max_delay``. With ``jitter`` enabled, a random amount of up to + half of that delay is removed ("equal jitter"), so the actual wait is always in + the range ``[delay / 2, delay]``. + Args: func (function): The function to be retried. initial_delay (float): Initial delay in seconds (default is 1). exponential_base (float): The base for exponential backoff (default is 2). jitter (bool): Whether to add jitter to the delay (default is True). max_retries (int): Maximum number of retries (default is 7). + max_delay (float): Maximum delay between two attempts, in seconds + (default is 60). Returns: function: The decorated function. @@ -265,13 +280,12 @@ def retry_with_exponential_backoff( try: return func(*args, **kwargs) - # Currently retrying on all exceptions as there is no way to know the - # format of the error msgs used by different APIs. We'll log the error - # and say that it is assumed that if this portion errors out, it's due - # to rate limit but the user should check the error message to be sure. + # Errors are retried unless they are known to be permanent (see + # `_is_non_retryable_error`). Custom embedding functions can raise + # arbitrary exceptions, so unknown errors are assumed to be transient. except Exception as e: # noqa: PERF203 - # Don't retry on authentication errors (e.g., OpenAI 401) - # These are permanent failures that won't be fixed by retrying + # Don't retry permanent failures (e.g., OpenAI 401 or 404). They + # won't be fixed by retrying and would only delay the real error. if _is_non_retryable_error(e): raise @@ -280,24 +294,72 @@ def retry_with_exponential_backoff( if num_retries > max_retries: raise Exception( f"Maximum number of retries ({max_retries}) exceeded.", e - ) + ) from e - delay *= exponential_base * (1 + jitter * random.random()) + delay = min(delay * exponential_base, max_delay) + sleep_for = delay + if jitter: + sleep_for = delay / 2 + random.random() * delay / 2 logging.warning( "Error occurred: %s \n Retrying in %s seconds (retry %s of %s) \n", e, - delay, + sleep_for, num_retries, max_retries, ) - time.sleep(delay) + time.sleep(sleep_for) return wrapper +# HTTP statuses in the 4xx range that are worth retrying: request timeout, +# conflict (the OpenAI SDK treats it as transient) and rate limiting. +_RETRYABLE_CLIENT_STATUSES = (408, 409, 429) + +# AWS error codes that are reported with a 4xx status but are transient. +_RETRYABLE_AWS_ERROR_CODES = ( + "Throttl", + "TooManyRequests", + "RequestLimitExceeded", + "ProvisionedThroughputExceeded", + "ModelNotReady", + "ModelTimeout", +) + + +def _get_http_status(error: Exception): + """Best-effort extraction of an HTTP status code from an exception. + + Provider SDKs expose it in different ways: ``status_code`` (openai, cohere), + ``http_status`` (voyageai), ``response.status_code`` (requests, httpx), + ``code`` (google-api-core, urllib) or ``response["ResponseMetadata"]`` + (botocore). Returns None if no status can be found. + """ + candidates = [ + getattr(error, "status_code", None), + getattr(error, "http_status", None), + ] + response = getattr(error, "response", None) + if isinstance(response, dict): + candidates.append(response.get("ResponseMetadata", {}).get("HTTPStatusCode")) + else: + candidates.append(getattr(response, "status_code", None)) + candidates.append(getattr(error, "code", None)) + + for status in candidates: + if isinstance(status, int) and not isinstance(status, bool): + if 100 <= status <= 599: + return status + return None + + def _is_non_retryable_error(error: Exception) -> bool: """Check if an error should not be retried. + An error is non-retryable if it is an authentication error or if it carries an + HTTP client-error status (4xx) other than 408, 409 and 429. Errors without a + recognisable status are treated as retryable. + Args: error: The exception to check @@ -314,18 +376,25 @@ def _is_non_retryable_error(error: Exception) -> bool: if error_type == "AuthenticationError": return True - # Check for other common non-retryable HTTP status codes - # 401 Unauthorized, 403 Forbidden - if hasattr(error, "status_code"): - if error.status_code in (401, 403): - return True - # A missing API key or a missing optional dependency fails the same way on # every attempt if isinstance(error, (MissingApiKeyError, ImportError)): return True - return False + status = _get_http_status(error) + if status is None or not 400 <= status < 500: + return False + if status in _RETRYABLE_CLIENT_STATUSES: + return False + + # AWS reports throttling with a 400 status; the error code tells them apart. + response = getattr(error, "response", None) + if isinstance(response, dict): + aws_code = str(response.get("Error", {}).get("Code", "")) + if aws_code.startswith(_RETRYABLE_AWS_ERROR_CODES): + return False + + return True def url_retrieve(url: str): diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 2438bfe80..129fa6742 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -633,6 +633,269 @@ def test_openai_no_retry_on_401(mock_sleep): assert mock_sleep.call_count == 0 +class _HttpStatusError(Exception): + """Stand-in for provider SDK errors (openai, cohere, ...) with a status code.""" + + def __init__(self, status_code, message="error"): + super().__init__(message) + self.status_code = status_code + + +class _ResponseStatusError(Exception): + """Stand-in for requests/httpx errors, which expose `response.status_code`.""" + + def __init__(self, status_code): + super().__init__(f"HTTP {status_code}") + self.response = MagicMock(status_code=status_code) + + +def _run_with_retry(side_effect, **kwargs): + """Run a function through the retry helper with a mocked `time.sleep`. + + Returns (mock_func, sleeps, error). + """ + from lancedb.embeddings.utils import retry_with_exponential_backoff + + mock_func = MagicMock(side_effect=side_effect) + sleeps = [] + error = None + with patch("time.sleep", side_effect=sleeps.append): + try: + retry_with_exponential_backoff(mock_func, **kwargs)() + except Exception as e: + error = e + return mock_func, sleeps, error + + +@pytest.mark.parametrize("status", [400, 401, 403, 404, 405, 410, 413, 422]) +@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError]) +def test_retry_does_not_retry_permanent_http_errors(status, make_error): + err = make_error(status) + mock_func, sleeps, error = _run_with_retry(err, max_retries=7) + + assert error is err + assert mock_func.call_count == 1 + assert sleeps == [] + + +@pytest.mark.parametrize("status", [408, 409, 429, 500, 502, 503, 504]) +@pytest.mark.parametrize("make_error", [_HttpStatusError, _ResponseStatusError]) +def test_retry_retries_transient_http_errors(status, make_error): + mock_func, sleeps, error = _run_with_retry( + [make_error(status), make_error(status), "result"], max_retries=7 + ) + + assert error is None + assert mock_func.call_count == 3 + assert len(sleeps) == 2 + + +@pytest.mark.parametrize( + "exc", + [ + ConnectionError("connection reset"), + TimeoutError("timed out"), + OSError("network is unreachable"), + ValueError("custom embedding function failure"), + RuntimeError("boom"), + Exception("plain exception"), + ], +) +def test_retry_retries_errors_without_http_status(exc): + mock_func, sleeps, error = _run_with_retry([exc, exc, "result"], max_retries=7) + + assert error is None + assert mock_func.call_count == 3 + assert len(sleeps) == 2 + + +def test_retry_error_status_attributes_of_other_types_are_ignored(): + # `code` is a string for openai errors ("model_not_found"), and arbitrary + # exceptions may carry unrelated numeric attributes. Neither is an HTTP status. + class WeirdError(Exception): + code = "model_not_found" + status_code = None + + class ExitCodeError(Exception): + code = 2 + + for exc in (WeirdError("a"), ExitCodeError("b")): + mock_func, sleeps, error = _run_with_retry(exc, max_retries=2) + assert mock_func.call_count == 3 + assert len(sleeps) == 2 + + +def test_retry_botocore_style_errors(): + # botocore ClientError exposes `response["ResponseMetadata"]["HTTPStatusCode"]`. + # Bedrock reports throttling with a 400, which must stay retryable. + class ClientError(Exception): + def __init__(self, status, code): + super().__init__(code) + self.response = { + "ResponseMetadata": {"HTTPStatusCode": status}, + "Error": {"Code": code}, + } + + throttled = ClientError(400, "ThrottlingException") + mock_func, sleeps, error = _run_with_retry([throttled, "result"], max_retries=3) + assert error is None + assert mock_func.call_count == 2 + + invalid = ClientError(400, "ValidationException") + mock_func, sleeps, error = _run_with_retry(invalid, max_retries=3) + assert error is invalid + assert mock_func.call_count == 1 + assert sleeps == [] + + +def test_retry_authentication_error_by_name_not_retried(): + class AuthenticationError(Exception): + pass + + mock_func, sleeps, error = _run_with_retry(AuthenticationError("bad key")) + assert isinstance(error, AuthenticationError) + assert mock_func.call_count == 1 + assert sleeps == [] + + +def test_retry_gives_up_after_max_retries_with_original_error_attached(): + original = _HttpStatusError(503, "overloaded") + mock_func, sleeps, error = _run_with_retry(original, max_retries=3) + + assert mock_func.call_count == 4 + assert len(sleeps) == 3 + assert "Maximum number of retries (3) exceeded" in str(error) + assert error.args[1] is original + assert error.__cause__ is original + + +def test_retry_delay_is_capped_by_max_delay(): + # Default settings (7 retries, 60s cap): the old implementation could sleep for + # up to ~47 minutes in a single retry and hours in total. + exc = _HttpStatusError(503) + for seed in range(50): + import random + + random.seed(seed) + _, sleeps, _ = _run_with_retry(exc) + assert len(sleeps) == 7 + assert max(sleeps) <= 60 + assert sum(sleeps) <= 7 * 60 + + # Custom cap + _, sleeps, _ = _run_with_retry(exc, max_retries=10, max_delay=5) + assert len(sleeps) == 10 + assert max(sleeps) <= 5 + + +def test_retry_delay_without_jitter_is_exact_exponential_then_capped(): + _, sleeps, _ = _run_with_retry( + _HttpStatusError(500), + initial_delay=1, + exponential_base=2, + jitter=False, + max_retries=8, + max_delay=30, + ) + assert sleeps == [2, 4, 8, 16, 30, 30, 30, 30] + + +def test_retry_jitter_is_bounded(): + # Each wait is within [delay / 2, delay], where delay grows by exponential_base + # per retry (no random factor is compounded across retries). + exc = _HttpStatusError(500) + for seed in range(50): + import random + + random.seed(seed) + _, sleeps, _ = _run_with_retry(exc, max_retries=8, max_delay=100) + for i, slept in enumerate(sleeps): + delay = min(2 ** (i + 1), 100) + assert delay / 2 <= slept <= delay + + # Extremes of the random source map to the ends of the range + with patch("random.random", return_value=0.0): + _, sleeps, _ = _run_with_retry(exc, max_retries=3) + assert sleeps == [1, 2, 4] + with patch("random.random", return_value=1.0): + _, sleeps, _ = _run_with_retry(exc, max_retries=3) + assert sleeps == [2, 4, 8] + + +def test_retry_large_max_retries_does_not_overflow(): + _, sleeps, _ = _run_with_retry(_HttpStatusError(503), max_retries=1500) + assert len(sleeps) == 1500 + assert max(sleeps) <= 60 + + +def _register_failing_embedding(name, error): + """Register an embedding function that always raises `error`. + + Returns the class and a list that records every call. The function is + re-created from table metadata, so the call count has to live outside of it. + """ + calls = [] + + @register(name) + class FailingEmbeddingFunction(TextEmbeddingFunction): + def ndims(self): + return 4 + + def generate_embeddings(self, texts): + calls.append(len(texts)) + raise error + + return FailingEmbeddingFunction, calls + + +@pytest.mark.parametrize("status", [400, 404, 422]) +def test_permanent_embedding_error_is_not_retried(tmp_path, status): + """A permanent provider error (e.g. a misspelled model name) must surface + immediately from both `table.add` and `table.search` instead of being retried + with minutes of backoff.""" + error = _HttpStatusError(status, "The model `typo-embed` does not exist") + cls, calls = _register_failing_embedding(f"failing-embedding-{status}", error) + func = cls.create() # default max_retries=7 + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table("t", schema=Schema) + + with patch("time.sleep") as mock_sleep: + with pytest.raises(Exception, match="The model `typo-embed` does not exist"): + table.add([{"text": "hello"}]) + assert len(calls) == 1 + + with pytest.raises(_HttpStatusError, match="does not exist"): + table.search("hello").limit(1).to_list() + assert len(calls) == 2 + + mock_sleep.assert_not_called() + + +def test_transient_embedding_error_is_retried_in_table_add(tmp_path): + cls, calls = _register_failing_embedding( + "failing-embedding-503", _HttpStatusError(503) + ) + func = cls.create(max_retries=2) + + class Schema(LanceModel): + text: str = func.SourceField() + vector: Vector(func.ndims()) = func.VectorField() + + db = lancedb.connect(tmp_path) + table = db.create_table("t", schema=Schema) + + with patch("time.sleep") as mock_sleep: + with pytest.raises(Exception, match="Maximum number of retries"): + table.add([{"text": "hello"}]) + + assert len(calls) == 3 + assert mock_sleep.call_count == 2 + assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list) @patch("time.sleep") @pytest.mark.parametrize( "error", From bc4c32b779de70c4c62a8411ea7eca1c165c9f7e Mon Sep 17 00:00:00 2001 From: Avanish Garg Date: Fri, 23 Oct 2026 00:52:39 +0530 Subject: [PATCH 49/68] fix: keep tagged versions when pruning in optimize (#4544) `optimize()` fails once a tagged version is older than the cleanup cutoff, although the docs say tagged versions are exempt from cleanup, and the failed prune also skips the index-optimize step (repro and numbers in the issue). Change: `OptimizeAction::Prune` already exposes `error_if_tagged_old_versions: Option`, so no new API. `cleanup_old_versions` and `cleanup_old_versions_before` in `rust/lancedb/src/table/optimize.rs` now use `unwrap_or(false)`, which covers `OptimizeAction::Prune`, `OptimizeAction::All` and `optimize_prune_before`. The Python and Node bindings pass `None`, which now means "retain tagged versions"; I left `python/src/table.rs` and `nodejs/src/table.rs` untouched (this also avoids a conflict with #3870, which edits that region). `Some(true)` still opts in to the error. Doc comments updated; the sentence in `docs/web/management/table.mdx` is accurate once this lands. Tests: Rust `test_optimize_prune_retains_tagged_versions` (`Some(true)` still errors; the default retains the tag; `optimize_prune_before`; `checkout_tag` plus a data read), a Python sync test with a scalar index (index step runs, a second optimize works, untagged versions are pruned, `checkout("release-1")` works), a Python async test and a Node test. **What I could and could not verify (please read):** - I could NOT compile Rust here (not enough disk), so `cargo check`, `cargo test`, the Python/Node binding builds and therefore the fixed-path assertions of the new Rust, Python and Node tests have not been run. This is why the PR is a draft; I will mark it ready once CI is green and fix whatever it finds. - Ran: both new Python tests fail on the released 0.40.0 wheel with the exact error above (so they do catch the bug); pylance 13.0.0 confirms the semantics (`cleanup_old_versions(error_if_tagged_old_versions=False)` removes the untagged versions and keeps the tagged one, still readable with `lance.dataset(version="keep")`); in Lance v14.0.0-beta.7 `cleanup.rs` the policy field is `bool`, the builder default is `true`, and Lance's own `auto_cleanup_hook` uses `Some(false)`; `rustfmt --edition 2024 --check` passes on `optimize.rs`. - The Rust and Node tests rely on timestamp ordering (`nextMillisecond()`, `days(0)`/`now()`), so they could be flaky: tell me if you prefer another approach. Behaviour change: anyone who relied on the error to learn that a tag blocks cleanup now gets a silent retain. Not changed: the deprecated sync `LanceTable.cleanup_old_versions` calls pylance directly and still errors on tagged versions (the keyword may not exist in the oldest supported pylance); `freshness.rs:1377` and `refresh.rs:1077` still pass `None` (I did not check what those paths do with tagged versions). Fixes #4543 --- nodejs/__test__/table.test.ts | 15 +++++ python/python/tests/test_table.py | 51 ++++++++++++++ rust/lancedb/src/table/optimize.rs | 105 +++++++++++++++++++++++++++-- 3 files changed, 167 insertions(+), 4 deletions(-) diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 7d10f8b59..84340c72c 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3742,6 +3742,21 @@ describe("when optimizing a dataset", () => { expect(stats.prune.oldVersionsRemoved).toBe(2); }); + it("retains tagged versions older than the cleanup cutoff", async () => { + await (await table.tags()).create("keep", 1); + await table.add([{ id: 3 }]); + + const cutoff = await nextMillisecond(); + const stats = await table.optimize({ cleanupOlderThan: cutoff }); + // Versions 2 and 3 are pruned, tagged version 1 is retained + expect(stats.prune.oldVersionsRemoved).toBe(2); + + await table.checkout("keep"); + expect(await table.countRows()).toBe(1); + await table.checkoutLatest(); + expect(await table.countRows()).toBe(3); + }); + it("delete unverified", async () => { const version = await table.version(); const versionFile = `${tmpDir.name}/${table.name}.lance/_versions/${String( diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index e23b478f5..b55abfa7e 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4792,6 +4792,57 @@ def test_optimize_warns_on_zero_cleanup_sync(mem_db: DBConnection): assert all("concurrent" in str(w.message) for w in seen) +def test_optimize_retains_tagged_versions(tmp_db: DBConnection): + # A tagged version older than the cleanup cutoff must be retained, not make + # optimize() fail part-way (after compaction, before the index step). + table = tmp_db.create_table( + "test", data=pa.table({"id": list(range(100)), "text": ["a"] * 100}) + ) + table.create_index("id", config=BTree()) + table.tags.create("release-1", table.version) + tagged_version = table.version + for i in range(3): + table.add(pa.table({"id": [1000 + i], "text": ["new"]})) + assert table.stats()["fragment_stats"]["num_fragments"] == 4 + assert table.index_stats("id_idx").num_unindexed_rows == 3 + sleep(1.1) + + table.optimize(cleanup_older_than=timedelta(seconds=1)) + + # The index step ran after the prune step + assert table.index_stats("id_idx").num_unindexed_rows == 0 + # It stays optimizable + table.optimize(cleanup_older_than=timedelta(seconds=1)) + assert table.stats()["fragment_stats"]["num_fragments"] == 1 + + # Untagged old versions were pruned, the tagged one was retained + versions = [v["version"] for v in table.list_versions()] + assert tagged_version in versions + for added in range(tagged_version + 1, tagged_version + 4): + assert added not in versions + assert len(table) == 103 + + table.checkout("release-1") + assert len(table) == 100 + assert table.to_arrow().num_rows == 100 + + +@pytest.mark.asyncio +async def test_optimize_retains_tagged_versions_async(tmp_db_async: AsyncConnection): + table = await tmp_db_async.create_table("test", data=[{"x": [1]}]) + await table.tags.create("keep", await table.version()) + await table.add(data=[{"x": [2]}]) + await table.add(data=[{"x": [3]}]) + sleep(1.1) + + stats = await table.optimize(cleanup_older_than=timedelta(seconds=1)) + assert stats.prune.old_versions_removed >= 1 + + await table.checkout("keep") + assert await table.count_rows() == 1 + assert await table.query().to_arrow() == pa.table({"x": [[1]]}) + + def test_replace_field_metadata(tmp_path): db = lancedb.connect(tmp_path) table = db.create_table("my_table", data=[{"x": 0}]) diff --git a/rust/lancedb/src/table/optimize.rs b/rust/lancedb/src/table/optimize.rs index 2a502b059..ebe03359c 100644 --- a/rust/lancedb/src/table/optimize.rs +++ b/rust/lancedb/src/table/optimize.rs @@ -69,7 +69,11 @@ pub enum OptimizeAction { /// **WARNING**: This should only be set to true if you can guarantee that no other process is /// currently working on this dataset. Otherwise the dataset could be put into a corrupted state. delete_unverified: Option, + /// Tagged versions are never removed, regardless of their age. + /// /// If true, an error will be returned if there are any old versions that are still tagged. + /// If false or `None` (the default), tagged versions are silently retained and all other + /// old versions are removed. error_if_tagged_old_versions: Option, }, /// Optimize the indices @@ -125,6 +129,12 @@ pub(crate) async fn optimize_indices(table: &NativeTable, options: &OptimizeOpti /// no other process is currently working on this dataset. Otherwise the /// dataset could be put into a corrupted state. /// +/// * `error_if_tagged_old_versions` - Tagged versions are always retained. If +/// this is true, an error is returned when a tagged version is older than the +/// cutoff. Defaults to false (the tagged versions are retained silently), +/// unlike Lance's own default, so that `optimize` does not fail (and skip the +/// steps after pruning) once a tagged version ages past the cutoff. +/// /// This calls into [lance::dataset::Dataset::cleanup_old_versions] and /// returns the result. pub(crate) async fn cleanup_old_versions( @@ -136,7 +146,11 @@ pub(crate) async fn cleanup_old_versions( table.dataset.ensure_mutable()?; let dataset = table.dataset.get().await?; let stats = dataset - .cleanup_old_versions(older_than, delete_unverified, error_if_tagged_old_versions) + .cleanup_old_versions( + older_than, + delete_unverified, + Some(error_if_tagged_old_versions.unwrap_or(false)), + ) .await?; // Computed-column signature sidecars live outside lance's directories; // drop the ones the surviving versions no longer reference. @@ -161,9 +175,8 @@ pub(crate) async fn cleanup_old_versions_before( if let Some(delete_unverified) = delete_unverified { policy = policy.delete_unverified(delete_unverified); } - if let Some(error_if_tagged_old_versions) = error_if_tagged_old_versions { - policy = policy.error_if_tagged_old_versions(error_if_tagged_old_versions); - } + // Tagged versions are retained rather than raising an error unless asked. + policy = policy.error_if_tagged_old_versions(error_if_tagged_old_versions.unwrap_or(false)); let stats = dataset.cleanup_with_policy(policy.build()).await?; // Computed-column signature sidecars live outside lance's directories; // drop the ones the surviving versions no longer reference. @@ -511,6 +524,90 @@ mod tests { assert_eq!(all_values, expected); } + #[tokio::test] + async fn test_optimize_prune_retains_tagged_versions() { + let tmpdir = tempfile::tempdir().unwrap(); + let conn = connect(tmpdir.path().to_str().unwrap()) + .execute() + .await + .unwrap(); + + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let make_batch = |start: i32| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(start..start + 10))], + ) + .unwrap() + }; + + let table = conn + .create_table("test_prune_tagged", make_batch(0)) + .execute() + .await + .unwrap(); + // Tag the first version, then create more versions after it. + table.tags().await.unwrap().create("keep", 1).await.unwrap(); + for i in 1..4 { + table.add(make_batch(i * 10)).execute().await.unwrap(); + } + assert_eq!(table.list_versions().await.unwrap().len(), 4); + + // Explicitly asking for an error still errors, and leaves everything in place. + let err = table + .optimize(OptimizeAction::Prune { + older_than: Some(chrono::Duration::try_days(0).unwrap()), + delete_unverified: Some(true), + error_if_tagged_old_versions: Some(true), + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("tagged version"), "{err}"); + assert_eq!(table.list_versions().await.unwrap().len(), 4); + + // By default, tagged versions are retained and the other old versions are removed. + let stats = table + .optimize(OptimizeAction::Prune { + older_than: Some(chrono::Duration::try_days(0).unwrap()), + delete_unverified: Some(true), + error_if_tagged_old_versions: None, + }) + .await + .unwrap(); + assert_eq!(stats.prune.unwrap().old_versions, 2); + let mut versions = table + .list_versions() + .await + .unwrap() + .into_iter() + .map(|v| v.version) + .collect::>(); + versions.sort(); + assert_eq!(versions, vec![1, 4]); + + // Same through the absolute-cutoff entry point used by the Node binding. + table.add(make_batch(40)).execute().await.unwrap(); + let stats = table + .optimize_prune_before(chrono::Utc::now(), Some(true), None) + .await + .unwrap(); + assert_eq!(stats.prune.unwrap().old_versions, 1); + + // The tagged version is still readable, including its data files. + assert_eq!(table.count_rows(None).await.unwrap(), 50); + table.checkout_tag("keep").await.unwrap(); + let batches = table + .query() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 10); + } + #[tokio::test] async fn test_optimize_index() { let conn = connect("memory://").execute().await.unwrap(); From 5fa30ae404fc7df335c8031ae474400b834c62b2 Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:22:53 -0400 Subject: [PATCH 50/68] fix(python): honor explicit Cohere and VoyageAI reranker API keys (#4546) `CohereReranker` and `VoyageAIReranker` read the environment variable before the `api_key` argument. When `COHERE_API_KEY` or `VOYAGE_API_KEY` is set, the reranker ignores an explicit `api_key` and uses the account from the environment. The docs say the opposite. `docs/web/integrations/reranking/cohere.mdx` and `voyageai.mdx` use the environment variable only when `api_key` is not provided. `JinaReranker` and `WatsonxReranker` already put the explicit key first, and #4515 makes the same change for `OpenaiReranker`. ## Fix `CohereReranker._client` and `VoyageAIReranker._client` now use `self.api_key` first and the environment variable as the fallback. ## Tests `python/python/tests/test_reranker_credentials.py` replaces the SDK module with a stub whose `Client` records the API key. It covers both keys set, each key alone, and no key. With the stub, the test runs without `cohere` or `voyageai` installed. - On `main`, the 2 cases with both keys set fail and the other 6 pass. With the fix, all 8 pass. - `test_rerankers.py` with the new file: 24 passed, 13 skipped (the skipped tests need real API keys). - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed files copied over it. In that wheel, both reranker files are byte-identical to `main` before this change. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/rerankers/cohere.py | 2 +- python/python/lancedb/rerankers/voyageai.py | 2 +- .../python/tests/test_reranker_credentials.py | 74 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) create mode 100644 python/python/tests/test_reranker_credentials.py diff --git a/python/python/lancedb/rerankers/cohere.py b/python/python/lancedb/rerankers/cohere.py index efd752c22..a1e797f15 100644 --- a/python/python/lancedb/rerankers/cohere.py +++ b/python/python/lancedb/rerankers/cohere.py @@ -62,7 +62,7 @@ class CohereReranker(Reranker): "COHERE_API_KEY not set. Either set it in your environment or \ pass it as `api_key` argument to the CohereReranker." ) - return cohere.Client(os.environ.get("COHERE_API_KEY") or self.api_key) + return cohere.Client(self.api_key or os.environ.get("COHERE_API_KEY")) def _rerank(self, result_set: pa.Table, query: str): result_set = self._handle_empty_results(result_set) diff --git a/python/python/lancedb/rerankers/voyageai.py b/python/python/lancedb/rerankers/voyageai.py index d3dbd59ac..28847564b 100644 --- a/python/python/lancedb/rerankers/voyageai.py +++ b/python/python/lancedb/rerankers/voyageai.py @@ -64,7 +64,7 @@ class VoyageAIReranker(Reranker): pass it as `api_key` argument to the VoyageAIReranker." ) return voyageai.Client( - api_key=os.environ.get("VOYAGE_API_KEY") or self.api_key, + api_key=self.api_key or os.environ.get("VOYAGE_API_KEY"), ) def _rerank(self, result_set: pa.Table, query: str): diff --git a/python/python/tests/test_reranker_credentials.py b/python/python/tests/test_reranker_credentials.py new file mode 100644 index 000000000..c8f63a9ef --- /dev/null +++ b/python/python/tests/test_reranker_credentials.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import sys +import types + +import pytest + +from lancedb.rerankers import CohereReranker, VoyageAIReranker + +RERANKERS = [ + pytest.param("cohere", "COHERE_API_KEY", CohereReranker, {}, id="cohere"), + pytest.param( + "voyageai", + "VOYAGE_API_KEY", + VoyageAIReranker, + {"model_name": "rerank-2.5"}, + id="voyageai", + ), +] + + +def fake_sdk(monkeypatch, module_name): + """Replace the SDK module with one whose Client records the API key.""" + keys = [] + + def client(api_key=None): + keys.append(api_key) + return object() + + module = types.ModuleType(module_name) + module.Client = client + monkeypatch.setitem(sys.modules, module_name, module) + return keys + + +@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS) +@pytest.mark.parametrize( + "explicit_key, environment_key, expected_key", + [ + ("explicit-test-key", "environment-test-key", "explicit-test-key"), + ("explicit-test-key", None, "explicit-test-key"), + (None, "environment-test-key", "environment-test-key"), + ], +) +def test_reranker_api_key_precedence( + monkeypatch, + module_name, + env_var, + reranker_cls, + kwargs, + explicit_key, + environment_key, + expected_key, +): + keys = fake_sdk(monkeypatch, module_name) + if environment_key is None: + monkeypatch.delenv(env_var, raising=False) + else: + monkeypatch.setenv(env_var, environment_key) + + reranker_cls(api_key=explicit_key, **kwargs)._client + assert keys == [expected_key] + + +@pytest.mark.parametrize("module_name, env_var, reranker_cls, kwargs", RERANKERS) +def test_reranker_requires_api_key( + monkeypatch, module_name, env_var, reranker_cls, kwargs +): + keys = fake_sdk(monkeypatch, module_name) + monkeypatch.delenv(env_var, raising=False) + with pytest.raises(ValueError, match=f"{env_var} not set"): + reranker_cls(**kwargs)._client + assert keys == [] From 353ce4d80c9e66f8f1ac9e49dbcec0ec899f1f9d Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:23:07 -0400 Subject: [PATCH 51/68] fix(python): give each JinaEmbeddings instance its own session (#4547) `JinaEmbeddings` stores its `requests` session in a `ClassVar`. The first instance builds the session with its own API key, and every later instance reuses that session. So a second instance with a different `api_key` sends the first instance's key. ## Fix The session is now a per-instance `cached_property`, the same pattern `OpenAIEmbeddings` uses for `_openai_client`. `_init_client` and the `ClassVar` are gone. An instance with no `api_key` and no `JINA_API_KEY` now raises the missing-key error. Before, it reused the session of an earlier instance and sent that instance's key. ## Tests `test_jina_embeddings_use_their_own_api_key` creates two instances with different keys through the registry. It stubs `requests.Session.post` and records the `Authorization` header of each request. - On `main`, the test fails because both requests send `Bearer key-a`. With the fix, it passes. - `test_embeddings.py`: 18 passed, 8 skipped. - An instance still pickles after the session exists, and `safe_model_dump()` still returns only `api_key`. - On pydantic 2.7.4, the minimum version, a model with the same `cached_property` gets one session per instance and leaves it out of `model_dump()`. - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed file copied over it. In that wheel, `embeddings/jinaai.py` is byte-identical to `main` before this change. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/embeddings/jinaai.py | 26 +++++++++++----------- python/python/tests/test_embeddings.py | 25 +++++++++++++++++++++ 2 files changed, 38 insertions(+), 13 deletions(-) diff --git a/python/python/lancedb/embeddings/jinaai.py b/python/python/lancedb/embeddings/jinaai.py index f6ab601b3..25b5a940b 100644 --- a/python/python/lancedb/embeddings/jinaai.py +++ b/python/python/lancedb/embeddings/jinaai.py @@ -5,9 +5,10 @@ import os import io import base64 +from functools import cached_property from urllib.parse import urlparse from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, List, Union, Optional, Any, Dict +from typing import TYPE_CHECKING, List, Union, Optional, Any, Dict import numpy as np import pyarrow as pa @@ -51,7 +52,6 @@ class JinaEmbeddings(EmbeddingFunction): name: str = "jina-clip-v1" api_key: Optional[str] = None - _session: ClassVar = None def ndims(self): # TODO: fix hardcoding @@ -206,8 +206,7 @@ class JinaEmbeddings(EmbeddingFunction): texts: list[str] or np.ndarray (of str) The texts to embed """ - self._init_client() - resp = JinaEmbeddings._session.post( # type: ignore + resp = self._session.post( API_URL, json={"input": input, "model": self.name} ).json() if "data" not in resp: @@ -220,14 +219,15 @@ class JinaEmbeddings(EmbeddingFunction): return [result["embedding"] for result in sorted_embeddings] - def _init_client(self): + @cached_property + def _session(self): import requests - if JinaEmbeddings._session is None: - if self.api_key is None and os.environ.get("JINA_API_KEY") is None: - api_key_not_found_help("jina") - api_key = self.api_key or os.environ.get("JINA_API_KEY") - JinaEmbeddings._session = requests.Session() - JinaEmbeddings._session.headers.update( - {"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"} - ) + if self.api_key is None and os.environ.get("JINA_API_KEY") is None: + api_key_not_found_help("jina") + api_key = self.api_key or os.environ.get("JINA_API_KEY") + session = requests.Session() + session.headers.update( + {"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"} + ) + return session diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index 129fa6742..e319d3048 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -1053,3 +1053,28 @@ def test_jina_generate_image_input_dict_local_path(tmp_path): image_dict = JinaEmbeddings._generate_image_input_dict(image) assert "image" in image_dict assert isinstance(image_dict["image"], str) and len(image_dict["image"]) > 0 + + +def test_jina_embeddings_use_their_own_api_key(monkeypatch): + """Each JinaEmbeddings instance must send its own API key.""" + pytest.importorskip("requests") + monkeypatch.delenv("JINA_API_KEY", raising=False) + registry = get_registry() + registry.set_var("jina_key_a", "key-a") + registry.set_var("jina_key_b", "key-b") + func_a = registry.get("jina").create(api_key="$var:jina_key_a") + func_b = registry.get("jina").create(api_key="$var:jina_key_b") + + sent = [] + + def fake_post(session, url, json=None, **kwargs): + sent.append(session.headers["Authorization"]) + response = MagicMock() + response.json.return_value = {"data": [{"index": 0, "embedding": [0.0]}]} + return response + + with patch("requests.Session.post", fake_post): + func_a.compute_source_embeddings(["hello"]) + func_b.compute_source_embeddings(["hello"]) + + assert sent == ["Bearer key-a", "Bearer key-b"] From 2af7d88ec7e974484e1be0d9f7bfba53732e0ef6 Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:23:21 -0400 Subject: [PATCH 52/68] fix(python): keep OpenaiReranker scores on the rows they rate (#4549) `OpenaiReranker._rerank` asks the chat model for a relevance score per document. The model can return the documents in any order, for example sorted by score. The reranker then replaces the `text` column with the documents in that order and appends the scores in that order. The other columns keep the input order. So after a rerank, the `text` and `_relevance_score` of a row belong to a different row. The `id`, the vector and every other column still describe the original row. If the model leaves out a document, `append_column` raises `ArrowInvalid`, because the new column has fewer rows than the table. ## Fix `_rerank` keeps the `text` column unchanged, so it also keeps its position and type. It maps each score back to its row by the document content. If the model leaves out a document, `_rerank` raises a `ValueError` that gives the number of documents with no score. ## Tests `python/python/tests/test_openai_reranker.py` replaces the chat client with a stub that returns a fixed reply. The tests do not need `openai` or an API key. - One test returns the documents sorted by score. It checks that each `id` keeps its own `text` and score. On `main`, the row with id 1 gets the text `cats purr`, which belongs to id 2. - One test returns 2 of 3 documents. It checks for the `ValueError` with the count of missing documents. On `main`, the call raises `ArrowInvalid: Added column's length must match table's length`. - Both tests fail on `main` and pass with the fix. - `test_rerankers.py`: 16 passed, 13 skipped, on the lancedb 0.40.0 wheel with and without the fix. The skipped tests need real API keys. - `ruff format --check` and `ruff check` (0.9.9) pass on the changed files. - I did not build the native extension. I ran the tests against the lancedb 0.40.0 wheel with the changed file copied over it. In that wheel, `rerankers/openai.py` and `rerankers/base.py` are byte-identical to `main`. The Node SDK has no OpenAI reranker, so there is no parity change. This change does not conflict with #4515, which edits `_client` in the same file. Co-authored-by: Claude Opus 5.5 --- python/python/lancedb/rerankers/openai.py | 21 ++++--- python/python/tests/test_openai_reranker.py | 66 +++++++++++++++++++++ 2 files changed, 78 insertions(+), 9 deletions(-) create mode 100644 python/python/tests/test_openai_reranker.py diff --git a/python/python/lancedb/rerankers/openai.py b/python/python/lancedb/rerankers/openai.py index 4cbbb0822..da0394d1c 100644 --- a/python/python/lancedb/rerankers/openai.py +++ b/python/python/lancedb/rerankers/openai.py @@ -70,15 +70,18 @@ class OpenaiReranker(Reranker): ], ) results = json.loads(response.choices[0].message.content)["documents"] - docs, scores = list( - zip(*[(result["content"], result["relevance_score"]) for result in results]) - ) # tuples - # replace the self.column column with the docs - result_set = result_set.drop(self.column) - result_set = result_set.append_column( - self.column, pa.array(docs, type=pa.string()) - ) - # add the scores + # The model returns the documents in its own order, so match each score + # back to its row by content. + scores_by_content = { + result["content"]: result["relevance_score"] for result in results + } + missing = sum(doc not in scores_by_content for doc in docs) + if missing: + raise ValueError( + f"The model returned no relevance_score for {missing} of " + f"{len(docs)} documents" + ) + scores = [scores_by_content[doc] for doc in docs] result_set = result_set.append_column( "_relevance_score", pa.array(scores, type=pa.float32()) ) diff --git a/python/python/tests/test_openai_reranker.py b/python/python/tests/test_openai_reranker.py new file mode 100644 index 000000000..43475103f --- /dev/null +++ b/python/python/tests/test_openai_reranker.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json +from types import SimpleNamespace + +import pyarrow as pa +import pytest + +from lancedb.rerankers import OpenaiReranker + + +def reranker_replying(documents): + """Return an OpenaiReranker whose chat client replies with these documents.""" + reply = json.dumps({"documents": documents}) + + def create(**kwargs): + message = SimpleNamespace(content=reply) + return SimpleNamespace(choices=[SimpleNamespace(message=message)]) + + reranker = OpenaiReranker(api_key="test-key") + reranker.__dict__["_client"] = SimpleNamespace( + chat=SimpleNamespace(completions=SimpleNamespace(create=create)) + ) + return reranker + + +def vector_results(): + return pa.table( + { + "id": [1, 2, 3], + "text": ["dogs bark", "cats purr", "birds sing"], + "_distance": [0.1, 0.2, 0.3], + } + ) + + +def test_scores_stay_with_their_rows_when_the_model_reorders(): + reranker = reranker_replying( + [ + {"content": "cats purr", "relevance_score": 0.9}, + {"content": "birds sing", "relevance_score": 0.5}, + {"content": "dogs bark", "relevance_score": 0.1}, + ] + ) + + result = reranker.rerank_vector("which animal purrs?", vector_results()) + + assert result.select(["id", "text"]).to_pylist() == [ + {"id": 2, "text": "cats purr"}, + {"id": 3, "text": "birds sing"}, + {"id": 1, "text": "dogs bark"}, + ] + assert result["_relevance_score"].to_pylist() == pytest.approx([0.9, 0.5, 0.1]) + + +def test_a_row_the_model_leaves_out_raises(): + reranker = reranker_replying( + [ + {"content": "cats purr", "relevance_score": 0.9}, + {"content": "dogs bark", "relevance_score": 0.1}, + ] + ) + + with pytest.raises(ValueError, match="1 of 3"): + reranker.rerank_vector("which animal purrs?", vector_results()) From 0cb5bd0a5d7b96ebafa83dc0f98f98f6af70a189 Mon Sep 17 00:00:00 2001 From: Captain <42566883+po-et@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:23:36 +0800 Subject: [PATCH 53/68] fix(python): accept VoyageAI text batch containers (#4558) Direct VoyageAI `compute_source_embeddings()` calls reject `list[str]` and NumPy text batches before reaching the provider, although both belong to the declared `TEXT` input type. For example, `["first", "second"]` raises `ValueError: Input type not allowed with text model.` Normalize these containers to a list before the existing string-entry validation. The regression cases exercise the public embedding API with a stub provider: input order, list/NumPy batches, Arrow controls, empty batches, and rejection of non-text entries and ordinary two-dimensional arrays. The VoyageAI integration page now documents the accepted text containers. Verified on current main `7b3ba4c` with a locally compiled native extension (0.41.0b7, Python 3.12.12): - VoyageAI suite: 7 failed / 20 passed before the fix; 27 passed afterward. - Adjacent embedding suite: 13 passed, 4 skipped, 8 slow cases deselected. - Python CI's pinned Ruff 0.9.9 lint/format checks pass; configured Pyright reports zero errors/warnings; all applicable pre-commit hooks pass. The provider is stubbed; no live Voyage API request or model-quality evaluation was performed. AI disclosure: I used Codex for investigation, implementation, tests, review, and drafting this description. --- docs/web/integrations/embedding/voyageai.mdx | 4 +- python/python/lancedb/embeddings/voyageai.py | 2 + .../python/tests/test_voyageai_embeddings.py | 57 +++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/docs/web/integrations/embedding/voyageai.mdx b/docs/web/integrations/embedding/voyageai.mdx index cfa795841..c0741b39b 100644 --- a/docs/web/integrations/embedding/voyageai.mdx +++ b/docs/web/integrations/embedding/voyageai.mdx @@ -45,6 +45,8 @@ Supported parameters (to be passed in `create` method) are: | `output_dimension` | `int` | `None` | Output embedding dimension. Only supported by `voyage-multimodal-3.5`. Valid options: 256, 512, 1024 (default), 2048. | +For text models, `compute_source_embeddings()` accepts a single string, a list of strings, a one-dimensional NumPy array, or a PyArrow `Array` or `ChunkedArray`. Batch inputs retain their order. + Usage Example: @@ -61,4 +63,4 @@ The `voyage-multimodal-3.5` model can embed text alongside images. You can use i {PyEmbeddingVoyageaiMultimodal} - \ No newline at end of file + diff --git a/python/python/lancedb/embeddings/voyageai.py b/python/python/lancedb/embeddings/voyageai.py index 4a721c998..95525e07b 100644 --- a/python/python/lancedb/embeddings/voyageai.py +++ b/python/python/lancedb/embeddings/voyageai.py @@ -145,6 +145,8 @@ def sanitize_text_input(inputs: TEXT) -> List[str]: """ if isinstance(inputs, str): inputs = [inputs] + elif isinstance(inputs, (list, np.ndarray)): + inputs = list(inputs) elif isinstance(inputs, pa.Array): inputs = inputs.to_pylist() elif isinstance(inputs, pa.ChunkedArray): diff --git a/python/python/tests/test_voyageai_embeddings.py b/python/python/tests/test_voyageai_embeddings.py index 040cade1f..d49978fb3 100644 --- a/python/python/tests/test_voyageai_embeddings.py +++ b/python/python/tests/test_voyageai_embeddings.py @@ -6,6 +6,8 @@ These tests verify model registration and configuration without requiring API calls. """ +import numpy as np +import pyarrow as pa import pytest from unittest.mock import MagicMock, patch @@ -122,3 +124,58 @@ class TestVoyageAIModelRegistration: assert "voyage-4" not in func.multimodal_embedding_models assert "voyage-4-lite" not in func.multimodal_embedding_models assert "voyage-4-large" not in func.multimodal_embedding_models + + +@pytest.mark.parametrize( + "inputs", + [ + ["first", "second"], + np.array(["first", "second"]), + pa.array(["first", "second"]), + pa.chunked_array([["first"], ["second"]]), + ], + ids=["list", "numpy", "arrow", "chunked-arrow"], +) +def test_text_source_embeddings_accept_batch_containers(inputs): + from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + client = MagicMock() + client.tokenize.return_value = [[1], [2]] + client.embed.return_value.embeddings = [[0.1, 0.2], [0.3, 0.4]] + VoyageAIEmbeddingFunction.client = client + func = get_registry().get("voyageai").create(name="voyage-3") + + assert func.compute_source_embeddings(inputs) == [[0.1, 0.2], [0.3, 0.4]] + client.embed.assert_called_once_with( + texts=["first", "second"], model="voyage-3", input_type="document" + ) + + +@pytest.mark.parametrize("inputs", [[], np.array([], dtype=str)]) +def test_text_source_embeddings_accept_empty_batch(inputs): + from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + client = MagicMock() + VoyageAIEmbeddingFunction.client = client + func = get_registry().get("voyageai").create(name="voyage-3") + + assert func.compute_source_embeddings(inputs) == [] + client.tokenize.assert_not_called() + client.embed.assert_not_called() + + +@pytest.mark.parametrize( + "inputs", + [["text", 42], np.array(["text", None]), np.array([["first", "second"]])], +) +def test_text_source_embeddings_reject_non_text_batch_entries(inputs): + from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + client = MagicMock() + VoyageAIEmbeddingFunction.client = client + func = get_registry().get("voyageai").create(name="voyage-3") + + with pytest.raises(ValueError, match="Each input should be str"): + func.compute_source_embeddings(inputs) + client.tokenize.assert_not_called() + client.embed.assert_not_called() From 118e606e8e06b6aa3ad57839184f655d996ae831 Mon Sep 17 00:00:00 2001 From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com> Date: Fri, 23 Oct 2026 00:53:50 +0530 Subject: [PATCH 54/68] fix(python): align vector parent indices after flattening null lists (#4572) Arrow's `flatten()` skips child values hidden under null list parents, but `list_parent_indices()` can retain their indices. This makes `has_nan_values()` raise `ArrowInvalid` because the parent indices and NaN mask have different lengths. Filter out indices belonging to null parents before applying the mask. The regression test covers `ListArray` and `LargeListArray`, float16/32/64, and chunked inputs: 12 cases. NaN and null-component policies stay the same. The crash also affects #4503. This patch fixes it separately from that PR's changes to bad-vector policies. Validation: - All 12 regression cases failed before the fix. The utility suite now passes with 87 passed and 1 skipped. - Ruff lint and formatting pass for both changed files, as does `git diff --check`. - The patch applies cleanly to main at `ca275b6d`. The patched helper from that revision also passes all 12 cases. - A separate AI reviewer checked commit `343082a6` and found no blocking issues. They reran the utility suite and compared 4,059 generated cases plus 12 explicit null-parent cases with a row-by-row reference implementation. Those checks include slices, chunked inputs, and fixed-size lists. SDK tests used base `07e626ec` with the existing compiled extension. The broader Python run stopped after 474 passed and 23 skipped at `test_local_module_outside_code_warns_at_registration` with `ModuleNotFoundError`. That test also fails with the original vector helper restored. Repository-wide Ruff reports 49 lint errors and 11 files needing formatting, all in untouched documentation examples. Full SDK integration tests against current main require a newer native build and have not been run. --- python/python/lancedb/table.py | 5 +++++ python/python/tests/test_util.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c7de4161d..5eea55760 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5333,6 +5333,11 @@ def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray else: values_has_nan = pc.is_nan(values) values_indices = pc.list_parent_indices(arr) + if arr.null_count: + # Match flatten(), which excludes child spans under null parent lists. + values_indices = pc.filter( + values_indices, pc.take(pc.is_valid(arr), values_indices) + ) has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan)) indices = pa.array(range(len(arr)), type=pa.uint32()) return pc.is_in(indices, has_nan_indices) diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index 4671109b8..4fb16bb31 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -21,6 +21,7 @@ from lancedb.table import ( _infer_target_schema, _merge_metadata, _sanitize_data, + has_nan_values, sanitize_create_table, ) import pyarrow as pa @@ -489,6 +490,20 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("dtype", [pa.float16(), pa.float32(), pa.float64()]) +@pytest.mark.parametrize("large", [False, True]) +@pytest.mark.parametrize("chunked", [False, True]) +def test_has_nan_values_skips_hidden_children_of_null_lists(dtype, large, chunked): + array_type = pa.LargeListArray if large else pa.ListArray + values = pa.array([float("nan"), None, float("nan"), 3.0, 4.0, 5.0], type=dtype) + arr = array_type.from_arrays( + [0, 2, 4, 6], values, mask=pa.array([True, False, False]) + ) + if chunked: + arr = pa.chunked_array([arr.slice(0, 2), arr.slice(2, 1)]) + assert has_nan_values(arr).to_pylist() == [False, True, False] + + def test_handle_bad_vectors_noop(): # ChunkedArray should be preserved as-is vector = pa.chunked_array( From f2c6aa35c236e04b012c150ce94dcc6fc09fd2f7 Mon Sep 17 00:00:00 2001 From: Captain <42566883+po-et@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:24:04 +0800 Subject: [PATCH 55/68] fix(python): normalize VoyageAI image modes before JPEG encoding (#4576) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PIL images, PNG bytes and local image paths in VoyageAI source embeddings are saved directly as JPEG. RGBA, grayscale-with-alpha (LA) and palette (P) images therefore raise `OSError: cannot write mode ... as JPEG` before the provider is called. Convert to RGB at all three JPEG encoding sites. The caller’s original PIL image remains unchanged. Document that alpha is discarded and callers should composite onto a chosen background when it matters. Validation: - New image regressions + existing VoyageAI tests: **18 failed / 26 passed before → 44 passed after**. Tests cover both multimodal models, PIL/bytes/Path, RGBA/LA/P/RGB, decoded JPEG pixels and invalid images. - Independent local LanceDB table ingestion smoke: three PNG modes stored with 1024-dimensional vectors, original bytes preserved. Uses the freshly built current native extension. - Adjacent local embedding tests: 15 passed, 1 skipped (no OpenAI key), 9 deselected. The separate Flickr download test failed with `RemoteDisconnected`; it is excluded from the local result. - Pyright: 0 errors; CI-pinned Ruff 0.9.9 lint and Python formatting passed. External VoyageAI calls are stubbed; real Pillow conversion and LanceDB ingestion execute locally. No live provider or full-repository test result is claimed. AI assistance was used for implementation and validation. --- docs/web/integrations/embedding/voyageai.mdx | 7 +- python/python/lancedb/embeddings/voyageai.py | 6 +- .../tests/test_voyageai_image_inputs.py | 88 +++++++++++++++++++ 3 files changed, 97 insertions(+), 4 deletions(-) create mode 100644 python/python/tests/test_voyageai_image_inputs.py diff --git a/docs/web/integrations/embedding/voyageai.mdx b/docs/web/integrations/embedding/voyageai.mdx index c0741b39b..957ab2b4d 100644 --- a/docs/web/integrations/embedding/voyageai.mdx +++ b/docs/web/integrations/embedding/voyageai.mdx @@ -57,7 +57,12 @@ Usage Example: ### Multimodal Example {#multimodal-example} -The `voyage-multimodal-3.5` model can embed text alongside images. You can use image URLs, file paths, or PIL Image objects: +The `voyage-multimodal-3.5` model can embed text alongside images. You can use image URLs, local `pathlib.Path` objects, raw image bytes, or PIL Image objects. + +For source embeddings, PIL images and images loaded from bytes or local paths are +converted to RGB and encoded as JPEG. This supports RGBA, grayscale with alpha, +and palette images; transparency is discarded. Composite transparent images onto +your chosen background before embedding if the background color matters. diff --git a/python/python/lancedb/embeddings/voyageai.py b/python/python/lancedb/embeddings/voyageai.py index 95525e07b..bcb2fd5a3 100644 --- a/python/python/lancedb/embeddings/voyageai.py +++ b/python/python/lancedb/embeddings/voyageai.py @@ -75,7 +75,7 @@ def transform_input(input_data: Union[str, bytes, Path]): content = {"type": "text", "text": input_data} elif isinstance(input_data, PIL_Image.Image): buffered = BytesIO() - input_data.save(buffered, format="JPEG") + input_data.convert("RGB").save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") content = { "type": "image_base64", @@ -84,7 +84,7 @@ def transform_input(input_data: Union[str, bytes, Path]): elif isinstance(input_data, bytes): img = PIL_Image.open(BytesIO(input_data)) buffered = BytesIO() - img.save(buffered, format="JPEG") + img.convert("RGB").save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") content = { "type": "image_base64", @@ -103,7 +103,7 @@ def transform_input(input_data: Union[str, bytes, Path]): else: img = PIL_Image.open(input_data) buffered = BytesIO() - img.save(buffered, format="JPEG") + img.convert("RGB").save(buffered, format="JPEG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") content = { "type": "image_base64", diff --git a/python/python/tests/test_voyageai_image_inputs.py b/python/python/tests/test_voyageai_image_inputs.py new file mode 100644 index 000000000..885f94821 --- /dev/null +++ b/python/python/tests/test_voyageai_image_inputs.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import base64 +from io import BytesIO +from unittest.mock import MagicMock + +import pytest + +from lancedb.embeddings import get_registry +from lancedb.embeddings.voyageai import VoyageAIEmbeddingFunction + + +@pytest.mark.parametrize("model_name", ["voyage-multimodal-3", "voyage-multimodal-3.5"]) +@pytest.mark.parametrize("input_kind", ["pil", "bytes", "path"]) +@pytest.mark.parametrize( + "mode,color,expected_rgb", + [ + ("RGBA", (64, 128, 192, 128), (64, 128, 192)), + ("LA", (96, 128), (96, 96, 96)), + ("P", 0, (64, 128, 192)), + ("RGB", (64, 128, 192), (64, 128, 192)), + ], +) +def test_multimodal_source_embeddings_encode_image_modes( + model_name, input_kind, mode, color, expected_rgb, tmp_path, monkeypatch +): + Image = pytest.importorskip("PIL.Image") + image = Image.new(mode, (8, 6), color) + if mode == "P": + image.putpalette([64, 128, 192] + [0, 0, 0] * 255) + image.info["transparency"] = 0 + original_pixels = image.tobytes() + image_input = image + if input_kind != "pil": + png = BytesIO() + image.save(png, format="PNG") + if input_kind == "bytes": + image_input = png.getvalue() + else: + image_input = tmp_path / "image.png" + image_input.write_bytes(png.getvalue()) + + client = MagicMock() + client.multimodal_embed.return_value.embeddings = [[0.1, 0.2]] + monkeypatch.setattr(VoyageAIEmbeddingFunction, "client", client) + func = get_registry().get("voyageai").create(name=model_name) + + assert func.compute_source_embeddings(image_input) == [[0.1, 0.2]] + + client.multimodal_embed.assert_called_once() + request = client.multimodal_embed.call_args.kwargs + assert request["model"] == model_name + assert request["input_type"] == "document" + assert len(request["inputs"]) == 1 + (content,) = request["inputs"][0]["content"] + assert content["type"] == "image_base64" + prefix, encoded = content["image_base64"].split(",", 1) + assert prefix == "data:image/jpeg;base64" + with Image.open(BytesIO(base64.b64decode(encoded, validate=True))) as decoded: + assert decoded.format == "JPEG" + assert decoded.mode == "RGB" + assert decoded.size == (8, 6) + assert decoded.getpixel((0, 0)) == pytest.approx(expected_rgb, abs=3) + assert image.mode == mode + assert image.tobytes() == original_pixels + client.embed.assert_not_called() + + +@pytest.mark.parametrize("input_kind", ["bytes", "path"]) +def test_multimodal_source_embeddings_reject_invalid_image( + input_kind, tmp_path, monkeypatch +): + PIL = pytest.importorskip("PIL") + image_input = b"not an image" + if input_kind == "path": + image_input = tmp_path / "invalid.png" + image_input.write_bytes(b"not an image") + + client = MagicMock() + monkeypatch.setattr(VoyageAIEmbeddingFunction, "client", client) + func = get_registry().get("voyageai").create(name="voyage-multimodal-3") + + with pytest.raises(PIL.UnidentifiedImageError): + func.compute_source_embeddings(image_input) + + client.multimodal_embed.assert_not_called() + client.embed.assert_not_called() From 566517be05ff39d3fc6445c5031abccceecf11eb Mon Sep 17 00:00:00 2001 From: Raphael Fakhri Date: Thu, 22 Oct 2026 22:24:23 +0300 Subject: [PATCH 56/68] fix(python): apply pyarrow expression filters in PyarrowDatasetAdapter (#4592) `PyarrowDatasetAdapter` lets DuckDB query a LanceDB table, but DuckDB pushes `WHERE` clauses into `Dataset.scanner(filter=...)` as `pyarrow.compute.Expression` objects. The adapter forwarded them to `Table.count_rows` and `LanceQueryBuilder.where`, which accept only SQL strings, so every filtered DuckDB query failed with `TypeError: argument 'filter': 'Expression' object is not an instance of 'str'`. This change: - Applies a `pyarrow.compute.Expression` filter locally with `pyarrow.dataset.Scanner.from_batches`. The scan reads all columns and pyarrow applies the filter and projection, because DuckDB can filter on a column it does not project. - Applies `head` limits after the filter, so the limit counts matching rows. - Routes `PyarrowDatasetAdapter.count_rows(filter)` through the scanner, so it accepts both filter types. SQL string filters keep the existing path and still run inside LanceDB. Fixes #4591 ### Testing - New `test_dataset_adapter_expression_filter` in `test_pyarrow.py` covers `count_rows`, `to_table`, `head` and a projection that excludes the filtered column. - New `test_query_with_filter` in `test_duckdb.py` runs DuckDB queries with a `WHERE` clause. - Both tests fail without the fix and pass with it. `ruff check` and `ruff format --check` pass. --- python/python/lancedb/integrations/pyarrow.py | 24 ++++++++++++++++++- python/python/tests/test_duckdb.py | 16 +++++++++++++ python/python/tests/test_pyarrow.py | 17 +++++++++++++ 3 files changed, 56 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/integrations/pyarrow.py b/python/python/lancedb/integrations/pyarrow.py index 085964690..1626a6b57 100644 --- a/python/python/lancedb/integrations/pyarrow.py +++ b/python/python/lancedb/integrations/pyarrow.py @@ -52,6 +52,9 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): raise NotImplementedError("memory_pool not supported") def count_rows(self): + if isinstance(self.filter, pa.compute.Expression): + # stream the batches so memory does not grow with the matching rows + return sum(batch.num_rows for batch in self.to_reader()) return self.table.count_rows(self.filter) def from_batches(self, **kwargs): @@ -83,6 +86,8 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): return self.to_reader().read_all() def to_reader(self, *, limit: Optional[int] = None): + if isinstance(self.filter, pa.compute.Expression): + return self._to_reader_with_expression(limit=limit) query = self.table.search() # Disable the builtin limit if limit is None: @@ -98,13 +103,30 @@ class PyarrowScannerAdapter(pa.dataset.Scanner): query = query.where(self.filter, prefilter=True) return query.to_batches(batch_size=self.batch_size) + def _to_reader_with_expression(self, *, limit: Optional[int] = None): + # LanceDB filters are SQL strings, but engines such as DuckDB push down + # pyarrow compute expressions. Scan the table and apply the filter and + # projection with pyarrow, since the filter may reference columns that + # are not part of the projection. + if limit is not None and limit <= 0: + raise ValueError("limit must be positive") + query = self.table.search().limit(self.table.count_rows()) + scanner = pa.dataset.Scanner.from_batches( + query.to_batches(batch_size=self.batch_size), + columns=self.columns, + filter=self.filter, + ) + if limit is None: + return scanner.to_reader() + return scanner.head(limit).to_reader() + class PyarrowDatasetAdapter(pa.dataset.Dataset): def __init__(self, table: Table): self.table = table def count_rows(self, filter: Optional[Filter] = None): - return self.table.count_rows(filter) + return self.scanner(filter=filter).count_rows() def get_fragments(self, filter: Optional[Filter] = None): raise NotImplementedError diff --git a/python/python/tests/test_duckdb.py b/python/python/tests/test_duckdb.py index a307149ea..a24933210 100644 --- a/python/python/tests/test_duckdb.py +++ b/python/python/tests/test_duckdb.py @@ -22,3 +22,19 @@ def test_basic_query(tmp_path): results = duck_conn.sql("SELECT SUM(y) FROM adapter").fetchall() assert results[0][0] == 26 + + +def test_query_with_filter(tmp_path): + data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) + conn = lancedb.connect(tmp_path) + tbl = conn.create_table("test", data) + + adapter = PyarrowDatasetAdapter(tbl) # noqa: F841 + + duck_conn = duckdb.connect() + + results = duck_conn.sql("SELECT SUM(y) FROM adapter WHERE x > 2").fetchall() + assert results[0][0] == 15 + + results = duck_conn.sql("SELECT y FROM adapter WHERE x > 2 ORDER BY y").fetchall() + assert results == [(7,), (8,)] diff --git a/python/python/tests/test_pyarrow.py b/python/python/tests/test_pyarrow.py index 7d99d19c6..272eee2ea 100644 --- a/python/python/tests/test_pyarrow.py +++ b/python/python/tests/test_pyarrow.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import pyarrow as pa +import pyarrow.compute as pc import lancedb from lancedb.integrations.pyarrow import PyarrowDatasetAdapter @@ -48,3 +49,19 @@ def test_dataset_adapter(tmp_path): assert adapter.head(10).num_rows == 0 assert adapter.scanner().projected_schema == pa.schema({"x": pa.int64()}) + + +def test_dataset_adapter_expression_filter(tmp_path): + data = pa.table({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) + conn = lancedb.connect(tmp_path) + tbl = conn.create_table("test", data) + + adapter = PyarrowDatasetAdapter(tbl) + filt = pc.field("x") > 2 + + assert adapter.count_rows(filt) == 2 + assert adapter.to_table(filter=filt) == data.slice(2, 2) + assert adapter.head(1, filter=filt) == data.slice(2, 1) + # The filter can reference a column that is not projected + assert adapter.to_table(columns=["y"], filter=filt) == pa.table({"y": [7, 8]}) + assert adapter.scanner(filter=filt).count_rows() == 2 From a6011dcd9501a2b50a3e9d753bd0b343140d5662 Mon Sep 17 00:00:00 2001 From: Captain <42566883+po-et@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:24:40 +0800 Subject: [PATCH 57/68] fix(python): accept NumPy text arrays in Ollama embeddings (#4605) `OllamaEmbeddings.compute_source_embeddings(np.array(["hello", "world"]))` raises a validation error before sending a request. The embedding interface accepts NumPy arrays, but the adapter passes the array directly to Ollama's string-or-sequence request field. Convert NumPy text arrays to lists before the batched request. This preserves input order, the single request, model options and `keep_alive`. The tests cover Unicode/object/sliced arrays, empty input, existing list/Arrow/query inputs and API error propagation. The added input-contract tests also cover two-dimensional arrays and byte strings, with equivalent list inputs as controls. The real Ollama SDK rejects 2-D inputs and invalid UTF-8 locally, before transport. Valid UTF-8 bytes are decoded normally. No extra adapter validation is needed for these cases. Validation on Python 3.12.14 / macOS arm64: - Built the Rust/PyO3 bindings using the repository's uv/maturin environment. - The initial regression run had 5 failures and 4 passing controls on the base; all 9 passed with the fix. With the 6 additional input-contract cases, the related embedding and trust suites now total 55 passed / 11 skipped. - All Python Ruff lint and formatting checks pass. Root checks still report the same 49 lint errors and 11 documentation formatting files reproduced on the untouched base. - Separate review exercised 13 public API cases. Another 7 groups passed through a real loopback HTTP transport, including local validation, UTF-8 serialization and empty-batch model-error propagation. The regression tests use the real Ollama 0.6.2 SDK with HTTPX MockTransport. The loopback check uses an HTTP surrogate, not an Ollama daemon. No live model inference was run. A live Ollama server was unavailable locally. [The current server handler](https://github.com/ollama/ollama/blob/2cd63f7a15039ccd790e97b31edc091fa1707c56/server/routes.go#L1212) explicitly returns HTTP 200 with empty embeddings for empty input after model lookup and runner setup. The adapter preserves the existing empty-list request and its model/setup errors. This is a source check, not a live-server test. Developed with OpenAI Codex, with separate agent review and executed local checks. --- docs/web/integrations/embedding/ollama.mdx | 3 + python/python/lancedb/embeddings/ollama.py | 2 + python/python/tests/test_ollama_embeddings.py | 126 ++++++++++++++++++ 3 files changed, 131 insertions(+) create mode 100644 python/python/tests/test_ollama_embeddings.py diff --git a/docs/web/integrations/embedding/ollama.mdx b/docs/web/integrations/embedding/ollama.mdx index 120cde8f2..bd161a692 100644 --- a/docs/web/integrations/embedding/ollama.mdx +++ b/docs/web/integrations/embedding/ollama.mdx @@ -10,6 +10,9 @@ Generate embeddings via the [ollama](https://github.com/ollama/ollama-python) py - [Ollama docs on embeddings](https://github.com/ollama/ollama/blob/main/docs/api.md#generate-embeddings) - [Ollama blog on embeddings](https://ollama.com/blog/embedding-models) +For direct calls, `compute_source_embeddings` accepts a list of strings or a +one-dimensional NumPy string array. Both inputs are sent to Ollama in one batch. + | Parameter | Type | Default Value | Description | |------------------------|----------------------------|--------------------------|------------------------------------------------------------------------------------------------------------------------------------------------| | `name` | `str` | `nomic-embed-text` | The name of the model. | diff --git a/python/python/lancedb/embeddings/ollama.py b/python/python/lancedb/embeddings/ollama.py index aad8742de..3151c7c72 100644 --- a/python/python/lancedb/embeddings/ollama.py +++ b/python/python/lancedb/embeddings/ollama.py @@ -70,6 +70,8 @@ class OllamaEmbeddings(TextEmbeddingFunction): The texts to embed """ # TODO retry, rate limit, token limit + if isinstance(texts, np.ndarray): + texts = texts.tolist() embeddings = self._compute_embedding(texts) return list(embeddings) diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py new file mode 100644 index 000000000..9d8568ca7 --- /dev/null +++ b/python/python/tests/test_ollama_embeddings.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import json + +import httpx +import numpy as np +import pyarrow as pa +import pytest +from pydantic import ValidationError + +from lancedb.embeddings import get_registry + + +@pytest.fixture +def ollama_embeddings(): + pytest.importorskip("ollama") + requests = [] + + def respond(request): + body = json.loads(request.content) + requests.append(body) + if body["model"] == "missing-model": + return httpx.Response(404, json={"error": "model not found"}) + return httpx.Response( + 200, + json={ + "embeddings": [ + [float(index), float(len(text))] + for index, text in enumerate(body["input"]) + ] + }, + ) + + function = ( + get_registry() + .get("ollama") + .create( + name="nomic-embed-text", + options={"num_ctx": 4096}, + keep_alive="10m", + ollama_client_kwargs={ + "transport": httpx.MockTransport(respond), + "trust_env": False, + }, + ) + ) + yield function, requests + function._ollama_client._client.close() + + +@pytest.mark.parametrize( + "texts", + [ + pytest.param(["hello", "世界"], id="list"), + pytest.param(np.array(["hello", "世界"]), id="numpy-unicode"), + pytest.param(np.array(["hello", "世界"], dtype=object), id="numpy-object"), + pytest.param( + np.array([b"hello", "世界".encode("utf-8")]), id="numpy-utf8-bytes" + ), + pytest.param([b"hello", "世界".encode("utf-8")], id="list-utf8-bytes"), + pytest.param(np.array(["hello", "unused", "世界"])[::2], id="numpy-slice"), + pytest.param(pa.array(["hello", "世界"]), id="arrow"), + pytest.param(pa.chunked_array([["hello"], ["世界"]]), id="chunked-arrow"), + ], +) +def test_ollama_source_embeddings_inputs(ollama_embeddings, texts): + function, requests = ollama_embeddings + + assert function.compute_source_embeddings(texts) == [[0.0, 5.0], [1.0, 2.0]] + assert requests == [ + { + "model": "nomic-embed-text", + "input": ["hello", "世界"], + "options": {"num_ctx": 4096}, + "keep_alive": "10m", + } + ] + + +@pytest.mark.parametrize( + "texts", + [ + pytest.param(np.array([["hello"], ["world"]]), id="numpy-two-dimensional"), + pytest.param([["hello"], ["world"]], id="list-two-dimensional"), + pytest.param(np.array([b"\xff"]), id="numpy-invalid-utf8"), + pytest.param([b"\xff"], id="list-invalid-utf8"), + ], +) +def test_ollama_source_embeddings_rejects_invalid_input_locally( + ollama_embeddings, texts +): + function, requests = ollama_embeddings + + with pytest.raises(ValidationError, match="valid string"): + function.compute_source_embeddings(texts) + + assert requests == [] + + +def test_ollama_generate_embeddings_empty_array(ollama_embeddings): + function, requests = ollama_embeddings + + assert function.generate_embeddings(np.array([], dtype=str)) == [] + assert len(requests) == 1 + assert requests[0]["input"] == [] + + +def test_ollama_query_embeddings(ollama_embeddings): + function, requests = ollama_embeddings + + assert function.compute_query_embeddings("hello") == [[0.0, 5.0]] + assert len(requests) == 1 + assert requests[0]["input"] == ["hello"] + + +def test_ollama_source_embeddings_preserves_api_error(ollama_embeddings): + ollama = pytest.importorskip("ollama") + function, requests = ollama_embeddings + function.name = "missing-model" + + with pytest.raises(ollama.ResponseError, match="model not found") as error: + function.compute_source_embeddings(np.array(["hello"])) + + assert error.value.status_code == 404 + assert len(requests) == 1 From e5873d976350d74d7d276794abb8bde1c45fcb47 Mon Sep 17 00:00:00 2001 From: Ishita Singh Date: Fri, 23 Oct 2026 00:54:56 +0530 Subject: [PATCH 58/68] fix(node): support nested blobs with blobMode 'bytes' (#4610) Closes #4609. `blobMode: "bytes"` only resolved top-level blob columns, so a blob nested in a struct came back as its descriptor. Now struct projections that contain blob columns are resolved too: the bytes are fetched per blob column and the struct is rebuilt with `LargeBinary` leaves, preserving struct validity so a null struct stays null. Sibling fields pass through untouched, aliases work for both `info.image AS image` and `info AS details`, and blobs nested in lists still keep their descriptors. Tested with `npx tsc`, `biome check`, and the full `table.test.ts` suite (580 tests, including new regressions for select-all, struct selection, aliases, null structs, null and empty blobs, sibling fields, empty results, and single-row batches). Not run against remote tables. --- .../js/interfaces/QueryExecutionOptions.md | 11 +- nodejs/__test__/table.test.ts | 125 ++++++++++++++ nodejs/lancedb/arrow.ts | 1 + nodejs/lancedb/query.ts | 154 ++++++++++++++---- 4 files changed, 258 insertions(+), 33 deletions(-) diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index 8d3749f38..c0b33dcf0 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,12 +20,13 @@ How to return blob v2 columns. Applies to `toArrow()` and `toArray()`. - `"descriptions"` (default): each blob is a descriptor. Read the bytes with [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles). -- `"bytes"`: each top-level blob column holds the blob bytes as - `LargeBinary`. The query also reads `_rowid` to fetch the bytes, and - leaves it out of the result unless [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was - called. All the bytes are held in memory. +- `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, + including blobs nested in structs. The query also reads `_rowid` to + fetch the bytes, and leaves it out of the result unless + [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was called. All the bytes are held in + memory. -Blobs nested in a struct or a list keep their descriptors. +Blobs nested in a list keep their descriptors. The rows and the bytes come from the same table version. If the table changes while the query runs, the query runs again, up to three times. diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 84340c72c..92e3df456 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3385,6 +3385,106 @@ describe("when dealing with blob columns", () => { expect(bytes[0]!.equals(payload)).toBe(true); }); + it("returns bytes for blobs nested in structs", async () => { + const { table, alpha, empty } = await openNestedBlobTable(); + const rows = await table.query().toArray({ blobMode: "bytes" }); + const byId = new Map(rows.map((row) => [Number(row.id), row])); + expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha); + expect(byId.get(1)!.info.caption).toBe("one"); + expect(byId.get(2)!.info.image).toBeNull(); + expect(byId.get(2)!.info.caption).toBe("two"); + expect(byId.get(3)!.info).toBeNull(); + expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty); + }); + + it("lines up nested rows and bytes across single-row batches", async () => { + const { table, alpha, empty } = await openNestedBlobTable(); + const result = await table + .query() + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.batches.length).toBe(4); + const byId = new Map(result.toArray().map((row) => [Number(row.id), row])); + expect(Buffer.from(byId.get(1)!.info.image)).toEqual(alpha); + expect(byId.get(1)!.info.caption).toBe("one"); + expect(byId.get(2)!.info.image).toBeNull(); + expect(byId.get(2)!.info.caption).toBe("two"); + expect(byId.get(3)!.info).toBeNull(); + expect(Buffer.from(byId.get(4)!.info.image)).toEqual(empty); + }); + + it("handles struct fields named like Object prototypes", async () => { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + new Field( + "info", + new Struct([ + new Field("constructor", new Struct([blob("image")]), true), + ]), + true, + ), + ]); + const payload = Buffer.from("proto"); + const table = await db.createTable( + "prototype_named_blobs", + [{ id: 1n, info: { constructor: { image: payload } } }], + { schema }, + ); + expect(await table.blobColumns()).toEqual(["info.constructor.image"]); + const [row] = await table.query().toArray({ blobMode: "bytes" }); + // Arrow's StructRow shadows the `constructor` property, so read the + // field through the JSON form instead of property access. + const info = JSON.parse(JSON.stringify(row.info)); + expect(Buffer.from(Object.values(info.constructor.image))).toEqual( + payload, + ); + }); + + it("returns nested blob bytes for a struct selection", async () => { + const { table, alpha } = await openNestedBlobTable(); + const [row] = await table + .query() + .where("id = 1") + .select(["info"]) + .toArray({ blobMode: "bytes" }); + expect(Object.keys(row)).toEqual(["info"]); + expect(Buffer.from(row.info.image)).toEqual(alpha); + expect(row.info.caption).toBe("one"); + }); + + it("returns nested blob bytes through aliases", async () => { + const { table, alpha } = await openNestedBlobTable(); + const [row] = await table + .query() + .where("id = 1") + .select( + new Map([ + ["picture", "info.image"], + ["details", "info"], + ]), + ) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.picture)).toEqual(alpha); + expect(Buffer.from(row.details.image)).toEqual(alpha); + expect(row.details.caption).toBe("one"); + }); + + it("returns a nested bytes schema for an empty result", async () => { + const { table } = await openNestedBlobTable(); + const result = await table + .query() + .where("id > 100") + .toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + const field = result.schema.fields.find((f) => f.name === "info")!; + expect(arrow.DataType.isStruct(field.type)).toBe(true); + const image = (field.type as Struct).children.find( + (f) => f.name === "image", + )!; + expect(arrow.DataType.isLargeBinary(image.type)).toBe(true); + expect(isBlobField(image)).toBe(false); + }); + it("creates list blobs without advertising unfetchable elements", async () => { const db = await connect(tmpDir.name); const schema = new Schema([ @@ -3593,6 +3693,31 @@ describe("when dealing with blob columns", () => { return { table, rowIds, alpha, beta }; } + async function openNestedBlobTable() { + const db = await connect(tmpDir.name); + const schema = new Schema([ + new Field("id", new Int64(), true), + new Field( + "info", + new Struct([blob("image"), new Field("caption", new Utf8(), true)]), + true, + ), + ]); + const alpha = Buffer.from("alpha"); + const empty = Buffer.alloc(0); + const table = await db.createTable( + "nested_blobs", + [{ id: 1n, info: { image: alpha, caption: "one" } }], + { schema }, + ); + await table.add([ + { id: 2n, info: { image: null, caption: "two" } }, + { id: 3n, info: null }, + { id: 4n, info: { image: empty, caption: "four" } }, + ]); + return { table, alpha, empty }; + } + async function openExternalBlobTable() { const db = await connect(path.join(tmpDir.name, "db")); const payload = Buffer.from("external-payload"); diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index 68b35e256..56c8e1897 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -61,6 +61,7 @@ import { import { inferSchema } from "./schema"; export * from "apache-arrow"; +export type { Data as ArrowData } from "apache-arrow"; export type SchemaLike = | Schema | { diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 2abe6db97..823239ef9 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -2,7 +2,9 @@ // SPDX-FileCopyrightText: Copyright The LanceDB Authors import { + type ArrowData, Table as ArrowTable, + DataType, Field, type IntoVector, LargeBinary, @@ -70,16 +72,80 @@ function blobSources( blobColumns: string[], projection: [string, string][] | undefined, ): Map { - const columns = new Set(blobColumns); + const holds = (source: string) => + blobColumns.some( + (column) => column === source || column.startsWith(`${source}.`), + ); const pairs = projection ?? - blobColumns.map((column): [string, string] => [column, column]); - return new Map(pairs.filter(([, source]) => columns.has(source))); + Array.from(new Set(blobColumns.map((column) => column.split(".")[0]))).map( + (name): [string, string] => [name, name], + ); + return new Map(pairs.filter(([, source]) => holds(source))); +} + +// Fetched bytes for the blobs under one output column. A blob column is a +// leaf; a struct holds the blobs nested in it, keyed by field name. A Map +// keeps struct field names from reaching shared JavaScript prototypes. +type BlobBytes = (Buffer | null)[] | Map; + +function isLeaf(bytes: BlobBytes): bytes is (Buffer | null)[] { + return Array.isArray(bytes); +} + +function nestedFieldType(type: Struct, bytes: Map): Struct { + return new Struct( + type.children.map((field) => { + const inner = bytes.get(field.name); + if (inner === undefined) { + return field; + } + if (isLeaf(inner)) { + return new Field(field.name, new LargeBinary(), true); + } + return new Field( + field.name, + nestedFieldType(field.type as Struct, inner), + field.nullable, + ); + }), + ); +} + +// Rebuilds a struct with the blobs under it as bytes. The struct keeps its own +// validity, so a null struct stays null. +function nestedData( + data: ArrowData, + bytes: Map, + offset: number, +): ArrowData { + const type = data.type as Struct; + const children = data.children.map((child, index) => { + const inner = bytes.get(type.children[index].name); + if (inner === undefined) { + return child; + } + if (isLeaf(inner)) { + return vectorFromArray( + inner.slice(offset, offset + data.length), + new LargeBinary(), + ).data[0]; + } + return nestedData(child, inner, offset); + }); + return makeData({ + type: nestedFieldType(type, bytes), + length: data.length, + nullCount: data.nullCount, + nullBitmap: data.nullBitmap, + children, + }); } async function readBlobBytes( table: ArrowTable, sources: Map, + blobColumns: string[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -91,7 +157,11 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const replacements = new Map(); + const fetch = async (column: string) => + (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( + (value) => value ?? null, + ); + const replacements = new Map(); for (const [output, source] of sources) { const index = table.schema.fields.findIndex( (field) => field.name === output, @@ -99,12 +169,30 @@ async function readBlobBytes( if (index < 0) { continue; } - const values = - rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); - replacements.set( - index, - values.map((value) => value ?? null), - ); + if (blobColumns.includes(source)) { + replacements.set(index, await fetch(source)); + continue; + } + if (!DataType.isStruct(table.schema.fields[index].type)) { + continue; + } + const nested = new Map(); + for (const column of blobColumns.filter((c) => + c.startsWith(`${source}.`), + )) { + const path = column.slice(source.length + 1).split("."); + let level = nested; + for (const name of path.slice(0, -1)) { + let next = level.get(name); + if (next === undefined) { + next = new Map(); + level.set(name, next); + } + level = next as Map; + } + level.set(path[path.length - 1], await fetch(column)); + } + replacements.set(index, nested); } const result = withBinaryColumns(table, replacements); if (!dropRowId) { @@ -121,26 +209,37 @@ async function readBlobBytes( // the bytes. function withBinaryColumns( table: ArrowTable, - replacements: Map, + replacements: Map, ): ArrowTable { if (replacements.size === 0) { return table; } - const fields = table.schema.fields.map((field, index) => - replacements.has(index) + const fields = table.schema.fields.map((field, index) => { + const bytes = replacements.get(index); + if (bytes === undefined) { + return field; + } + return isLeaf(bytes) ? new Field(field.name, new LargeBinary(), true) - : field, - ); + : new Field( + field.name, + nestedFieldType(field.type as Struct, bytes), + field.nullable, + ); + }); const schema = new Schema(fields, table.schema.metadata); let offset = 0; const batches = table.batches.map((batch) => { const children = batch.data.children.map((child, index) => { - const values = replacements.get(index); - if (values === undefined) { + const bytes = replacements.get(index); + if (bytes === undefined) { return child; } + if (!isLeaf(bytes)) { + return nestedData(child, bytes, offset); + } return vectorFromArray( - values.slice(offset, offset + batch.numRows), + bytes.slice(offset, offset + batch.numRows), new LargeBinary(), ).data[0]; }); @@ -175,12 +274,13 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each top-level blob column holds the blob bytes as - * `LargeBinary`. The query also reads `_rowid` to fetch the bytes, and - * leaves it out of the result unless {@link QueryBase.withRowId} was - * called. All the bytes are held in memory. + * - `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, + * including blobs nested in structs. The query also reads `_rowid` to + * fetch the bytes, and leaves it out of the result unless + * {@link QueryBase.withRowId} was called. All the bytes are held in + * memory. * - * Blobs nested in a struct or a list keep their descriptors. + * Blobs nested in a list keep their descriptors. * * The rows and the bytes come from the same table version. If the table * changes while the query runs, the query runs again, up to three times. @@ -429,17 +529,15 @@ export class QueryBase< // unchanged version afterwards proves it read the snapshot's version. const snapshot = await context.table.querySnapshot(); const version = await snapshot.version(); - const sources = blobSources( - await snapshot.blobColumns(), - context.projection, - ); + const blobColumns = await snapshot.blobColumns(); + const sources = blobSources(blobColumns, context.projection); if (sources.size === 0) { return this.collect(options, false); } const addRowId = !context.withRowId; const table = await this.collect(options, addRowId); if ((await context.table.version()) === version) { - return readBlobBytes(table, sources, snapshot, addRowId); + return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error( From f5f4463b15bf1e6b08725da26f05cf8dd5188888 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 03:42:22 +0800 Subject: [PATCH 59/68] fix(node): materialize nested struct blobs in bytes mode (#4611) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `blobMode: "bytes"` left blobs inside structs as descriptors because projection matching ignored parent structs and byte replacement only rebuilt top-level fields. Expand direct struct projections and aliases to their blob leaves, then recursively replace the corresponding Arrow fields and data with `LargeBinary`. Track literal output names separately from descendant paths so scalar aliases stay intact and every overlapping blob projection is converted. Preserve null structs, null and empty blobs, sibling fields, schema metadata, and the existing row ID and snapshot behavior. Blobs inside lists retain their descriptors. Update the API docs to describe struct support. Regression coverage includes full and parent-struct projections, leaf and struct aliases, dotted output names, overlapping projections in both orders, nested structs, nulls, empty blobs and results, multiple batches, list siblings, query reuse, row IDs, and concurrent writes. Regression tests reproduce the nested-blob and projection-collision failures before the corresponding changes. Validation: - `pnpm build` - `pnpm lint` - `pnpm run docs` - `pnpm test __test__/nested_blob.test.ts __test__/blob.test.ts __test__/query.test.ts --runInBand --verbose=false` — 142 passed - `pnpm test __test__/table.test.ts --runInBand --verbose=false -t 'when dealing with blob columns'` — 60 passed Fixes #4609 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- .../js/interfaces/QueryExecutionOptions.md | 13 +- nodejs/__test__/nested_blob.test.ts | 327 ++++++++++++++++++ nodejs/__test__/table.test.ts | 125 ++++--- nodejs/lancedb/query.ts | 259 +++++++------- 4 files changed, 528 insertions(+), 196 deletions(-) create mode 100644 nodejs/__test__/nested_blob.test.ts diff --git a/docs/src/js/interfaces/QueryExecutionOptions.md b/docs/src/js/interfaces/QueryExecutionOptions.md index c0b33dcf0..f18ada7aa 100644 --- a/docs/src/js/interfaces/QueryExecutionOptions.md +++ b/docs/src/js/interfaces/QueryExecutionOptions.md @@ -20,13 +20,14 @@ How to return blob v2 columns. Applies to `toArrow()` and `toArray()`. - `"descriptions"` (default): each blob is a descriptor. Read the bytes with [Table.fetchBlobs](../classes/Table.md#fetchblobs) or [Table.fetchBlobFiles](../classes/Table.md#fetchblobfiles). -- `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - including blobs nested in structs. The query also reads `_rowid` to - fetch the bytes, and leaves it out of the result unless - [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was called. All the bytes are held in - memory. +- `"bytes"`: each blob column, including blobs nested in structs, holds + the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + leaves it out of the result unless [QueryBase.withRowId](../classes/QueryBase.md#withrowid) was + called. All the bytes are held in memory. -Blobs nested in a list keep their descriptors. +Selecting a parent struct or directly aliasing a blob or struct also returns +bytes. Blobs inside lists keep their descriptors. Arbitrary projection +expressions are not materialized as bytes. The rows and the bytes come from the same table version. If the table changes while the query runs, the query runs again, up to three times. diff --git a/nodejs/__test__/nested_blob.test.ts b/nodejs/__test__/nested_blob.test.ts new file mode 100644 index 000000000..508b57c77 --- /dev/null +++ b/nodejs/__test__/nested_blob.test.ts @@ -0,0 +1,327 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import * as tmp from "tmp"; +import { type Table, blob, connect, isBlobField } from "../lancedb"; +import { + DataType, + Field, + Int64, + List, + Schema, + Struct, + Utf8, +} from "../lancedb/arrow"; + +describe('nested blobs with blobMode "bytes"', () => { + let tmpDir: tmp.DirResult; + let table: Table; + + const caption = new Field( + "caption", + new Utf8(), + false, + new Map([["label", "caption"]]), + ); + const info = new Field( + "info", + new Struct([ + blob("image"), + caption, + new Field("preview", new Struct([blob("image")]), true), + ]), + true, + new Map([["label", "info"]]), + ); + const schema = new Schema( + [new Field("id", new Int64(), false), info, blob("cover")], + new Map([["source", "nested-blobs"]]), + ); + + beforeEach(async () => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + const db = await connect(tmpDir.name); + table = await db.createTable( + "nested_blobs", + [ + { + id: 1n, + info: { + image: Buffer.from("hello"), + caption: "first", + preview: { image: Buffer.from("preview") }, + }, + cover: Buffer.from("cover"), + }, + { id: 2n, info: null, cover: null }, + { + id: 3n, + info: { image: null, caption: "third", preview: null }, + cover: null, + }, + { + id: 4n, + info: { + image: Buffer.alloc(0), + caption: "fourth", + preview: { image: null }, + }, + cover: Buffer.alloc(0), + }, + ], + { schema }, + ); + }); + + afterEach(() => tmpDir.removeCallback()); + + it.each([false, true])( + "reads struct leaves across batches (select parent: %s)", + async (selectParent) => { + const query = table + .query() + .orderBy({ columnName: "id", ascending: true }); + if (selectParent) { + query.select(["info"]); + } + const result = await query.toArrow({ + blobMode: "bytes", + maxBatchLength: 1, + }); + expect(result.batches.length).toBeGreaterThan(1); + expect(result.schema.fields.map((field) => field.name)).toEqual( + selectParent ? ["info"] : ["id", "info", "cover"], + ); + expect(result.schema.metadata).toEqual(schema.metadata); + const infoField = result.schema.fields.find( + (field) => field.name === "info", + )!; + expect(infoField.nullable).toBe(true); + expect(infoField.metadata).toEqual(info.metadata); + const fields = (infoField.type as Struct).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[0].metadata.size).toBe(0); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + + const rows = result.toArray(); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(Buffer.from(rows[0].info.preview.image).toString()).toBe( + "preview", + ); + expect(rows[1].info).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(rows[2].info.caption).toBe("third"); + expect(rows[2].info.preview).toBeNull(); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + expect(rows[3].info.preview.image).toBeNull(); + if (!selectParent) { + expect(Buffer.from(rows[0].cover).toString()).toBe("cover"); + expect(rows[1].cover).toBeNull(); + } + }, + ); + + it("reads direct leaf and struct aliases", async () => { + const rows = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select({ + image: "info.image", + details: "info", + preview: "info.preview", + caption: "info.caption", + }) + .toArray({ blobMode: "bytes", maxBatchLength: 1 }); + expect(Object.keys(rows[0])).toEqual([ + "image", + "details", + "preview", + "caption", + ]); + expect(Buffer.from(rows[0].image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.image).toString()).toBe("hello"); + expect(Buffer.from(rows[0].details.preview.image).toString()).toBe( + "preview", + ); + expect(Buffer.from(rows[0].preview.image).toString()).toBe("preview"); + expect(rows[0].caption).toBe("first"); + expect(rows[1].details).toBeNull(); + expect(rows[1].image).toBeNull(); + expect(rows[2].image).toBeNull(); + expect(Buffer.from(rows[3].image)).toEqual(Buffer.alloc(0)); + }); + + it.each([false, true])( + "preserves a literal scalar alias matching a nested blob path (reverse: %s)", + async (reverse) => { + const projection: [string, string][] = [ + ["info", "info"], + ["info.image", "id"], + ]; + if (reverse) { + projection.reverse(); + } + const [row] = await table + .query() + .where("id = 1") + .select(new Map(projection)) + .toArray({ blobMode: "bytes" }); + expect(row["info.image"]).toBe(1n); + expect(Buffer.from(row.info.image).toString()).toBe("hello"); + }, + ); + + it.each([false, true])( + "materializes every overlapping blob projection (reverse: %s)", + async (reverse) => { + const columns = ["info.image", "info"]; + if (reverse) { + columns.reverse(); + } + const result = await table + .query() + .orderBy({ columnName: "id", ascending: true }) + .select(columns) + .toArrow({ blobMode: "bytes", maxBatchLength: 1 }); + expect(result.schema.fields.map((field) => field.name)).toEqual(columns); + expect(DataType.isLargeBinary(result.getChild("info.image")!.type)).toBe( + true, + ); + expect( + DataType.isLargeBinary( + result.getChild("info")!.getChild("image")!.type, + ), + ).toBe(true); + const rows = result.toArray(); + expect(Buffer.from(rows[0]["info.image"]).toString()).toBe("hello"); + expect(Buffer.from(rows[0].info.image).toString()).toBe("hello"); + expect(rows[0].info.caption).toBe("first"); + expect(rows[1]["info.image"]).toBeNull(); + expect(rows[1].info).toBeNull(); + expect(rows[2]["info.image"]).toBeNull(); + expect(rows[2].info.image).toBeNull(); + expect(Buffer.from(rows[3]["info.image"])).toEqual(Buffer.alloc(0)); + expect(Buffer.from(rows[3].info.image)).toEqual(Buffer.alloc(0)); + }, + ); + + it("materializes dotted aliases independently of nested fields", async () => { + const [row] = await table + .query() + .where("id = 1") + .select({ + details: "info", + "details.image": "info.image", + "details.preview": "info.preview", + "details.preview.image": "info.preview.image", + }) + .toArray({ blobMode: "bytes" }); + expect(Buffer.from(row.details.image).toString()).toBe("hello"); + expect(Buffer.from(row["details.image"]).toString()).toBe("hello"); + expect(Buffer.from(row.details.preview.image).toString()).toBe("preview"); + expect(Buffer.from(row["details.preview"].image).toString()).toBe( + "preview", + ); + expect(Buffer.from(row["details.preview.image"]).toString()).toBe( + "preview", + ); + }); + + it.each([false, true])( + "converts the empty result schema (alias: %s)", + async (alias) => { + const query = table.query().where("id < 0"); + if (alias) { + query.select({ image: "info.image", details: "info" }); + } + const result = await query.toArrow({ blobMode: "bytes" }); + expect(result.numRows).toBe(0); + expect(result.schema.fields.map((field) => field.name)).toEqual( + alias ? ["image", "details"] : ["id", "info", "cover"], + ); + const fields = ( + result.schema.fields.find( + (field) => field.name === (alias ? "details" : "info"), + )!.type as Struct + ).children; + expect(DataType.isLargeBinary(fields[0].type)).toBe(true); + expect(isBlobField(fields[0])).toBe(false); + expect(fields[1]).toEqual(caption); + expect( + DataType.isLargeBinary((fields[2].type as Struct).children[0].type), + ).toBe(true); + if (alias) { + expect(DataType.isLargeBinary(result.schema.fields[0].type)).toBe(true); + } + }, + ); + + it("preserves requested row IDs and leaves subsequent descriptors intact", async () => { + const query = table.query().where("id = 1").select(["info"]); + const [bytes] = await query.toArray({ blobMode: "bytes" }); + expect(Object.keys(bytes)).toEqual(["info"]); + const [description] = await query.toArray(); + expect(Object.keys(description)).toEqual(["info"]); + expect(Number(description.info.image.size)).toBe(5); + const [withRowId] = await query.withRowId().toArray({ blobMode: "bytes" }); + expect(typeof withRowId._rowid).toBe("bigint"); + expect(await table.fetchBlobs("info.image", [withRowId._rowid])).toEqual([ + Buffer.from(withRowId.info.image), + ]); + }); + + it("keeps descriptors in description mode", async () => { + const [description] = await table + .query() + .where("id = 1") + .toArray({ blobMode: "descriptions" }); + expect(Number(description.info.image.size)).toBe(5); + }); + + it("keeps list elements as descriptors alongside struct blob bytes", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable( + "list_blobs", + [ + { + info: { image: Buffer.from("hello"), images: [Buffer.from("list")] }, + }, + ], + { + schema: new Schema([ + new Field( + "info", + new Struct([ + blob("image"), + new Field("images", new List(blob("item")), true), + ]), + true, + ), + ]), + }, + ); + const [bytes] = await table.query().toArray({ blobMode: "bytes" }); + expect(Buffer.from(bytes.info.image).toString()).toBe("hello"); + expect(Number(bytes.info.images.get(0).size)).toBe(4); + const [row] = await table + .query() + .select(["info.images"]) + .toArray({ blobMode: "bytes" }); + expect(Number(row["info.images"].get(0).size)).toBe(4); + expect(Object.keys(row)).toEqual(["info.images"]); + const result = await table.query().toArrow({ blobMode: "bytes" }); + const fields = (result.schema.fields[0].type as Struct).children; + const descriptions = await table.query().toArrow(); + expect(fields[1]).toEqual( + (descriptions.schema.fields[0].type as Struct).children[1], + ); + expect(DataType.isStruct((fields[1].type as List).children[0].type)).toBe( + true, + ); + }); +}); diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 92e3df456..d848676cd 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -3314,51 +3314,64 @@ describe("when dealing with blob columns", () => { ).rejects.toThrow(/blobMode must be/); }); - it("reruns the query when the table changes before the version check", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - await update(2n, "new"); - return liveVersion(); - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(2n); - expect(Buffer.from(row.image).toString()).toBe("new"); - }); + it.each([false, true])( + "reruns the query when the table changes before the version check (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + await update(2n, "new"); + return liveVersion(); + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(2n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "new", + ); + }, + ); - it("reads bytes from the query's version when a write lands before the fetch", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - native.version = async () => { - native.version = liveVersion; - const version = await liveVersion(); - await update(2n, "new"); - return version; - }; - const [row] = await table.query().toArray({ blobMode: "bytes" }); - expect(row.generation).toBe(1n); - expect(Buffer.from(row.image).toString()).toBe("old"); - }); + it.each([false, true])( + "reads bytes from the query's version when a write lands before the fetch (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + native.version = async () => { + native.version = liveVersion; + const version = await liveVersion(); + await update(2n, "new"); + return version; + }; + const [row] = await table.query().toArray({ blobMode: "bytes" }); + expect(row.generation).toBe(1n); + expect(Buffer.from(nested ? row.info.image : row.image).toString()).toBe( + "old", + ); + }, + ); - it("fails when the table changes during every attempt", async () => { - const { table, native, update } = await openVersionedBlobTable(); - const liveVersion = native.version.bind(native); - let generation = 1n; - native.version = async () => { - generation += 1n; - await update(generation, `generation ${generation}`); - return liveVersion(); - }; - try { - await expect( - table.query().toArray({ blobMode: "bytes" }), - ).rejects.toThrow(/changed during each of 3 attempts/); - expect(generation).toBe(4n); - } finally { - native.version = liveVersion; - } - }); + it.each([false, true])( + "fails when the table changes during every attempt (nested: %s)", + async (nested) => { + const { table, native, update } = await openVersionedBlobTable(nested); + const liveVersion = native.version.bind(native); + let generation = 1n; + native.version = async () => { + generation += 1n; + await update(generation, `generation ${generation}`); + return liveVersion(); + }; + try { + await expect( + table.query().toArray({ blobMode: "bytes" }), + ).rejects.toThrow(/changed during each of 3 attempts/); + expect(generation).toBe(4n); + } finally { + native.version = liveVersion; + } + }, + ); it("rejects fetchBlobs on a non-blob column", async () => { const { table, rowIds } = await openBlobTable(); @@ -3642,18 +3655,28 @@ describe("when dealing with blob columns", () => { ).map((value) => (value == null ? null : Number(value.size))); } - async function openVersionedBlobTable() { + async function openVersionedBlobTable(nested = false) { const db = await connect(tmpDir.name, { storageOptions: { newTableEnableStableRowIds: "true" }, }); const schema = new Schema([ new Field("id", new Int64(), false), new Field("generation", new Int64(), true), - blob("image"), + nested + ? new Field("info", new Struct([blob("image")]), true) + : blob("image"), ]); const table = await db.createTable( "versions", - [{ id: 1n, generation: 1n, image: Buffer.from("old") }], + [ + { + id: 1n, + generation: 1n, + ...(nested + ? { info: { image: Buffer.from("old") } } + : { image: Buffer.from("old") }), + }, + ], { schema }, ); const native = ( @@ -3663,7 +3686,15 @@ describe("when dealing with blob columns", () => { await table .mergeInsert("id") .whenMatchedUpdateAll() - .execute([{ id: 1n, generation, image: Buffer.from(image) }]); + .execute([ + { + id: 1n, + generation, + ...(nested + ? { info: { image: Buffer.from(image) } } + : { image: Buffer.from(image) }), + }, + ]); }; return { table, native, update }; } diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 823239ef9..b96e4fda6 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -1,8 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +import type { Data as ArrowData } from "apache-arrow"; import { - type ArrowData, Table as ArrowTable, DataType, Field, @@ -14,7 +14,6 @@ import { extractVectorBuffer, fromBufferToRecordBatch, fromRecordBatchToBuffer, - makeData, tableFromIPC, vectorFromArray, } from "./arrow"; @@ -68,84 +67,42 @@ export interface QueryBlobContext { const BLOB_BYTES_ATTEMPTS = 3; +interface BlobSource { + // Output names are literal: an alias containing dots is still one field. + output: string; + path: string[]; + source: string; +} + function blobSources( blobColumns: string[], projection: [string, string][] | undefined, -): Map { - const holds = (source: string) => - blobColumns.some( - (column) => column === source || column.startsWith(`${source}.`), - ); - const pairs = - projection ?? - Array.from(new Set(blobColumns.map((column) => column.split(".")[0]))).map( - (name): [string, string] => [name, name], - ); - return new Map(pairs.filter(([, source]) => holds(source))); -} - -// Fetched bytes for the blobs under one output column. A blob column is a -// leaf; a struct holds the blobs nested in it, keyed by field name. A Map -// keeps struct field names from reaching shared JavaScript prototypes. -type BlobBytes = (Buffer | null)[] | Map; - -function isLeaf(bytes: BlobBytes): bytes is (Buffer | null)[] { - return Array.isArray(bytes); -} - -function nestedFieldType(type: Struct, bytes: Map): Struct { - return new Struct( - type.children.map((field) => { - const inner = bytes.get(field.name); - if (inner === undefined) { - return field; +): BlobSource[] { + if (projection === undefined) { + return blobColumns.map((source) => { + const [output, ...path] = source.split("."); + return { output, path, source }; + }); + } + const sources: BlobSource[] = []; + for (const [output, source] of projection) { + for (const column of blobColumns) { + if (column === source || column.startsWith(`${source}.`)) { + sources.push({ + output, + path: + column === source ? [] : column.slice(source.length + 1).split("."), + source: column, + }); } - if (isLeaf(inner)) { - return new Field(field.name, new LargeBinary(), true); - } - return new Field( - field.name, - nestedFieldType(field.type as Struct, inner), - field.nullable, - ); - }), - ); -} - -// Rebuilds a struct with the blobs under it as bytes. The struct keeps its own -// validity, so a null struct stays null. -function nestedData( - data: ArrowData, - bytes: Map, - offset: number, -): ArrowData { - const type = data.type as Struct; - const children = data.children.map((child, index) => { - const inner = bytes.get(type.children[index].name); - if (inner === undefined) { - return child; } - if (isLeaf(inner)) { - return vectorFromArray( - inner.slice(offset, offset + data.length), - new LargeBinary(), - ).data[0]; - } - return nestedData(child, inner, offset); - }); - return makeData({ - type: nestedFieldType(type, bytes), - length: data.length, - nullCount: data.nullCount, - nullBitmap: data.nullBitmap, - children, - }); + } + return sources; } async function readBlobBytes( table: ArrowTable, - sources: Map, - blobColumns: string[], + sources: BlobSource[], snapshot: NativeTable, dropRowId: boolean, ): Promise { @@ -157,42 +114,27 @@ async function readBlobBytes( rowIdColumn === null ? [] : Array.from(rowIdColumn.toArray() as BigUint64Array); - const fetch = async (column: string) => - (rowIds.length === 0 ? [] : await snapshot.fetchBlobs(column, rowIds)).map( - (value) => value ?? null, + const outputFields = new Map( + table.schema.fields.map((field) => [field.name, field]), + ); + const replacements = new Map(); + for (const { output, path, source } of sources) { + let field = outputFields.get(output); + for (const name of path) { + field = + field !== undefined && DataType.isStruct(field.type) + ? field.type.children.find((child) => child.name === name) + : undefined; + } + if (field === undefined) { + continue; + } + const values = + rowIds.length === 0 ? [] : await snapshot.fetchBlobs(source, rowIds); + replacements.set( + field, + values.map((value) => value ?? null), ); - const replacements = new Map(); - for (const [output, source] of sources) { - const index = table.schema.fields.findIndex( - (field) => field.name === output, - ); - if (index < 0) { - continue; - } - if (blobColumns.includes(source)) { - replacements.set(index, await fetch(source)); - continue; - } - if (!DataType.isStruct(table.schema.fields[index].type)) { - continue; - } - const nested = new Map(); - for (const column of blobColumns.filter((c) => - c.startsWith(`${source}.`), - )) { - const path = column.slice(source.length + 1).split("."); - let level = nested; - for (const name of path.slice(0, -1)) { - let next = level.get(name); - if (next === undefined) { - next = new Map(); - level.set(name, next); - } - level = next as Map; - } - level.set(path[path.length - 1], await fetch(column)); - } - replacements.set(index, nested); } const result = withBinaryColumns(table, replacements); if (!dropRowId) { @@ -205,48 +147,76 @@ async function readBlobBytes( ); } -// A fresh field, not a clone: the descriptor field's metadata must not follow -// the bytes. function withBinaryColumns( table: ArrowTable, - replacements: Map, + replacements: Map, ): ArrowTable { if (replacements.size === 0) { return table; } - const fields = table.schema.fields.map((field, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return field; + const replaceField = (field: Field): Field => { + if (replacements.has(field)) { + // The descriptor's extension and encoding metadata must not follow bytes. + return new Field(field.name, new LargeBinary(), true); } - return isLeaf(bytes) - ? new Field(field.name, new LargeBinary(), true) - : new Field( - field.name, - nestedFieldType(field.type as Struct, bytes), - field.nullable, - ); - }); + if (DataType.isStruct(field.type)) { + const children = field.type.children.map(replaceField); + if (children.some((child, i) => child !== field.type.children[i])) { + return field.clone({ type: new Struct(children) }); + } + } + return field; + }; + const fields = table.schema.fields.map(replaceField); const schema = new Schema(fields, table.schema.metadata); let offset = 0; const batches = table.batches.map((batch) => { - const children = batch.data.children.map((child, index) => { - const bytes = replacements.get(index); - if (bytes === undefined) { - return child; + const replaceData = ( + data: ArrowData, + field: Field, + updated: Field, + ): ArrowData => { + if (field === updated) { + return data; } - if (!isLeaf(bytes)) { - return nestedData(child, bytes, offset); + const values = replacements.get(field); + if (values !== undefined) { + return vectorFromArray( + values.slice(offset, offset + batch.numRows), + new LargeBinary(), + ).data[0]; } - return vectorFromArray( - bytes.slice(offset, offset + batch.numRows), - new LargeBinary(), - ).data[0]; - }); + const children = data.children.map((child, index) => + replaceData( + child, + field.type.children[index], + updated.type.children[index], + ), + ); + // Keep the struct's validity bitmap and offset, including null parents. + return data.clone( + updated.type, + undefined, + undefined, + undefined, + undefined, + children, + ); + }; + const children = batch.data.children.map((child, index) => + replaceData(child, table.schema.fields[index], fields[index]), + ); offset += batch.numRows; return new RecordBatch( schema, - makeData({ type: new Struct(fields), length: batch.numRows, children }), + batch.data.clone( + new Struct(fields), + undefined, + undefined, + undefined, + undefined, + children, + ), ); }); return new ArrowTable(schema, batches); @@ -274,13 +244,14 @@ export interface QueryExecutionOptions { * * - `"descriptions"` (default): each blob is a descriptor. Read the bytes * with {@link Table.fetchBlobs} or {@link Table.fetchBlobFiles}. - * - `"bytes"`: each blob column holds the blob bytes as `LargeBinary`, - * including blobs nested in structs. The query also reads `_rowid` to - * fetch the bytes, and leaves it out of the result unless - * {@link QueryBase.withRowId} was called. All the bytes are held in - * memory. + * - `"bytes"`: each blob column, including blobs nested in structs, holds + * the bytes as `LargeBinary`. The query also reads `_rowid` to fetch them, and + * leaves it out of the result unless {@link QueryBase.withRowId} was + * called. All the bytes are held in memory. * - * Blobs nested in a list keep their descriptors. + * Selecting a parent struct or directly aliasing a blob or struct also returns + * bytes. Blobs inside lists keep their descriptors. Arbitrary projection + * expressions are not materialized as bytes. * * The rows and the bytes come from the same table version. If the table * changes while the query runs, the query runs again, up to three times. @@ -529,15 +500,17 @@ export class QueryBase< // unchanged version afterwards proves it read the snapshot's version. const snapshot = await context.table.querySnapshot(); const version = await snapshot.version(); - const blobColumns = await snapshot.blobColumns(); - const sources = blobSources(blobColumns, context.projection); - if (sources.size === 0) { + const sources = blobSources( + await snapshot.blobColumns(), + context.projection, + ); + if (sources.length === 0) { return this.collect(options, false); } const addRowId = !context.withRowId; const table = await this.collect(options, addRowId); if ((await context.table.version()) === version) { - return readBlobBytes(table, sources, blobColumns, snapshot, addRowId); + return readBlobBytes(table, sources, snapshot, addRowId); } if (attempt === BLOB_BYTES_ATTEMPTS) { throw new Error( From 88203d25156a0dd7fc7d2c51aee66410d546ae18 Mon Sep 17 00:00:00 2001 From: Christopher Krah Date: Thu, 22 Oct 2026 21:42:36 +0200 Subject: [PATCH 60/68] fix(python): keep a renamed Permutation column in its position (#4626) The `Permutation.rename_column` docstring says it "only changes the name of the column that is visible from this permutation". But `permutation.py:868-870` adds the new name at the end of the selection and then deletes the old key, so the renamed column also moves to the last position. Hugging Face `datasets`, which `Permutation` mirrors (#2725), keeps it in place. Positional formats feel it most: `with_format("numpy")` swaps the feature order of every row. Before, on `d97df0c`: ``` $ python repro.py /python/python # lancedb 0.41.0b12 extension, Python sources from the tree ['id', 'value'] ['value', 'row_id'] [{'value': 0, 'row_id': 0}] [[10 1]] ``` After: ``` $ python repro.py /python/python ['id', 'value'] ['row_id', 'value'] [{'row_id': 0, 'value': 0}] [[ 1 10]] ```
repro.py (its first output line, lancedb.__file__, is cut above) ```python import sys; sys.path.insert(0, sys.argv[1]) import tempfile import lancedb from lancedb.permutation import Permutation print(lancedb.__file__) db = lancedb.connect(tempfile.mkdtemp()) tbl = db.create_table("t", [{"id": i, "value": i * 10} for i in range(3)]) perm = Permutation.identity(tbl) print(perm.column_names) renamed = perm.rename_column("id", "row_id") print(renamed.column_names) print(renamed.__getitems__([0])) print(renamed.with_format("numpy").__getitems__([1])) ```
For comparison, `datasets` 5.1.0 gives `['new_id', 'value']` for `Dataset.from_dict({'id': [1, 2], 'value': [3, 4]}).rename_column('id', 'new_id').column_names`. No issue exists for this. The fix rebuilds the selection dict in order. The cost: code that renamed a column and then read it by position in numpy batches (say `[:, -1]`) gets a different column after this, with no error. `rename_columns` calls `rename_column` once per name, so it gets the fix too. The tests that check schema order only rename `value`, the last column, and `test_rename_rowid` checks membership but not order. The new test renames `id`. ### Verification ``` # Python 3.12.6; lancedb 0.41.0b12 wheel from pypi.fury.io, its _lancedb.abi3.so copied into the tree; pylance 14.0.0-rc.1 $ PYTHONPATH=python python -m pytest python/tests/test_permutation.py -q -p no:cacheprovider d97df0c: 72 passed in 2.53s this branch: 73 passed in 2.36s permutation.py reverted, test kept: FAILED python/tests/test_permutation.py::test_rename_column_keeps_position assert value: int64\nnew_id: int64 == new_id: int64\nvalue: int64 1 failed, 72 passed in 2.41s $ ruff format --check . && ruff check . # ruff 0.9.9, as python.yml pins it 152 files already formatted All checks passed! # not run: pyright (its include list covers neither permutation.py nor tests/), the Rust build ``` @westonpace, this is your class from #2725. Do you want the docstring to say outright that the column keeps its position? If so, I'll add that line to both rename methods in this PR. Signed-off-by: Christopher Krah <48338417+chrikrah@users.noreply.github.com> --- python/python/lancedb/permutation.py | 8 +++++--- python/python/tests/test_permutation.py | 7 +++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index 355ebf798..e95db27f6 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -909,9 +909,11 @@ class Permutation: f"Cannot rename column `{old_name}` to `{new_name}` because a column " "with that name already exists" ) - new_selection = self.selection.copy() - new_selection[new_name] = new_selection[old_name] - del new_selection[old_name] + # Rebuild the mapping so the renamed column keeps its position + new_selection = { + (new_name if name == old_name else name): value + for name, value in self.selection.items() + } return self._with_selection(new_selection) def rename_columns(self, column_map: dict[str, str]) -> "Permutation": diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index ed1894340..50ca53061 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -811,6 +811,13 @@ def test_rename_column(some_permutation: Permutation): some_permutation.rename_column("non_existent", "new_value") +def test_rename_column_keeps_position(some_permutation: Permutation): + renamed = some_permutation.rename_column("id", "new_id") + assert renamed.schema == pa.schema([("new_id", pa.int64()), ("value", pa.int64())]) + batch = renamed.with_format("arrow").__getitems__([0]) + assert batch.schema.names == ["new_id", "value"] + + def test_rename_columns(some_permutation: Permutation): assert some_permutation.rename_columns({"value": "new_value"}).schema == pa.schema( [("id", pa.int64()), ("new_value", pa.int64())] From 631ea41a635890218bf7ee1ab80df6caccc7606c Mon Sep 17 00:00:00 2001 From: Colin Patrick McCabe Date: Thu, 22 Oct 2026 15:41:30 -0700 Subject: [PATCH 61/68] feat!: return lazy iterators for table and namespace listings (#4599) Make table and namespace listings consistent with other resource listings across Rust, Python, and TypeScript. Callers can iterate over individual names across all pages, with configurable page sizes and continuation tokens. In order to avoid breaking existing Python users, this change adds some compatibility methods to the returned Python iterators. Code such as: ``` for t in db.list_tables().tables: print(t) ``` will continue to work as expected because the iterator has a `tables` field and a `page_token` field, which function as expected to reveal the first page of results. In other words, it is never necessary to use the iterator __as an iterator__. It may be used as if it were an old-style response object. (As the iterator progresses, the `tables` list may switch to the second, third, etc. page of results.) Similar compatibility code allows the async python API to keep working as well. --- docs/src/js/classes/Connection.md | 80 ++------ docs/src/js/globals.md | 4 - .../js/interfaces/ListNamespacesOptions.md | 27 --- .../js/interfaces/ListNamespacesResponse.md | 23 --- docs/src/js/interfaces/ListTablesOptions.md | 35 ---- docs/src/js/interfaces/ListTablesResponse.md | 23 --- docs/src/js/interfaces/TableNamesOptions.md | 2 +- docs/src/python/python.md | 23 ++- docs/web-tests/ts/connection.test.ts | 16 +- docs/web-tests/ts/quickstart.test.ts | 7 +- nodejs/__test__/connection.test.ts | 156 ++++++++-------- nodejs/lancedb/connection.ts | 133 ++++---------- nodejs/lancedb/index.ts | 4 - nodejs/src/connection.rs | 101 ++--------- python/python/lancedb/_lancedb.pyi | 11 +- python/python/lancedb/db.py | 70 ++++--- python/python/lancedb/listing.py | 171 ++++++++++++++++++ python/python/lancedb/namespace.py | 39 ++-- python/python/lancedb/remote/db.py | 16 +- python/python/tests/test_embeddings.py | 4 + python/python/tests/test_listing.py | 137 ++++++++++++++ python/python/tests/test_ollama_embeddings.py | 2 +- python/src/connection.rs | 56 ++---- python/src/listing.rs | 30 ++- rust/lancedb/src/connection.rs | 161 +++++++++-------- rust/lancedb/src/listing.rs | 130 +++++++++---- rust/lancedb/src/remote/db.rs | 66 ++----- 27 files changed, 814 insertions(+), 713 deletions(-) delete mode 100644 docs/src/js/interfaces/ListNamespacesOptions.md delete mode 100644 docs/src/js/interfaces/ListNamespacesResponse.md delete mode 100644 docs/src/js/interfaces/ListTablesOptions.md delete mode 100644 docs/src/js/interfaces/ListTablesResponse.md diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index d5ab8c3d1..c110eb38c 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -653,94 +653,48 @@ See [Listing](Listing.md) for pagination state and error behavior. ### listNamespaces() ```ts -abstract listNamespaces(namespacePath?, options?): Promise +abstract listNamespaces(namespacePath?, options?): Listing ``` -List the immediate child namespaces under the given parent. - -Results may be paginated. To retrieve subsequent pages, pass the -`pageToken` returned by a previous call. +Lazily list immediate child namespaces, 100 per page by default. #### Parameters * **namespacePath?**: `string`[] - The parent namespace path. Defaults - to the root namespace if omitted. -* **options?**: `Partial`<[`ListNamespacesOptions`](../interfaces/ListNamespacesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) #### Returns -`Promise`<[`ListNamespacesResponse`](../interfaces/ListNamespacesResponse.md)> - -Child namespace names and - an optional token for fetching the next page. +[`Listing`](Listing.md)<`string`> *** ### listTables() -#### listTables(options) - ```ts -abstract listTables(options?): Promise +abstract listTables(namespacePath?, options?): Listing ``` -List a page of the tables in this database. +Lazily list tables in a namespace, requesting 100 per page by default. -To retrieve the tables after the page, pass the `pageToken` the response -carries back in. A page can be shorter than `limit` without being the last -one, so walk until a response carries no page token: - -```ts -const names = []; -let pageToken = undefined; -do { - const page = await conn.listTables({ pageToken, limit: 100 }); - names.push(...page.tables); - pageToken = page.pageToken; -} while (pageToken); -``` - -##### Parameters - -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). - -##### Returns - -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> - -A page of table names and an - optional token for the tables after it. - -#### listTables(namespacePath, options) - -```ts -abstract listTables(namespacePath?, options?): Promise -``` - -List a page of the tables in this database. - -##### Parameters +#### Parameters * **namespacePath?**: `string`[] - The namespace path to list tables from - (defaults to root namespace) -* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)> - Pagination options - (`pageToken`, `limit`). +* **options?**: [`ListingOptions`](../interfaces/ListingOptions.md) -##### Returns +#### Returns -`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)> +[`Listing`](Listing.md)<`string`> -A page of table names and an - optional token for the tables after it. +#### Example + +```ts +for await (const name of conn.listTables([], { pageLimit: 20 })) { + console.log(name); +} +``` *** diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 646341190..e3830e9ad 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -108,10 +108,6 @@ - [JobEventsOptions](interfaces/JobEventsOptions.md) - [JobFailureInfo](interfaces/JobFailureInfo.md) - [JobInfo](interfaces/JobInfo.md) -- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) -- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) -- [ListTablesOptions](interfaces/ListTablesOptions.md) -- [ListTablesResponse](interfaces/ListTablesResponse.md) - [ListingOptions](interfaces/ListingOptions.md) - [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md) - [MemtableStats](interfaces/MemtableStats.md) diff --git a/docs/src/js/interfaces/ListNamespacesOptions.md b/docs/src/js/interfaces/ListNamespacesOptions.md deleted file mode 100644 index 79680417f..000000000 --- a/docs/src/js/interfaces/ListNamespacesOptions.md +++ /dev/null @@ -1,27 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesOptions - -# Interface: ListNamespacesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An optional limit to the number of results to return. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response for pagination. diff --git a/docs/src/js/interfaces/ListNamespacesResponse.md b/docs/src/js/interfaces/ListNamespacesResponse.md deleted file mode 100644 index a99ab480f..000000000 --- a/docs/src/js/interfaces/ListNamespacesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListNamespacesResponse - -# Interface: ListNamespacesResponse - -## Properties - -### namespaces - -```ts -namespaces: string[]; -``` - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` diff --git a/docs/src/js/interfaces/ListTablesOptions.md b/docs/src/js/interfaces/ListTablesOptions.md deleted file mode 100644 index b8dd8b5ef..000000000 --- a/docs/src/js/interfaces/ListTablesOptions.md +++ /dev/null @@ -1,35 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesOptions - -# Interface: ListTablesOptions - -## Properties - -### limit? - -```ts -optional limit: number; -``` - -An upper bound on how many tables to return. -Defaults to 100. Zero returns an empty page without a continuation token. - -A page may hold fewer than this and still not be the last one, so keep -going while the response carries a page token rather than while pages are -full. - -*** - -### pageToken? - -```ts -optional pageToken: string; -``` - -Token from a previous response, to resume listing where it left off. - -The token is opaque: it carries whatever the database needs to resume, and -callers should not construct or interpret one. diff --git a/docs/src/js/interfaces/ListTablesResponse.md b/docs/src/js/interfaces/ListTablesResponse.md deleted file mode 100644 index 76cac2b23..000000000 --- a/docs/src/js/interfaces/ListTablesResponse.md +++ /dev/null @@ -1,23 +0,0 @@ -[**@lancedb/lancedb**](../README.md) • **Docs** - -*** - -[@lancedb/lancedb](../globals.md) / ListTablesResponse - -# Interface: ListTablesResponse - -## Properties - -### pageToken? - -```ts -optional pageToken: string; -``` - -*** - -### tables - -```ts -tables: string[]; -``` diff --git a/docs/src/js/interfaces/TableNamesOptions.md b/docs/src/js/interfaces/TableNamesOptions.md index 123aa6c9f..e1fd8c382 100644 --- a/docs/src/js/interfaces/TableNamesOptions.md +++ b/docs/src/js/interfaces/TableNamesOptions.md @@ -8,7 +8,7 @@ ## Deprecated -Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) +Use [ListingOptions](ListingOptions.md) with [Connection.listTables](../classes/Connection.md#listtables) instead. ## Properties diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f4c697497..c465d0ac4 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -26,14 +26,17 @@ is also an [asynchronous API client](#connections-asynchronous). ::: lancedb.db.DBConnection -Resource listings (`list_secrets`, `list_views`, `list_jobs`, `list_functions`, and -`list_materialized_views`) return lazy iterators. Use `list(db.list_views())` to -collect synchronous results, or `[name async for name in db.list_views()]` with -an asynchronous connection. `page_limit` controls each request, and `page_token` -resumes a listing. Drain the iterator's cached results before saving its token. +Resource listings return iterators. Use `list(db.list_tables())` to collect all +table names. `db.list_tables().tables` is a list of the first page; `limit` +controls its page size and `page_token` resumes at a saved token. The same +patterns apply to `list_namespaces()`. Other resource listings use `page_limit`. ::: lancedb.listing.Listing +::: lancedb.listing.TableListing + +::: lancedb.listing.NamespaceListing + ::: lancedb.Session ## Catalogs (Synchronous) @@ -631,6 +634,16 @@ can be used to create, list, or open tables. ::: lancedb.listing.AsyncListing +`async for name in db.list_tables()` loads pages as needed. To inspect the +first page, use `page = await db.list_tables()` and then `page.tables`. Reading +the page list does not consume the iterator. `page.page_token` is the token for +the next request; drain cached results before using it to resume in a new +listing. The same behavior applies to namespaces. + +::: lancedb.listing.AsyncTableListing + +::: lancedb.listing.AsyncNamespaceListing + ## Namespaces (Asynchronous) ::: lancedb.connect_namespace_async diff --git a/docs/web-tests/ts/connection.test.ts b/docs/web-tests/ts/connection.test.ts index 891501d4c..651be7d36 100644 --- a/docs/web-tests/ts/connection.test.ts +++ b/docs/web-tests/ts/connection.test.ts @@ -4,6 +4,12 @@ import { afterEach, describe, expect, jest, test } from "@jest/globals"; import * as lancedb from "@lancedb/lancedb"; import { withTempDirectory } from "./util.ts"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + test("connect to a local database", async () => { await withTempDirectory(async (tempDir) => { const cwd = process.cwd(); @@ -13,12 +19,12 @@ test("connect to a local database", async () => { const db = await lancedb.connect("ex_lancedb"); // --8<-- [end:connect] await db.createTable("probe", [{ id: 1 }]); - expect((await db.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(db.listTables())).toEqual(["probe"]); } finally { process.chdir(cwd); } const reopened = await lancedb.connect(`${tempDir}/ex_lancedb`); - expect((await reopened.listTables()).tables).toEqual(["probe"]); + expect(await collectNames(reopened.listTables())).toEqual(["probe"]); }); }); @@ -109,8 +115,8 @@ async function namespaceTableOpsExample() { ); // Verify - console.log((await db.listNamespaces()).namespaces); // ["prod"] - console.log((await db.listNamespaces(["prod"])).namespaces); // ["recommendations", "search"] + console.log(await collectNames(db.listNamespaces())); // ["prod"] + console.log(await collectNames(db.listNamespaces(["prod"]))); // ["recommendations", "search"] console.log(await db.tableNames(["prod", "search"])); // ["user"] console.log(await db.tableNames(["prod", "recommendations"])); // ["user"] // --8<-- [end:namespace_table_ops] @@ -124,7 +130,7 @@ async function namespaceAdminOpsExample() { await db.createNamespace(["prod"]); await db.createNamespace(["prod", "search"]); - const childNamespaces = (await db.listNamespaces(["prod"])).namespaces; + const childNamespaces = await collectNames(db.listNamespaces(["prod"])); console.log(`Child namespaces under ${JSON.stringify(namespace)}:`, childNamespaces); // Child namespaces under ["prod","search"]: [ 'search' ] diff --git a/docs/web-tests/ts/quickstart.test.ts b/docs/web-tests/ts/quickstart.test.ts index e32a14f86..2c3ee8252 100644 --- a/docs/web-tests/ts/quickstart.test.ts +++ b/docs/web-tests/ts/quickstart.test.ts @@ -20,12 +20,7 @@ const connections = async function tableNames(db: lancedb.Connection) { const names: string[] = []; - let pageToken: string | undefined; - do { - const page = await db.listTables({ pageToken }); - names.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); + for await (const name of db.listTables()) names.push(name); return names; } diff --git a/nodejs/__test__/connection.test.ts b/nodejs/__test__/connection.test.ts index bbd0c3c52..73f99803a 100644 --- a/nodejs/__test__/connection.test.ts +++ b/nodejs/__test__/connection.test.ts @@ -5,15 +5,15 @@ import { mkdirSync, readdirSync } from "fs"; import { join } from "path"; import { Field, Float64, Schema } from "apache-arrow"; import * as tmp from "tmp"; -import { - Connection, - ListTablesResponse, - Table, - connect, - connectNamespace, -} from "../lancedb"; +import { Connection, Table, connect, connectNamespace } from "../lancedb"; import { LocalTable } from "../lancedb/table"; +async function collectNames(items: AsyncIterable): Promise { + const names: string[] = []; + for await (const name of items) names.push(name); + return names; +} + describe("when connecting", () => { let tmpDir: tmp.DirResult; beforeEach(() => { @@ -135,7 +135,7 @@ describe("given a connection", () => { await db.close(); expect(db.isOpen()).toBe(false); await expect(db.tableNames()).rejects.toThrow("Connection is closed"); - await expect(db.listTables()).rejects.toThrow("Connection is closed"); + expect(() => db.listTables()).toThrow("Connection is closed"); await expect(db.renameTable("a", "b")).rejects.toThrow( "Connection is closed", ); @@ -247,15 +247,15 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const first = await db.listTables(namespacePath, { limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables(namespacePath, { pageLimit: 1 }); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables(namespacePath, { - limit: 1, - pageToken: first.pageToken, + const second = db.listTables(namespacePath, { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); }); it("should keep the existing table when mode is exist_ok", async () => { @@ -299,19 +299,18 @@ describe("given a connection", () => { await db.createTable("a", [{ id: 1 }]); await db.createTable("c", [{ id: 1 }]); - const all = await db.listTables(); - expect(all.tables).toEqual(["a", "b", "c"]); - expect(all.pageToken).toBeUndefined(); + expect(await collectNames(db.listTables())).toEqual(["a", "b", "c"]); - const first = await db.listTables({ limit: 1 }); - expect(first.tables).toEqual(["a"]); - expect(first.pageToken).toBeDefined(); + const first = db.listTables([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect(await first.next()).toEqual({ done: false, value: "a" }); + expect(first.pageToken()).toBeDefined(); - const second = await db.listTables({ - limit: 1, - pageToken: first.pageToken, + const second = db.listTables([], { + pageLimit: 1, + pageToken: first.pageToken(), }); - expect(second.tables).toEqual(["b"]); + expect(await second.next()).toEqual({ done: false, value: "b" }); }); it("should visit every table exactly once when walking pages", async () => { @@ -322,18 +321,9 @@ describe("given a connection", () => { await db.createTable(name, [{ id: 1 }]); } - const seen: string[] = []; - let pageToken: string | undefined = undefined; - do { - const page: ListTablesResponse = await db.listTables({ - limit: 2, - pageToken, - }); - seen.push(...page.tables); - pageToken = page.pageToken; - } while (pageToken); - - expect(seen).toEqual(created); + expect(await collectNames(db.listTables([], { pageLimit: 2 }))).toEqual( + created, + ); }); it.each([2 ** 32, 2 ** 32 + 1, Number.MAX_SAFE_INTEGER])( @@ -347,11 +337,10 @@ describe("given a connection", () => { ]) { await expect(names).resolves.toEqual(["a", "b"]); } - for (const page of [ - db.listTables({ limit }), - db.listTables([], { limit }), - ]) { - await expect(page).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + expect(() => + db.listTables(namespacePath, { pageLimit: limit }), + ).toThrow(/limit/); } }, ); @@ -366,8 +355,11 @@ describe("given a connection", () => { ])("should reject an invalid table listing limit %s", async (limit) => { await expect(db.tableNames({ limit })).rejects.toThrow(/limit/); await expect(db.tableNames([], { limit })).rejects.toThrow(/limit/); - await expect(db.listTables({ limit })).rejects.toThrow(/limit/); - await expect(db.listTables([], { limit })).rejects.toThrow(/limit/); + for (const namespacePath of [undefined, []]) { + expect(() => db.listTables(namespacePath, { pageLimit: limit })).toThrow( + /limit/, + ); + } }); it("should preserve zero, small and unsigned-32-bit table listing limits", async () => { @@ -379,11 +371,21 @@ describe("given a connection", () => { [2 ** 32 - 1, ["a", "b"]], ] as const) { await expect(db.tableNames({ limit })).resolves.toEqual(names); - const page = await db.listTables({ limit }); - expect(page.tables).toEqual(names); } await expect(db.tableNames()).resolves.toEqual(["a", "b"]); - await expect(db.listTables()).resolves.toEqual({ tables: ["a", "b"] }); + for (const namespacePath of [undefined, []]) { + for (const pageLimit of [1, 2 ** 31 - 1]) { + expect( + await collectNames(db.listTables(namespacePath, { pageLimit })), + ).toEqual(["a", "b"]); + } + for (const pageLimit of [0, 2 ** 31, 2 ** 32 - 1]) { + expect(() => db.listTables(namespacePath, { pageLimit })).toThrow( + /limit/, + ); + } + } + expect(await collectNames(db.listTables())).toEqual(["a", "b"]); }); it("should list all names and use a default page size of 100", async () => { @@ -396,20 +398,17 @@ describe("given a connection", () => { } await expect(db.tableNames()).resolves.toEqual(names); await expect(db.tableNames({ limit: 0 })).resolves.toEqual([]); - const first = await db.listTables(); - expect(first.tables).toEqual(names.slice(0, 100)); - expect(first.pageToken).toBeDefined(); - const second = await db.listTables({ pageToken: first.pageToken }); - expect(second.tables).toEqual(names.slice(100)); - expect(second.pageToken).toBeUndefined(); - await expect(db.listTables({ limit: 200 })).resolves.toEqual({ - tables: names, - pageToken: undefined, - }); - await expect(db.listTables({ limit: 0 })).resolves.toEqual({ - tables: [], - pageToken: undefined, - }); + const first = db.listTables(); + const firstPage: string[] = []; + for (let i = 0; i < 100; i++) firstPage.push((await first.next()).value!); + expect(firstPage).toEqual(names.slice(0, 100)); + expect(first.pageToken()).toBeDefined(); + expect(first.numPageResults()).toBe(0); + expect(await collectNames(first)).toEqual(names.slice(100)); + expect(await collectNames(db.listTables([], { pageLimit: 200 }))).toEqual( + names, + ); + expect(() => db.listTables([], { pageLimit: 0 })).toThrow(/limit/); }); it("should list tables in a namespace", async () => { @@ -420,12 +419,8 @@ describe("given a connection", () => { await db.createNamespace(["child"]); await db.createTable("nested", [{ id: 1 }], ["child"]); - await expect(db.listTables(["child"])).resolves.toEqual( - expect.objectContaining({ tables: ["nested"] }), - ); - await expect(db.listTables()).resolves.toEqual( - expect.objectContaining({ tables: [] }), - ); + expect(await collectNames(db.listTables(["child"]))).toEqual(["nested"]); + expect(await collectNames(db.listTables())).toEqual([]); }); it("should create tables in v2 mode", async () => { @@ -719,22 +714,35 @@ describe("namespaces", () => { it("should list namespaces created at the root", async () => { await db.createNamespace(["alpha"]); await db.createNamespace(["beta"]); - const list = await db.listNamespaces(); - expect(list.namespaces).toEqual(expect.arrayContaining(["alpha", "beta"])); + const list = await collectNames(db.listNamespaces()); + expect(list).toEqual(expect.arrayContaining(["alpha", "beta"])); }); it("should list child namespaces under a parent", async () => { await db.createNamespace(["parent"]); await db.createNamespace(["parent", "child"]); - const list = await db.listNamespaces(["parent"]); - expect(list.namespaces).toContain("child"); + const list = await collectNames(db.listNamespaces(["parent"])); + expect(list).toContain("child"); + }); + + it("should resume a namespace listing from its page token", async () => { + await db.createNamespace(["alpha"]); + await db.createNamespace(["beta"]); + const first = db.listNamespaces([], { pageLimit: 1 }); + expect(first.numPageResults()).toBe(0); + expect((await first.next()).value).toBe("alpha"); + const second = db.listNamespaces([], { + pageLimit: 1, + pageToken: first.pageToken(), + }); + expect(await collectNames(second)).toEqual(["beta"]); }); it("should drop a namespace", async () => { await db.createNamespace(["ephemeral"]); await db.dropNamespace(["ephemeral"]); - const list = await db.listNamespaces(); - expect(list.namespaces).not.toContain("ephemeral"); + const list = await collectNames(db.listNamespaces()); + expect(list).not.toContain("ephemeral"); }); it("should raise an error on any namespace op after close", async () => { @@ -742,7 +750,7 @@ describe("namespaces", () => { await expect(db.describeNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); - await expect(db.listNamespaces()).rejects.toThrow("Connection is closed"); + expect(() => db.listNamespaces()).toThrow("Connection is closed"); await expect(db.createNamespace(["foo"])).rejects.toThrow( "Connection is closed", ); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index c93414fab..5f067eeef 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -30,16 +30,12 @@ import type { DescribeNamespaceResponse, DropNamespaceResponse, JobInfo, - ListNamespacesResponse, - ListTablesResponse, } from "./native"; import { ViewDescription, viewDescriptionFromNative } from "./view"; export type { CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, }; import { sanitizeTable } from "./sanitize"; import { LocalTable, Table } from "./table"; @@ -141,7 +137,7 @@ export interface OpenTableOptions { } /** - * @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables} + * @deprecated Use {@link ListingOptions} with {@link Connection.listTables} * instead. */ export interface TableNamesOptions { @@ -157,32 +153,6 @@ export interface TableNamesOptions { limit?: number; } -export interface ListTablesOptions { - /** - * Token from a previous response, to resume listing where it left off. - * - * The token is opaque: it carries whatever the database needs to resume, and - * callers should not construct or interpret one. - */ - pageToken?: string; - /** - * An upper bound on how many tables to return. - * Defaults to 100. Zero returns an empty page without a continuation token. - * - * A page may hold fewer than this and still not be the last one, so keep - * going while the response carries a page token rather than while pages are - * full. - */ - limit?: number; -} - -export interface ListNamespacesOptions { - /** Token from a previous response for pagination. */ - pageToken?: string; - /** An optional limit to the number of results to return. */ - limit?: number; -} - export interface CreateNamespaceOptions { /** Creation mode. */ mode?: "create" | "exist_ok" | "overwrite"; @@ -278,45 +248,18 @@ export abstract class Connection { options?: Partial, ): Promise; - /** - * List a page of the tables in this database. - * - * To retrieve the tables after the page, pass the `pageToken` the response - * carries back in. A page can be shorter than `limit` without being the last - * one, so walk until a response carries no page token: - * + /** Lazily list tables in a namespace, requesting 100 per page by default. + * @example * ```ts - * const names = []; - * let pageToken = undefined; - * do { - * const page = await conn.listTables({ pageToken, limit: 100 }); - * names.push(...page.tables); - * pageToken = page.pageToken; - * } while (pageToken); + * for await (const name of conn.listTables([], { pageLimit: 20 })) { + * console.log(name); + * } * ``` - * - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. - */ - abstract listTables( - options?: Partial, - ): Promise; - /** - * List a page of the tables in this database. - * - * @param {string[]} namespacePath - The namespace path to list tables from - * (defaults to root namespace) - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} A page of table names and an - * optional token for the tables after it. */ abstract listTables( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Open a table in the database. @@ -542,23 +485,11 @@ export abstract class Connection { namespacePath: string[], ): Promise; - /** - * List the immediate child namespaces under the given parent. - * - * Results may be paginated. To retrieve subsequent pages, pass the - * `pageToken` returned by a previous call. - * - * @param {string[]} namespacePath - The parent namespace path. Defaults - * to the root namespace if omitted. - * @param {Partial} options - Pagination options - * (`pageToken`, `limit`). - * @returns {Promise} Child namespace names and - * an optional token for fetching the next page. - */ + /** Lazily list immediate child namespaces, 100 per page by default. */ abstract listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise; + options?: ListingOptions, + ): Listing; /** * Create a new namespace at the given path. @@ -821,22 +752,17 @@ export class LocalConnection extends Connection { ); } - async listTables( - namespacePathOrOptions?: string[] | Partial, - options?: Partial, - ): Promise { - // Detect if first argument is namespacePath array or options object - const namespacePath = Array.isArray(namespacePathOrOptions) - ? namespacePathOrOptions - : undefined; - const listTablesOptions = Array.isArray(namespacePathOrOptions) - ? options - : (namespacePathOrOptions ?? options); - - return this.inner.listTables( - namespacePath ?? [], - listTablesOptions?.pageToken, - normalizeTableListingLimit(listTablesOptions?.limit), + listTables( + namespacePath?: string[], + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listTables( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } @@ -1039,12 +965,15 @@ export class LocalConnection extends Connection { listNamespaces( namespacePath?: string[], - options?: Partial, - ): Promise { - return this.inner.listNamespaces( - namespacePath ?? [], - options?.pageToken, - options?.limit, + options: ListingOptions = {}, + ): Listing { + validateListingOptions(options); + return new Listing( + this.inner.listNamespaces( + namespacePath ?? [], + options.pageToken, + options.pageLimit, + ), ); } diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 4ad8b2a93..0f1b5f333 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -97,13 +97,9 @@ export { Connection, CreateTableOptions, TableNamesOptions, - ListTablesOptions, OpenTableOptions, - ListNamespacesOptions, CreateNamespaceOptions, DropNamespaceOptions, - ListNamespacesResponse, - ListTablesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 67c4706e4..914f9e4c9 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -16,8 +16,7 @@ use crate::table::Table; use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, connect_namespace}; use lance_namespace::models::{ - CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest, - ListTablesRequest, + CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, }; use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema}; @@ -31,18 +30,6 @@ pub struct DescribeNamespaceResponse { pub properties: Option>, } -#[napi(object)] -pub struct ListNamespacesResponse { - pub namespaces: Vec, - pub page_token: Option, -} - -#[napi(object)] -pub struct ListTablesResponse { - pub tables: Vec, - pub page_token: Option, -} - #[napi(object)] pub struct CreateNamespaceResponse { pub properties: Option>, @@ -239,31 +226,20 @@ impl Connection { op.execute().await.default_error() } - /// List a page of tables in the database. + /// Lazily list tables in the database. #[napi(catch_unwind)] - pub async fn list_tables( + pub fn list_tables( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let request = ListTablesRequest { - // The root namespace is an empty path, not an absent one: a namespace-backed - // database rejects a request that names no namespace. - id: Some(namespace_path.unwrap_or_default()), - page_token, - limit: limit.map(|limit| i32::try_from(limit).unwrap_or(i32::MAX)), - ..Default::default() - }; - let response = self - .get_inner()? - .list_tables(request) - .await - .default_error()?; - Ok(ListTablesResponse { - tables: response.tables, - page_token: response.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Create table from a Apache Arrow IPC (file) buffer. @@ -663,28 +639,19 @@ impl Connection { } #[napi(catch_unwind)] - /// List child namespaces under the given namespace path - pub async fn list_namespaces( + /// Lazily list child namespaces under the given namespace path + pub fn list_namespaces( &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> napi::Result { - let req = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: checked_limit(limit)?, - ..Default::default() - }; - let resp = self - .get_inner()? - .list_namespaces(req) - .await - .default_error()?; - Ok(ListNamespacesResponse { - namespaces: resp.namespaces, - page_token: resp.page_token, - }) + ) -> napi::Result { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[napi(catch_unwind)] @@ -787,33 +754,3 @@ impl Connection { .default_error() } } - -fn checked_limit(limit: Option) -> napi::Result> { - limit - .map(i32::try_from) - .transpose() - .map_err(|_| napi::Error::from_reason("limit must not exceed 2147483647")) -} - -#[cfg(test)] -mod tests { - use super::checked_limit; - - #[test] - fn checked_limit_rejects_values_above_i32_max() { - // u32 values above i32::MAX wrapped negative through `as i32`, - // sending a negative page size to the server. - assert!(checked_limit(Some(3_000_000_000)).is_err()); - assert!(checked_limit(Some(u32::MAX)).is_err()); - } - - #[test] - fn checked_limit_passes_boundary_and_missing_values() { - assert_eq!( - checked_limit(Some(i32::MAX as u32)).unwrap(), - Some(i32::MAX) - ); - assert_eq!(checked_limit(Some(0)).unwrap(), Some(0)); - assert_eq!(checked_limit(None).unwrap(), None); - } -} diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 96eef064c..6ef23abc7 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -19,11 +19,9 @@ from .index import ( FTS, ) from lance_namespace import ( - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, ) from .remote import ClientConfig from .types import BaseTokenizerType @@ -114,12 +112,12 @@ class Connection(object): uri: str async def is_open(self): ... async def close(self): ... - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: ... + ) -> NameListing: ... async def create_namespace( self, namespace_path: List[str], @@ -136,12 +134,12 @@ class Connection(object): self, namespace_path: List[str], ) -> DescribeNamespaceResponse: ... - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: ... + ) -> NameListing: ... async def table_names( self, namespace_path: Optional[List[str]], @@ -1010,6 +1008,7 @@ async def connect_catalog( class NameListing: def num_page_results(self) -> int: ... def page_token(self) -> Optional[str]: ... + async def next_page(self) -> Optional[dict[str, Any]]: ... def __aiter__(self) -> NameListing: ... async def __anext__(self) -> str: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index a74763afb..90dc1cccf 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -32,18 +32,24 @@ if sys.version_info >= (3, 12): else: from overrides import EnforceOverrides, override # type: ignore -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, + validate_page_limit, +) from lancedb.embeddings.registry import EmbeddingFunctionRegistry from lancedb.common import data_to_reader, sanitize_uri, validate_schema from lancedb.background_loop import LOOP from lance_namespace import ( LanceNamespace, - ListNamespacesResponse, CreateNamespaceResponse, DropNamespaceResponse, DescribeNamespaceResponse, - ListTablesResponse, connect as namespace_connect, ) from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError @@ -171,7 +177,7 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -187,12 +193,12 @@ class DBConnection(EnforceOverrides): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return ListNamespacesResponse(namespaces=[], page_token=None) + return NamespaceListing(AsyncNamespaceListing(None, empty=True)) def create_namespace( self, @@ -317,8 +323,8 @@ class DBConnection(EnforceOverrides): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -334,7 +340,7 @@ class DBConnection(EnforceOverrides): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ raise NotImplementedError( @@ -1436,7 +1442,7 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -1452,16 +1458,14 @@ class LanceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( - namespace_path=namespace_path, - page_token=page_token, - limit=limit, + namespace_path=namespace_path, page_token=page_token, limit=limit ) ) @@ -1512,8 +1516,8 @@ class LanceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> TableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -1529,12 +1533,12 @@ class LanceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -2286,12 +2290,12 @@ class AsyncConnection(object): else: return None - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -2306,15 +2310,18 @@ class AsyncConnection(object): Returns ------- - ListNamespacesResponse + AsyncNamespaceListing Response containing namespace names and optional pagination token """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_namespaces( + validate_page_limit(limit) + inner = self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListNamespacesResponse(**result) + if limit == 0: + return AsyncNamespaceListing(None, empty=True) + return AsyncNamespaceListing(inner, page_token) async def create_namespace( self, @@ -2393,13 +2400,13 @@ class AsyncConnection(object): result = await self._inner.describe_namespace(namespace_path) return DescribeNamespaceResponse(**result) - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: - """List a page of tables in this database. + ) -> AsyncTableListing: + """List tables in this database, one page at a time. Parameters ---------- @@ -2415,15 +2422,18 @@ class AsyncConnection(object): Returns ------- - ListTablesResponse + AsyncTableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - result = await self._inner.list_tables( + validate_page_limit(limit) + inner = self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) - return ListTablesResponse(**result) + if limit == 0: + return AsyncTableListing(None, empty=True) + return AsyncTableListing(inner, page_token) async def table_names( self, diff --git a/python/python/lancedb/listing.py b/python/python/lancedb/listing.py index 0620aa7d8..8abc9d85b 100644 --- a/python/python/lancedb/listing.py +++ b/python/python/lancedb/listing.py @@ -5,6 +5,7 @@ from __future__ import annotations +import asyncio from typing import Any, AsyncIterator, Callable, Generic, Iterator, Optional, TypeVar from .background_loop import LOOP @@ -83,3 +84,173 @@ class Listing(Iterator[T], Generic[T]): for initial, cached, final-page, and error behavior. """ return self._inner.page_token() + + +def validate_page_limit(limit: Optional[int]) -> None: + """Validate the legacy limit argument before passing it to the bindings.""" + if limit is not None and ( + not isinstance(limit, int) or limit < 0 or limit > 2**31 - 1 + ): + raise ValueError("limit must be between 0 and 2147483647") + + +class _AsyncPagedNameListing(AsyncIterator[str]): + """An async listing with an inspectable copy of its current backend page.""" + + def __init__( + self, inner: Any, page_token: Optional[str] = None, *, empty: bool = False + ): + self._inner = inner + self._token = page_token + self._page: Optional[list[str]] = [] if empty else None + self._position = 0 + self._done = empty + self._failure: Optional[BaseException] = None + self._lock = asyncio.Lock() + + @property + def page_token(self) -> Optional[str]: + """The continuation token for the next backend request.""" + return self._token + + def num_page_results(self) -> int: + """Count the remaining names in the current page.""" + return 0 if self._page is None else len(self._page) - self._position + + def _page_copy(self) -> list[str]: + if self._page is None: + raise RuntimeError("Await or iterate the listing before reading its page") + return self._page.copy() + + async def _fetch_page(self) -> None: + try: + page = await self._inner.next_page() + except asyncio.CancelledError: + raise + except BaseException as exc: + self._failure = exc + self._done = True + raise + if page is None: + self._page = [] + self._token = None + self._done = True + else: + self._page = page["items"] + self._token = page["page_token"] + self._done = self._token is None + self._position = 0 + + async def _load_initial(self): + async with self._lock: + if self._failure is not None: + raise self._failure + if self._page is None: + await self._fetch_page() + return self + + def __await__(self): + return self._load_initial().__await__() + + def __aiter__(self): + return self + + async def __anext__(self) -> str: + async with self._lock: + if self._failure is not None: + raise StopAsyncIteration + if self._page is None: + await self._fetch_page() + assert self._page is not None + while self._position >= len(self._page): + if self._done: + raise StopAsyncIteration + await self._fetch_page() + assert self._page is not None + item = self._page[self._position] + self._position += 1 + return item + + +class AsyncTableListing(_AsyncPagedNameListing): + """Awaitable table iterator; awaiting loads one page. + + For example:: + + page = await db.list_tables(limit=20) + print(page.tables, page.page_token) + async for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + """A copy of every table name in the most recently fetched page.""" + return self._page_copy() + + +class AsyncNamespaceListing(_AsyncPagedNameListing): + """Awaitable namespace iterator; awaiting loads one page. + + For example:: + + page = await db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + """A copy of every name in the most recently fetched page.""" + return self._page_copy() + + +class _PagedNameListing(Iterator[str]): + def __init__(self, inner: _AsyncPagedNameListing): + self._inner = inner + LOOP.run(inner._load_initial()) + + def __iter__(self): + return self + + def __next__(self) -> str: + try: + return LOOP.run(self._inner.__anext__()) + except StopAsyncIteration: + raise StopIteration from None + + @property + def page_token(self) -> Optional[str]: + return self._inner.page_token + + def num_page_results(self) -> int: + return self._inner.num_page_results() + + +class TableListing(_PagedNameListing): + """Table iterator retaining the current page as a list. + + For example:: + + page = db.list_tables(limit=20) + print(page.tables, page.page_token) + for name in db.list_tables(limit=20): + print(name) + """ + + @property + def tables(self) -> list[str]: + return self._inner.tables + + +class NamespaceListing(_PagedNameListing): + """Namespace iterator retaining the current page as a list. + + For example:: + + page = db.list_namespaces(limit=20) + print(page.namespaces, page.page_token) + """ + + @property + def namespaces(self) -> list[str]: + return self._inner.namespaces diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index a5d88eb1f..55baa39a3 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -52,7 +52,14 @@ from lancedb.background_loop import LOOP from lancedb.arrow import AsyncRecordBatchReader from lancedb.db import AsyncConnection, DBConnection from lancedb.job import AsyncJob, Job -from lancedb.listing import AsyncListing, Listing +from lancedb.listing import ( + AsyncListing, + AsyncNamespaceListing, + AsyncTableListing, + Listing, + NamespaceListing, + TableListing, +) from lancedb.sql import AsyncQuery as AsyncSqlQuery from lancedb.sql import QueryDescription, QueryParameters from lance_namespace import ( @@ -61,8 +68,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, NamespaceExistsRequest, TableExistsRequest, ) @@ -854,7 +859,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """ List child namespaces under the given namespace. @@ -871,12 +876,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1023,7 +1028,7 @@ class LanceNamespaceDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """ List all tables in this database with pagination support. @@ -1040,12 +1045,12 @@ class LanceNamespaceDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1459,12 +1464,12 @@ class AsyncLanceNamespaceDBConnection: namespace_path = [] await self._inner.drop_all_tables(namespace_path=namespace_path) - async def list_namespaces( + def list_namespaces( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> AsyncNamespaceListing: """ List child namespaces under the given namespace. @@ -1481,12 +1486,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_namespaces( + return self._inner.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -1616,12 +1621,12 @@ class AsyncLanceNamespaceDBConnection: except TableNotFoundError: return False - async def list_tables( + def list_tables( self, namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> AsyncTableListing: """ List all tables in this database with pagination support. @@ -1638,12 +1643,12 @@ class AsyncLanceNamespaceDBConnection: Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return await self._inner.list_tables( + return self._inner.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 89c7f8a78..08125cc42 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -36,7 +36,7 @@ from lancedb.remote import ClientConfig, RetryConfig, TimeoutConfig, TlsConfig import pyarrow as pa from ..common import DATA -from lancedb.listing import Listing +from lancedb.listing import Listing, NamespaceListing, TableListing from ..db import DBConnection, LOOP from ..functions import FunctionVersion, UdfDefinition from ..job import AsyncJob, Job @@ -54,8 +54,6 @@ from lance_namespace import ( CreateNamespaceResponse, DescribeNamespaceResponse, DropNamespaceResponse, - ListNamespacesResponse, - ListTablesResponse, ) from ..pydantic import LanceModel from ..table import Table @@ -268,7 +266,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListNamespacesResponse: + ) -> NamespaceListing: """List immediate child namespace names in the given namespace. Parameters @@ -284,12 +282,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListNamespacesResponse + NamespaceListing Response containing namespace names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return NamespaceListing( self._conn.list_namespaces( namespace_path=namespace_path, page_token=page_token, limit=limit ) @@ -379,7 +377,7 @@ class RemoteDBConnection(DBConnection): namespace_path: Optional[List[str]] = None, page_token: Optional[str] = None, limit: Optional[int] = None, - ) -> ListTablesResponse: + ) -> TableListing: """List a page of tables in this database. Parameters @@ -396,12 +394,12 @@ class RemoteDBConnection(DBConnection): Returns ------- - ListTablesResponse + TableListing Response containing table names and optional page_token for pagination. """ if namespace_path is None: namespace_path = [] - return LOOP.run( + return TableListing( self._conn.list_tables( namespace_path=namespace_path, page_token=page_token, limit=limit ) diff --git a/python/python/tests/test_embeddings.py b/python/python/tests/test_embeddings.py index e319d3048..207a3d253 100644 --- a/python/python/tests/test_embeddings.py +++ b/python/python/tests/test_embeddings.py @@ -896,6 +896,8 @@ def test_transient_embedding_error_is_retried_in_table_add(tmp_path): assert len(calls) == 3 assert mock_sleep.call_count == 2 assert all(call.args[0] <= 60 for call in mock_sleep.call_args_list) + + @patch("time.sleep") @pytest.mark.parametrize( "error", @@ -941,6 +943,8 @@ def test_retry_on_wrapped_provider_error(mock_sleep): assert wrapped_func() == [0.1] assert mock_func.call_count == 2 assert mock_sleep.call_count == 1 + + def test_openai_bad_request_partial_batch_recovery(): # Regression test for issue #1677 diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 4a2745eb8..b0b973bbf 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -88,6 +88,143 @@ async def advance(items, asynchronous): KINDS = ["secrets", "views", "jobs", "functions", "materialized_views"] +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_compatible_name_listing_pages(listing_server, kind, asynchronous): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a", "b"], "next")), + (200, response_page(kind, ["c"], None)), + ] + ) + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)( + namespace_path=["team"], page_token="start", limit=2 + ) + field = kind + assert len(requests) == (0 if asynchronous else 1) + if asynchronous: + with pytest.raises(RuntimeError, match="Await or iterate"): + getattr(listing, field) + assert listing.page_token == "start" + assert await listing is listing + assert len(requests) == 1 + page_names = getattr(listing, field) + assert type(page_names) is list + assert page_names == ["a", "b"] + page_names.append("changed") + assert getattr(listing, field) == ["a", "b"] + assert listing.page_token == "next" + assert listing.num_page_results() == 2 + assert await advance(listing, asynchronous) == "a" + assert getattr(listing, field) == ["a", "b"] + assert await advance(listing, asynchronous) == "b" + assert len(requests) == 1 + assert await advance(listing, asynchronous) == "c" + assert getattr(listing, field) == ["c"] + assert listing.page_token is None + assert len(requests) == 2 + assert ([item async for item in listing] if asynchronous else list(listing)) == [] + if asynchronous: + assert await listing is listing + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_skips_empty_page_only_during_iteration( + listing_server, kind +): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, [], "next")), + (200, response_page(kind, ["a"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + await listing + assert getattr(listing, kind) == [] + assert listing.page_token == "next" + assert len(requests) == 1 + assert [name async for name in listing] == ["a"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_name_listing_zero_limit_uses_no_request( + listing_server, kind, asynchronous +): + endpoint, requests, _ = listing_server + db = await connection(endpoint, asynchronous) + listing = getattr(db, "list_" + kind)(limit=0) + if asynchronous: + assert await listing is listing + assert [name async for name in listing] == [] + else: + assert list(listing) == [] + assert getattr(listing, kind) == [] + assert listing.page_token is None + assert requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_failure_retains_token(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((401, {})) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(page_token="retry") + with pytest.raises(HttpError): + await listing + assert listing.page_token == "retry" + with pytest.raises(HttpError): + await listing + assert [name async for name in listing] == [] + assert len(requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_direct_async_iteration(listing_server, kind): + endpoint, requests, responses = listing_server + responses.extend( + [ + (200, response_page(kind, ["a"], "next")), + (200, response_page(kind, ["b"], None)), + ] + ) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)(limit=1) + assert requests == [] + assert [name async for name in listing] == ["a", "b"] + assert getattr(listing, kind) == ["b"] + assert len(requests) == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["tables", "namespaces"]) +async def test_name_listing_serializes_concurrent_advances(listing_server, kind): + endpoint, requests, responses = listing_server + responses.append((200, response_page(kind, ["a", "b"], None))) + db = await connection(endpoint, True) + listing = getattr(db, "list_" + kind)() + results = await asyncio.gather( + listing.__anext__(), + listing.__anext__(), + listing.__anext__(), + return_exceptions=True, + ) + assert results[:2] == ["a", "b"] + assert isinstance(results[2], StopAsyncIteration) + assert len(requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", KINDS) @pytest.mark.parametrize("asynchronous", [False, True]) diff --git a/python/python/tests/test_ollama_embeddings.py b/python/python/tests/test_ollama_embeddings.py index 9d8568ca7..de36dc24b 100644 --- a/python/python/tests/test_ollama_embeddings.py +++ b/python/python/tests/test_ollama_embeddings.py @@ -3,7 +3,6 @@ import json -import httpx import numpy as np import pyarrow as pa import pytest @@ -14,6 +13,7 @@ from lancedb.embeddings import get_registry @pytest.fixture def ollama_embeddings(): + httpx = pytest.importorskip("httpx") pytest.importorskip("ollama") requests = [] diff --git a/python/src/connection.rs b/python/src/connection.rs index baeb29780..2cdf71f4a 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -590,29 +590,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_namespaces( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListNamespacesRequest; - let request = ListNamespacesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_namespaces(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("namespaces", response.namespaces)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_namespaces( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } #[pyo3(signature = (namespace_path, mode=None, properties=None))] @@ -725,29 +713,17 @@ impl Connection { #[pyo3(signature = (namespace_path=None, page_token=None, limit=None))] pub fn list_tables( - self_: PyRef<'_, Self>, + &self, namespace_path: Option>, page_token: Option, limit: Option, - ) -> PyResult> { - let inner = self_.get_inner()?.clone(); - let py = self_.py(); - future_into_py(py, async move { - use lance_namespace::models::ListTablesRequest; - let request = ListTablesRequest { - id: namespace_path, - page_token, - limit: limit.map(|l| l as i32), - ..Default::default() - }; - let response = inner.list_tables(request).await.infer_error()?; - Python::attach(|py| -> PyResult> { - let dict = PyDict::new(py); - dict.set_item("tables", response.tables)?; - dict.set_item("page_token", response.page_token)?; - Ok(dict.unbind()) - }) - }) + ) -> PyResult { + Ok(crate::listing::NameListing::new( + self.get_inner()?.list_tables( + &namespace_path.unwrap_or_default(), + crate::listing::options(page_token, limit), + ), + )) } /// Get the configuration for constructing an equivalent namespace client. diff --git a/python/src/listing.rs b/python/src/listing.rs index 2d1a3916e..79a7effdd 100644 --- a/python/src/listing.rs +++ b/python/src/listing.rs @@ -3,10 +3,11 @@ use crate::error::PythonErrorExt; use crate::runtime::future_into_py; -use futures::{StreamExt, future::poll_fn}; +use futures::{Future, StreamExt, future::poll_fn}; use lancedb::listing::{Listing, ListingOptions}; use pyo3::exceptions::PyStopAsyncIteration; -use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyDict, PyDictMethods}; +use pyo3::{Bound, Py, PyAny, PyRef, PyResult, Python, pyclass, pymethods}; use std::sync::{Arc, Mutex}; pub fn options(page_token: Option, page_limit: Option) -> ListingOptions { @@ -52,6 +53,31 @@ impl NameListing { Ok(value) }) } + + fn next_page(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.inner.clone(); + let next_lock = self_.next_lock.clone(); + future_into_py(self_.py(), async move { + let _guard = next_lock.lock().await; + let page = futures::future::poll_fn(|cx| { + let mut inner = inner.lock().unwrap(); + let future = inner.next_page(); + futures::pin_mut!(future); + future.poll(cx) + }) + .await + .infer_error()?; + Python::attach(|py| -> PyResult>> { + page.map(|page| { + let dict = PyDict::new(py); + dict.set_item("items", page.items)?; + dict.set_item("page_token", page.page_token)?; + Ok(dict.unbind()) + }) + .transpose() + }) + }) + } } #[pyclass] diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 8a0e489c2..124f6a54a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -13,7 +13,7 @@ use lance::dataset::refs::MAIN_BRANCH; use lance_namespace::models::{ CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest, DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest, - ListNamespacesResponse, ListTablesRequest, ListTablesResponse, + ListTablesRequest, }; #[cfg(feature = "aws")] use object_store::aws::AwsCredential; @@ -28,7 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; -use crate::listing::{Listing, ListingOptions}; +use crate::listing::{Listing, ListingOptions, ListingPage}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -1126,12 +1126,43 @@ impl Connection { self.internal.drop_all_tables(namespace_path).await } - /// List immediate child namespace names in the given namespace - pub async fn list_namespaces( + /// Lazily list immediate child namespace names. Each request asks for at + /// most 100 names unless `options.page_limit` is set. + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_namespaces(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_namespaces( &self, - request: ListNamespacesRequest, - ) -> Result { - self.internal.list_namespaces(request).await + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_namespaces(ListNamespacesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.namespaces, + page_token: response.page_token, + }) + } + }) } /// Create a new namespace @@ -1182,23 +1213,43 @@ impl Connection { self.internal.namespace_client_config().await } - /// List tables with pagination support + /// Lazily list tables. Each request asks for at most 100 names unless + /// `options.page_limit` is set. /// - /// The default limit is 100 tables per page for both local and remote connections, - /// including namespaces. Zero returns an empty page without a continuation token. - /// Follow the response's opaque `page_token` until it is absent to retrieve every table; - /// a page can contain fewer than the limit even when more tables remain. - pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result { - // Apply the SDK's page size here so every backend receives the same request. - let limit = request.limit.get_or_insert(100); - if *limit <= 0 { - return Ok(ListTablesResponse { - context: None, - tables: Vec::new(), - page_token: None, - }); - } - self.internal.list_tables(request).await + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// use futures::TryStreamExt; + /// let names: Vec = connection.list_tables(&[], Default::default()) + /// .try_collect().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn list_tables( + &self, + namespace_path: &[String], + options: ListingOptions, + ) -> Listing { + let database = self.database().clone(); + let namespace_path = namespace_path.to_vec(); + Listing::new(options, move |options| { + let database = database.clone(); + let namespace_path = namespace_path.clone(); + async move { + validate_namespace(&namespace_path)?; + let response = database + .list_tables(ListTablesRequest { + id: Some(namespace_path), + page_token: options.page_token, + limit: Some(options.page_limit.unwrap_or(100) as i32), + ..Default::default() + }) + .await?; + Ok(ListingPage { + items: response.tables, + page_token: response.page_token, + }) + } + }) } /// Get the in-memory embedding registry. @@ -2260,26 +2311,11 @@ mod tests { } names.sort(); - // Walking in pages has to reach every table exactly once, with nothing lost at a - // page boundary. - let mut seen = Vec::with_capacity(names.len()); - let mut page_token = None; - loop { - let page = db - .list_tables(ListTablesRequest { - id: Some(Vec::new()), - limit: Some(2), - page_token, - ..Default::default() - }) - .await - .unwrap(); - seen.extend(page.tables); - page_token = page.page_token.filter(|token| !token.is_empty()); - if page_token.is_none() { - break; - } - } + let seen = futures::TryStreamExt::try_collect::>( + db.list_tables(&[], ListingOptions::default().page_limit(2)), + ) + .await + .unwrap(); assert_eq!(seen, names); } @@ -2305,37 +2341,20 @@ mod tests { .is_empty() ); - let first = db.list_tables(ListTablesRequest::default()).await.unwrap(); - assert_eq!(first.tables, names[..100]); + let mut listing = db.list_tables(&[], ListingOptions::default()); + assert_eq!(listing.num_page_results(), 0); + let first = listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items, names[..100]); assert!(first.page_token.is_some()); - let second = db - .list_tables(ListTablesRequest { - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(second.tables, names[100..]); + let second = listing.next_page().await.unwrap().unwrap(); + assert_eq!(second.items, names[100..]); assert!(second.page_token.is_none()); + assert!(listing.next_page().await.unwrap().is_none()); - let all = db - .list_tables(ListTablesRequest { - limit: Some(200), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(all.tables, names); - assert!(all.page_token.is_none()); - let empty = db - .list_tables(ListTablesRequest { - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(empty.tables.is_empty()); - assert!(empty.page_token.is_none()); + let mut all = db.list_tables(&[], ListingOptions::default().page_limit(200)); + assert_eq!(all.next_page().await.unwrap().unwrap().items, names); + let mut invalid = db.list_tables(&[], ListingOptions::default().page_limit(0)); + assert!(invalid.next_page().await.is_err()); } #[tokio::test] diff --git a/rust/lancedb/src/listing.rs b/rust/lancedb/src/listing.rs index 710202d9d..a0f6e5d5e 100644 --- a/rust/lancedb/src/listing.rs +++ b/rust/lancedb/src/listing.rs @@ -51,8 +51,9 @@ type FetchPage = /// A lazy stream of resources with inspectable pagination state. /// -/// No requests are made until the stream is polled. Empty pages with continuation -/// tokens are skipped. Errors terminate the stream and retain the failed request's +/// No requests are made until the stream is polled or [`Self::next_page`] is called. +/// Item iteration skips empty pages with continuation tokens; `next_page` returns +/// them. Errors terminate the listing and retain the failed request's /// token, allowing a new stream to retry it. Repeated tokens are rejected. /// /// A token advances when its page is fetched. Drain [`Self::num_page_results`] @@ -100,6 +101,64 @@ impl Listing { pub fn page_token(&self) -> Option<&str> { self.options.page_token.as_deref() } + + /// Return exactly one page without consuming a following page. + /// + /// If item iteration has started, this returns the remaining items in its + /// cached page. An empty page with a continuation token is returned as is. + pub async fn next_page(&mut self) -> Result>> { + if self.items.len() == 0 { + if self.done { + return Ok(None); + } + futures::future::poll_fn(|cx| self.poll_fetch_page(cx)).await?; + } + let items = std::mem::replace(&mut self.items, Vec::new().into_iter()).collect(); + Ok(Some(ListingPage { + items, + page_token: self.options.page_token.clone(), + })) + } + + fn poll_fetch_page(&mut self, cx: &mut Context<'_>) -> Poll> { + if self + .options + .page_limit + .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) + { + self.done = true; + return Poll::Ready(Err(Error::InvalidInput { + message: "Listing page limit must be between 1 and 2147483647".into(), + })); + } + let pending = self + .pending + .get_or_insert_with(|| (self.fetch)(self.options.clone())); + let response = futures::ready!(pending.as_mut().poll(cx)); + self.pending = None; + match response { + Ok(response) => { + let token = response.page_token.filter(|token| !token.is_empty()); + if token + .as_ref() + .is_some_and(|token| !self.seen_tokens.insert(token.clone())) + { + self.done = true; + return Poll::Ready(Err(Error::Runtime { + message: "Listing response repeated a page_token".into(), + })); + } + self.done = token.is_none(); + self.options.page_token = token; + self.items = response.items.into_iter(); + Poll::Ready(Ok(())) + } + Err(error) => { + self.done = true; + Poll::Ready(Err(error)) + } + } + } } impl Stream for Listing { @@ -114,41 +173,8 @@ impl Stream for Listing { if this.done { return Poll::Ready(None); } - if this - .options - .page_limit - .is_some_and(|limit| limit == 0 || limit > i32::MAX as u32) - { - this.done = true; - return Poll::Ready(Some(Err(Error::InvalidInput { - message: "Listing page limit must be between 1 and 2147483647".into(), - }))); - } - let pending = this - .pending - .get_or_insert_with(|| (this.fetch)(this.options.clone())); - let response = futures::ready!(pending.as_mut().poll(cx)); - this.pending = None; - match response { - Ok(response) => { - let token = response.page_token.filter(|token| !token.is_empty()); - if token - .as_ref() - .is_some_and(|token| !this.seen_tokens.insert(token.clone())) - { - this.done = true; - return Poll::Ready(Some(Err(Error::Runtime { - message: "Listing response repeated a page_token".into(), - }))); - } - this.done = token.is_none(); - this.options.page_token = token; - this.items = response.items.into_iter(); - } - Err(error) => { - this.done = true; - return Poll::Ready(Some(Err(error))); - } + if let Err(error) = futures::ready!(this.poll_fetch_page(cx)) { + return Poll::Ready(Some(Err(error))); } } } @@ -224,6 +250,36 @@ mod tests { assert!(listing.next().await.is_none()); } + #[tokio::test] + async fn next_page_preserves_empty_pages_and_remaining_items() { + let requests = Arc::new(Mutex::new(0)); + let count = requests.clone(); + let mut listing = Listing::new(ListingOptions::default(), move |options| { + *count.lock().unwrap() += 1; + async move { + let (items, page_token) = match options.page_token.as_deref() { + None => (vec![], Some("second".to_string())), + Some("second") => (vec!["a", "b"], Some("last".to_string())), + Some("last") => (vec!["c"], None), + _ => unreachable!(), + }; + Ok(ListingPage { items, page_token }) + } + }); + assert_eq!(*requests.lock().unwrap(), 0); + assert!(listing.next_page().await.unwrap().unwrap().items.is_empty()); + assert_eq!(*requests.lock().unwrap(), 1); + assert_eq!(listing.try_next().await.unwrap(), Some("a")); + assert_eq!(*requests.lock().unwrap(), 2); + let remaining = listing.next_page().await.unwrap().unwrap(); + assert_eq!(remaining.items, vec!["b"]); + assert_eq!(remaining.page_token.as_deref(), Some("last")); + assert_eq!(*requests.lock().unwrap(), 2); + assert_eq!(listing.next_page().await.unwrap().unwrap().items, vec!["c"]); + assert!(listing.next_page().await.unwrap().is_none()); + assert_eq!(*requests.lock().unwrap(), 3); + } + #[tokio::test] async fn invalid_limits_do_not_fetch() { for limit in [0, u32::MAX] { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 4a8cbf32e..8a456ab3b 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -2349,16 +2349,8 @@ mod tests { .unwrap() .is_empty() ); - let page = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - limit: Some(0), - ..Default::default() - }) - .await - .unwrap(); - assert!(page.tables.is_empty()); - assert!(page.page_token.is_none()); + let mut listing = conn.list_tables(&namespace, ListingOptions::default().page_limit(0)); + assert!(listing.next_page().await.is_err()); } #[rstest::rstest] @@ -2393,25 +2385,17 @@ mod tests { ) .unwrap() }); - let first = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(first.tables.len(), 100); + let mut first_listing = conn.list_tables(&namespace, ListingOptions::default()); + let first = first_listing.next_page().await.unwrap().unwrap(); + assert_eq!(first.items.len(), 100); assert_eq!(first.page_token.as_deref(), Some("opaque-token")); - let second = conn - .list_tables(lance_namespace::models::ListTablesRequest { - id: Some(namespace), - page_token: first.page_token, - ..Default::default() - }) - .await - .unwrap(); + let mut second_listing = conn.list_tables( + &namespace, + ListingOptions::default().page_token(first.page_token.unwrap()), + ); + let second = second_listing.next_page().await.unwrap().unwrap(); let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect(); - assert_eq!(second.tables, expected); + assert_eq!(second.items, expected); assert!(second.page_token.is_none()); } @@ -3680,14 +3664,9 @@ mod tests { assert!(table.is_ok(), "Failed to create table: {:?}", table.err()); // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .expect("Failed to list tables"); - assert_eq!(list_response.tables, vec!["test_table"]); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names, vec!["test_table"]); // Get namespace client and verify it can also list tables let namespace_client = conn.namespace_client().await.unwrap(); @@ -3754,17 +3733,12 @@ mod tests { } // List tables in the child namespace - let list_response = conn - .list_tables(ListTablesRequest { - id: Some(namespace.clone()), - ..Default::default() - }) - .await - .unwrap(); - assert_eq!(list_response.tables.len(), 3); - assert!(list_response.tables.contains(&"table1".to_string())); - assert!(list_response.tables.contains(&"table2".to_string())); - assert!(list_response.tables.contains(&"table3".to_string())); + let list_response = conn.list_tables(&namespace, ListingOptions::default()); + let names = list_response.try_collect::>().await.unwrap(); + assert_eq!(names.len(), 3); + assert!(names.contains(&"table1".to_string())); + assert!(names.contains(&"table2".to_string())); + assert!(names.contains(&"table3".to_string())); } } From 444f7f3bde8d18aa6ef39cba714dbe1bb448fb1d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 12:45:25 +0800 Subject: [PATCH 62/68] ci: refresh apt lists before installing protobuf in docs build (#4636) The docs deploy workflow installs protobuf-compiler without refreshing apt lists first. Ubuntu noble-updates replaced `protobuf 3.21.12-8.2ubuntu0.3` with `.0.4`, so stale package lists in older runner images make `apt install` fail with 404 and exit code 100. Run `apt update` before installing, matching the other workflows. --- .github/workflows/docs.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index d0ec583bf..fb3fb1d6a 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -38,6 +38,7 @@ jobs: uses: actions/checkout@v6 - name: Install dependencies needed for ubuntu run: | + sudo apt update sudo apt install -y protobuf-compiler libssl-dev rustup update && rustup default - name: Set up Python From 45c4b159cab990a749d592b42e21e030f7479192 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Thu, 22 Oct 2026 22:51:40 -0700 Subject: [PATCH 63/68] fix: cap the remote job poll interval at one second (#4629) job.wait() polled on a backoff that doubled up to 5 s, so a job's completion was seen 2.5 s late on average, as long as a small refresh itself takes. Cap the interval at 1 s; the first polls still back off from 200 ms, so a quick job is seen quickly and a long one costs one request a second. Linear: OSS-2704 https://linear.app/lancedb/issue/OSS-2704 --- rust/lancedb/src/remote/job.rs | 36 ++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/rust/lancedb/src/remote/job.rs b/rust/lancedb/src/remote/job.rs index 41d8971ac..7d2e70ac9 100644 --- a/rust/lancedb/src/remote/job.rs +++ b/rust/lancedb/src/remote/job.rs @@ -18,7 +18,9 @@ use crate::remote::client::{HttpSend, RequestResultExt, RestfulLanceDbClient}; /// Delay before the second job-state poll; doubles up to [`MAX_POLL_INTERVAL`]. const INITIAL_POLL_INTERVAL: Duration = Duration::from_millis(200); -const MAX_POLL_INTERVAL: Duration = Duration::from_secs(5); +/// A job's completion is seen at most this long after it happens. A 5 s cap +/// added 2.5 s on average to every short refresh, as much as the job itself. +const MAX_POLL_INTERVAL: Duration = Duration::from_secs(1); #[derive(Debug, Clone, PartialEq, Eq)] enum JobState { @@ -268,13 +270,18 @@ impl JobHandle for RemoteJob { #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + use async_trait::async_trait; use crate::Result; use crate::function::{FunctionVersion, RefreshColumnResult}; use crate::job::{Job, JobHandle, TerminalResult}; + use crate::remote::client::test_utils::client_with_handler; - use super::DescribeJobResponse; + use super::{DescribeJobResponse, RemoteJob}; const FUNCTION_JOB: &str = include_str!("../../tests/fixtures/first_class_functions/v1/remote_function_job.json"); @@ -333,6 +340,31 @@ mod tests { ); } + /// Five in-progress polls then done: 0.2 + 0.4 + 0.8 + 1 + 1 s of waiting. + /// Under the old 5 s cap the same job would have waited 10 s. + #[tokio::test(start_paused = true)] + async fn wait_polls_at_most_every_second() { + let describes = Arc::new(AtomicUsize::new(0)); + let seen = describes.clone(); + let client = client_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/jobs/describe"); + let n = seen.fetch_add(1, Ordering::SeqCst); + let body = if n < 5 { + r#"{"job_state":"IN_PROGRESS"}"#.to_string() + } else { + UNIT_JOB.to_string() + }; + http::Response::builder().status(200).body(body).unwrap() + }); + let started = tokio::time::Instant::now(); + RemoteJob::new(client, "job-1".to_string()) + .wait() + .await + .expect("unit job completes"); + assert_eq!(describes.load(Ordering::SeqCst), 6); + assert_eq!(started.elapsed(), Duration::from_millis(3400)); + } + #[test] fn remote_wire_unknown_fields_are_forward_decodable() { let response: DescribeJobResponse = From a9535684cc1d9901c54c852f103dd0d492d61c3f Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 23 Oct 2026 20:17:38 +0800 Subject: [PATCH 64/68] chore(deps): bump ctutils to 0.4.4 for RUSTSEC-2026-0385 (#4638) Fixes RUSTSEC-2026-0385: `ctutils` < 0.4.4 generates masks with branches on secret values (`u128` mask generation / comparisons are not constant-time on some targets). - Advisory: https://rustsec.org/advisories/RUSTSEC-2026-0385 - Affected: `ctutils 0.4.2` pinned in `Cargo.lock`, reached transitively via `digest 0.11` (RustCrypto) - Change: `cargo update -p ctutils --precise 0.4.4`; single lockfile-only change (patch bump, no API impact) The `deny` CI job started failing for every PR head after the advisory was published at 06:45Z today; `main`'s last `deny` run (05:51Z) predates it and does not yet cover this advisory. --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 34ee05cb1..8bad6ca92 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2230,9 +2230,9 @@ dependencies = [ [[package]] name = "ctutils" -version = "0.4.2" +version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +checksum = "12218662ccac37d9714f2191252553d67585b97df8c839cbe49e0bde278dc865" dependencies = [ "cmov", ] From b762ec221ba5d99ab368f2959f50381c589c762b Mon Sep 17 00:00:00 2001 From: LuQQiu Date: Fri, 23 Oct 2026 10:06:20 -0700 Subject: [PATCH 65/68] docs: add Enterprise distributed search reference page (#4625) Adds the Enterprise reference page at `docs/web/enterprise/distributed-search.mdx`, following the same arrangement as Enterprise architecture. The page describes distributed search and the embedded execution model, with links to shared search guides. Enterprise architecture links to it, and the official sidebar lists it immediately after Architecture in the existing Deployment group. Adds stable anchors for the RaBitQ and `approx_mode` sections linked by the Enterprise guide. Companion: https://github.com/lancedb/sophon/pull/8986 supplies the detailed Enterprise overlay at the same path. Merge this reference-page PR first so the overlay has an existing route to replace. Validation: - Combined-site assembly with both PR branches passed (309 files). - Mintlify build validation and link, anchor, and redirect checks passed. --- docs/web/docs.json | 1 + docs/web/enterprise/architecture.mdx | 3 ++ docs/web/enterprise/distributed-search.mdx | 33 ++++++++++++++++++++++ docs/web/indexing/vector-index.mdx | 2 +- docs/web/search/vector-search.mdx | 2 +- 5 files changed, 39 insertions(+), 2 deletions(-) create mode 100644 docs/web/enterprise/distributed-search.mdx diff --git a/docs/web/docs.json b/docs/web/docs.json index 9c9468a3b..54203f1a0 100644 --- a/docs/web/docs.json +++ b/docs/web/docs.json @@ -350,6 +350,7 @@ "pages": [ "enterprise/index", "enterprise/architecture", + "enterprise/distributed-search", "enterprise/deployment/index", "deployment/storage", "deployment/monitoring", diff --git a/docs/web/enterprise/architecture.mdx b/docs/web/enterprise/architecture.mdx index c47dcbd2e..8f887c142 100644 --- a/docs/web/enterprise/architecture.mdx +++ b/docs/web/enterprise/architecture.mdx @@ -18,3 +18,6 @@ run on machines that can be replaced without any of them owning the data. [Architecture](/architecture) covers this alongside the embedded form, which is the comparison most readers want first. + +See [Distributed search](/enterprise/distributed-search) for how Enterprise +distributes index execution and table-data reads. diff --git a/docs/web/enterprise/distributed-search.mdx b/docs/web/enterprise/distributed-search.mdx new file mode 100644 index 000000000..ebac960ba --- /dev/null +++ b/docs/web/enterprise/distributed-search.mdx @@ -0,0 +1,33 @@ +--- +title: "Distributed search" +sidebarTitle: "Distributed search" +description: "How LanceDB Enterprise distributes index execution and table-data reads." +icon: "network-wired" +availability: + oss: unavailable + enterprise: available + summary: >- + LanceDB Enterprise distributes index execution and table-data reads across + plan executors. Embedded LanceDB executes searches in the application process. +--- + +Distributed search spreads index execution, table-data reads, and caching across +multiple serving machines in LanceDB Enterprise. Query nodes plan queries and +combine results. Plan executors cache index structures and table data and execute +the remote work. See [Enterprise architecture](/enterprise/architecture). + +## Embedded search {#embedded-search} + +LanceDB OSS runs inside your application process. It searches indexes and reads +table data using that process's resources, including its index cache. Reading a +table from object storage does not by itself distribute query execution across +machines. + +## Search interfaces {#search-interfaces} + +The shared search guides describe the query APIs and index settings: + +- [Vector search](/search/vector-search) and [vector indexes](/indexing/vector-index). +- [Full-text search](/search/full-text-search) and [FTS indexes](/indexing/fts-index). +- [Filtering](/search/filtering) and [scalar indexes](/indexing/scalar-index). +- [Query plans](/search/explain) for inspecting execution and runtime measurements. diff --git a/docs/web/indexing/vector-index.mdx b/docs/web/indexing/vector-index.mdx index 2742a992f..b26e0c295 100644 --- a/docs/web/indexing/vector-index.mdx +++ b/docs/web/indexing/vector-index.mdx @@ -384,7 +384,7 @@ The inner workings of RaBitQ quantization are rather mathematically dense. It ge applying a uniformly random, approximately distance-preserving orthogonal transformation of the vertices of the $d$-dimensional hypercube, where $d$ is the dimensionality of the dataset. We defer the details, and an elegant theoretical error bound, to [the original paper.](https://arxiv.org/pdf/2405.12497) -##### Using RaBitQ +##### Using RaBitQ {#using-rabitq} Use RaBitQ quantization by selecting quantized index types ending in the suffix `RQ`. For example, call `create_index` with `index_type="IVF_RQ"`. Note that when using `IVF_RQ`, the dimension of the dataset must be a multiple of `8`. diff --git a/docs/web/search/vector-search.mdx b/docs/web/search/vector-search.mdx index 3519df553..cc6ae5421 100644 --- a/docs/web/search/vector-search.mdx +++ b/docs/web/search/vector-search.mdx @@ -235,7 +235,7 @@ The table below summarizes the behavior of `_distance` in search results based o See [vector indexes](/indexing/vector-index) for index configuration. The following sections cover query-time settings. -### Tuning `approx_mode` +### Tuning `approx_mode` {#tuning-approx-mode} Use `approx_mode` when you want to adjust the speed/recall trade-off for approximate vector search at query time. This setting currently applies only to RQ-quantized indexes, such as `IVF_RQ`; other index types ignore it. From 5a79c7d22da9f9729d27adc7e9d0b166767bbd4e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 01:52:45 +0800 Subject: [PATCH 66/68] ci: run aarch64 Windows tests on main only (#4637) Run the aarch64 Windows build/test cycle only on pushes to main so PRs wait for a single Windows platform (x86_64). Main keeps the same native aarch64 coverage as before. GitHub evaluates job-level `if` before the matrix strategy is applied, so the aarch64 target is split into its own `windows-arm` job instead of being filtered inside the existing matrix. The job keeps the historical check name. --- .github/workflows/rust.yml | 43 ++++++++++++++++++++++++++++++++------ 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index cbfefc83d..930ae67e2 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -245,12 +245,6 @@ jobs: include: - target: x86_64-pc-windows-msvc runner: windows-2022 - # windows-11-arm is a standard runner, so it is free on public repos. - # Running natively lets the aarch64 tests actually execute -- this - # job used to cross-compile them and then skip the test step, paying - # full codegen and link cost for a compile check. - - target: aarch64-pc-windows-msvc - runner: windows-11-arm runs-on: ${{ matrix.runner }} defaults: run: @@ -279,6 +273,43 @@ jobs: # dependency graph a second time. cargo test --profile ci --features aws,remote --locked --target ${{ matrix.target }} + # Run the aarch64 Windows tests only on main so PRs wait for a single + # Windows platform. windows-11-arm is a standard runner, so it is free on + # public repos. Running natively lets the aarch64 tests actually execute -- + # this job used to cross-compile them and then skip the test step, paying + # full codegen and link cost for a compile check. + windows-arm: + # Keep the historical name so the check shows up as before on main. + name: windows (aarch64-pc-windows-msvc, windows-11-arm) + if: github.ref == 'refs/heads/main' + runs-on: windows-11-arm + defaults: + run: + working-directory: rust/lancedb + steps: + - uses: actions/checkout@v6 + - name: Set target + run: rustup target add aarch64-pc-windows-msvc + - uses: Swatinem/rust-cache@v2 + with: + # Restore everywhere, but only save from main. Per-PR saves are + # unreadable outside their own branch anyway, since GitHub scopes + # caches to the creating ref. + save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Install Protoc v21.12 + run: choco install --no-progress protoc + - name: Build + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + cargo build --profile ci --features aws,remote --tests --locked --target aarch64-pc-windows-msvc + - name: Run tests + run: | + $env:VCPKG_ROOT = $env:VCPKG_INSTALLATION_ROOT + # `--target` has to match the build step above. Without it cargo uses + # target/ci/ rather than target//ci/ and rebuilds the entire + # dependency graph a second time. + cargo test --profile ci --features aws,remote --locked --target aarch64-pc-windows-msvc + msrv: # Check the minimum supported Rust version name: MSRV Check - Rust v${{ matrix.msrv }} From 702389654e11d1af5cd120d64866f028f41d6b90 Mon Sep 17 00:00:00 2001 From: KarthikS Date: Fri, 23 Oct 2026 15:27:32 -0400 Subject: [PATCH 67/68] fix(python): align pylance extra and preview install indexes (#4556) Fixes #4248. The preview index still contains pylance 1.0.0b6. uv's default first-index resolution can select it, or backtrack to an old LanceDB release when installing the extra. Meanwhile, the extra still permits pylance 5 even though current main bundles Lance 14. Raise the optional dependency to `pylance>=14.0.0b12`, the version already pinned in the test extra, and keep the lockfile metadata consistent. Document both Fury indexes and uv's `unsafe-best-match` strategy in the Python README and the source-owned quickstart (`docs/web/quickstart.mdx`). Explain the index trust requirement and recommend installing the extra to enforce the version floor. Add CI checks that reject older Lance majors and ensure the pylance and tests extras remain compatible. Run these checks when either dependency manifest changes, so a future Lance major bump cannot silently leave the Python minimum behind. ### Validation - Recreated the reported install on Ubuntu 26.04 / CPython 3.12.15 with uv 0.12.23: `lancedb==0.40.0b5` plus the old index selects `pylance==1.0.0b6`. Reproduced both the missing `blob_array` import and the unexpected `blob_handling` scanner argument using a real blob table. - Resolved the modified `python/pyproject.toml` with the pylance extra against both indexes: selects `pylance==14.0.0b18`, matching the current Rust pin. - Ran both documented pip and uv installations in isolated environments: `lancedb==0.41.0b7`, `pylance==14.0.0b18`; blob-array construction and a blob-bytes round trip through `to_pandas(blob_mode="bytes")` pass. Both package dependency checks pass. - New regression fails on the original dependency bound (27 rejected-version cases incorrectly accepted), then passes with the fix. Full CI-script suite: 3 tests pass. - `uv lock --check` with both indexes passes; focused Ruff checks pass; `git diff --check` passes. - Repository-wide Ruff reports 49 existing errors in untouched `docs/web-tests` files. Unrelated formatting changes were discarded. Runtime smoke checks used released native wheels; source dependency resolution was tested directly from the modified pyproject. No Rust or Python runtime code changes. --------- Co-authored-by: Xuanwo --- .../lancedb-update-lance-dependency/SKILL.md | 27 +++++++++++ .github/workflows/ci-scripts.yml | 6 +++ ci/tests/test_python_dependencies.py | 47 +++++++++++++++++++ docs/web/quickstart.mdx | 27 +++++++++++ python/README.md | 27 +++++++++++ python/pyproject.toml | 5 +- python/uv.lock | 18 +++---- 7 files changed, 146 insertions(+), 11 deletions(-) create mode 100644 ci/tests/test_python_dependencies.py diff --git a/.agents/skills/lancedb-update-lance-dependency/SKILL.md b/.agents/skills/lancedb-update-lance-dependency/SKILL.md index e97fae7ac..3cba2ae75 100644 --- a/.agents/skills/lancedb-update-lance-dependency/SKILL.md +++ b/.agents/skills/lancedb-update-lance-dependency/SKILL.md @@ -48,6 +48,23 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r This updates the Rust workspace Lance dependencies through `ci/set_lance_version.py`, updates `java/pom.xml`, refreshes Cargo metadata, and prints JSON metadata containing `branch_name`, `commit_message`, and `pr_title`. + Also synchronize the Python requirements in `python/pyproject.toml`: set the + `tests` extra's exact `pylance` pin to the matching published Python release, + and update the `pylance` extra's minimum when needed to exclude older bundled + Lance majors. The minimum must still accept the test pin. Python prerelease + spelling differs from Rust (for example, `15.0.0-beta.1` becomes `15.0.0b1`). + Confirm the matching Python release is published before changing the pin. + + Refresh and check `python/uv.lock` from the repository root: + + ```bash + uv lock --project python --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + uv lock --project python --check --index https://pypi.fury.io/lancedb/ --index https://pypi.fury.io/lance-format/ --index-strategy unsafe-best-match + ``` + + This strategy considers versions across indexes; use it only with trusted + indexes. Commit the Python requirements and lockfile together with the bump. + 5. Run validation: ```bash @@ -57,6 +74,16 @@ Inputs can be a version (`7.2.0-beta.1`), a tag (`v7.2.0-beta.1`), a tag ref (`r Fix real diagnostics and rerun clippy until it succeeds. Do not skip warnings. + Run the Python dependency guards with Python 3.11+ and `packaging` installed: + + ```bash + python3 -m unittest discover -s ci/tests -p test_python_dependencies.py -v + ``` + + These checks reject older Lance majors and ensure the `pylance` and `tests` + extras remain compatible. The full CI-script suite uses + `python3 -m unittest discover -s ci/tests -v`. + 6. Inspect `git status --short` and `git diff` to ensure only the Lance dependency update and required compatibility fixes are present. 7. If the task only asks to prepare local changes, stop here and report the changed files and validation result. diff --git a/.github/workflows/ci-scripts.yml b/.github/workflows/ci-scripts.yml index f7f0dbb07..d45289c2d 100644 --- a/.github/workflows/ci-scripts.yml +++ b/.github/workflows/ci-scripts.yml @@ -5,6 +5,8 @@ on: branches: - main paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -12,6 +14,8 @@ on: - .github/workflows/ci-scripts.yml pull_request: paths: + - Cargo.toml + - python/pyproject.toml - ci/check_cargo_lock.py - ci/check_sdist.py - ci/set_lance_version.py @@ -30,5 +34,7 @@ jobs: - uses: actions/setup-python@v6 with: python-version: "3.13" + - name: Install test dependencies + run: python -m pip install packaging - name: Run tests run: python -m unittest discover -s ci/tests -v diff --git a/ci/tests/test_python_dependencies.py b/ci/tests/test_python_dependencies.py new file mode 100644 index 000000000..97fb672c6 --- /dev/null +++ b/ci/tests/test_python_dependencies.py @@ -0,0 +1,47 @@ +import tomllib +import unittest +from pathlib import Path + +from packaging.requirements import Requirement +from packaging.version import Version + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +class PythonDependenciesTest(unittest.TestCase): + def setUp(self): + with (REPO_ROOT / "python/pyproject.toml").open("rb") as f: + self.project = tomllib.load(f)["project"] + self.pylance = Requirement(self.project["optional-dependencies"]["pylance"][0]) + + def test_pylance_extra_excludes_older_lance_majors(self): + with (REPO_ROOT / "Cargo.toml").open("rb") as f: + lance = tomllib.load(f)["workspace"]["dependencies"]["lance"] + if "version" not in lance: + self.skipTest("Local Lance checkout has no version requirement") + major = Version(lance["version"].lstrip("=")).major + for older_major in range(major): + for suffix in (".0.0", ".0.0b6", ".999.999"): + version = f"{older_major}{suffix}" + with self.subTest(version=version): + self.assertFalse( + self.pylance.specifier.contains(version, prereleases=True), + f"The pylance extra must exclude Lance {version}; " + f"the bundled Lance is major {major}", + ) + + def test_pylance_extra_accepts_tested_version(self): + requirements = [ + Requirement(dep) for dep in self.project["optional-dependencies"]["tests"] + ] + tested = next(dep for dep in requirements if dep.name == "pylance") + pin = next(spec for spec in tested.specifier if spec.operator == "==") + self.assertTrue( + self.pylance.specifier.contains(pin.version, prereleases=True), + f"The tests and pylance extras must be installable together: {tested}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/web/quickstart.mdx b/docs/web/quickstart.mdx index 788dd6585..3f466041d 100644 --- a/docs/web/quickstart.mdx +++ b/docs/web/quickstart.mdx @@ -58,6 +58,33 @@ uv add --prerelease allow --index https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra and +include the Lance preview index: + + +```bash Python (pip) icon="terminal" +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +```bash Python (uv) icon="terminal" +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + Pre-release builds receive the same level of testing as stable releases, but their availability is not guaranteed for more than 6 months after release. For real-world workloads, we recommend you use the latest stable release diff --git a/python/README.md b/python/README.md index f941c9e20..907f7c099 100644 --- a/python/README.md +++ b/python/README.md @@ -40,6 +40,33 @@ Stable releases are created about every 2 weeks. For the latest features and bug pip install --pre --extra-index-url https://pypi.fury.io/lancedb/ lancedb ``` +If you also need the optional `lance` Python API, install the `pylance` extra +and include the Lance preview index: + +```bash +pip install --pre \ + --extra-index-url https://pypi.fury.io/lancedb/ \ + --extra-index-url https://pypi.fury.io/lance-format/ \ + "lancedb[pylance]" +``` + +With uv, also select the best matching version across the indexes: + +```bash +uv pip install --prerelease allow \ + --index https://pypi.fury.io/lancedb/ \ + --index https://pypi.fury.io/lance-format/ \ + --index-strategy unsafe-best-match \ + "lancedb[pylance]" +``` + +The LanceDB preview index contains old `pylance` releases; current Lance previews +are published to the `lance-format` index. By default, uv only considers versions +from the first index containing a package, which can select an old `pylance` or +backtrack to an older LanceDB release. `unsafe-best-match` searches all configured +indexes, like pip; use it only with indexes you trust. Install the extra rather +than an unconstrained `pylance` so LanceDB's minimum supported version is enforced. + ### Threading in CPU-limited containers LanceDB uses separate pools for compute work and storage I/O. On a container with diff --git a/python/pyproject.toml b/python/pyproject.toml index ed54ae536..a0de6509c 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -45,7 +45,8 @@ repository = "https://github.com/lancedb/lancedb" [project.optional-dependencies] pylance = [ - "pylance>=5.0.0b5", + # Keep the minimum in the bundled Lance major and compatible with tests below. + "pylance>=15.0.0b1", ] # A library only needs the OpenTelemetry API; the application supplies and # configures the SDK (the actual exporter/reader). See @@ -63,7 +64,7 @@ tests = [ "polars>=0.19, <=1.32.3", "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance==14.0.0b12", + "pylance==15.0.0b1", "requests>=2.31.0", "datafusion>=54,<55", "opentelemetry-sdk>=1.30.0", diff --git a/python/uv.lock b/python/uv.lock index e808f5f9e..99bcb439c 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -2006,8 +2006,8 @@ requires-dist = [ { name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" }, { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=2.7.4,<3" }, - { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = "==14.0.0b12" }, + { name = "pylance", marker = "extra == 'pylance'", specifier = ">=15.0.0b1" }, + { name = "pylance", marker = "extra == 'tests'", specifier = "==15.0.0b1" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" }, { name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" }, @@ -3938,7 +3938,7 @@ crypto = [ [[package]] name = "pylance" -version = "14.0.0b12" +version = "15.0.0b1" source = { registry = "https://pypi.fury.io/lance-format/" } dependencies = [ { name = "lance-namespace" }, @@ -3948,12 +3948,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_1Ixasc/pylance-14.0.0b12-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:84350a3c5c1eea3398f4c4fad7cfe85b831d831a5e43b5ca2a783dbdb0e11969" }, - { url = "https://pypi.fury.io/lance-format/-/ver_sCEkG/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76b386c182a6a8e7b86fbef28f84d1dd39e4626ac8681a0cf36f40a45dd7cb87" }, - { url = "https://pypi.fury.io/lance-format/-/ver_1vdT1C/pylance-14.0.0b12-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6aa548c7937fcdfa8f1943a2d59041a42f8b500199880dbb509c7195fa253381" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2fXu1V/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:79f82d6a6c50f37aab338470e3733436e6a27fbb5972d29fcf453ece56d3b3c6" }, - { url = "https://pypi.fury.io/lance-format/-/ver_24UavN/pylance-14.0.0b12-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:de0305999257002785f60169e6926896584d66f1da050bb5afd1d48c6e3d9bb9" }, - { url = "https://pypi.fury.io/lance-format/-/ver_XNynX/pylance-14.0.0b12-cp310-abi3-win_amd64.whl", hash = "sha256:4d0cdb689e898ce71e8e434d00a3b97639cea8fb97da21f1c921543be363268f" }, + { url = "https://pypi.fury.io/lance-format/-/ver_23L1Nd/pylance-15.0.0b1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:19605d3daca1224cc0db079054be5d4eec403dcf286bf4617430bfc259ab0a52" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1Ho3RK/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be22b0b6f1e1d250ffaa82b449e982c817db33103c483d1005020285310623cf" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1QmyF0/pylance-15.0.0b1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5b2761d60b1b9db112058213f1f7b2de82b08cbe65888fab08b0a56e702482e" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1in2Ph/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:294a0fedae312ed32d1b66cddbf7e69770ec9b4c82a873af60c619ed88450718" }, + { url = "https://pypi.fury.io/lance-format/-/ver_21xpLH/pylance-15.0.0b1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:ff20a16703d4d3f3bc675b25ce1871ba1d54137929dc46305bffdd07f7b35a6b" }, + { url = "https://pypi.fury.io/lance-format/-/ver_AAGzF/pylance-15.0.0b1-cp310-abi3-win_amd64.whl", hash = "sha256:3b134e24a5becfc940ca134cf78aa3bae8b80ec40f8927a1b856273ab76213f2" }, ] [[package]] From e74a7c15b5f322f8ac15b96bdfc391616518d711 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Sat, 24 Oct 2026 03:47:16 +0800 Subject: [PATCH 68/68] fix: bump tokio to address RUSTSEC-2026-0386 (#4642) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upgrade the locked Tokio version from 1.53.1 to 1.53.3 to fix [RUSTSEC-2026-0386](https://rustsec.org/advisories/RUSTSEC-2026-0386) and restore the Rust workflow's `deny` job. The diff changes only Tokio's version and checksum. Validation: - `cargo deny check advisories` — passed. - `cargo deny check advisories bans licenses sources` — passed with the repository's `deny.toml`. - `cargo check --quiet -p lancedb` — passed (dead-code warnings). - PR `deny` job — [passed](https://github.com/lancedb/lancedb/actions/runs/40000920646/job/120576455117). --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8bad6ca92..fac89c4da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10446,9 +10446,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" dependencies = [ "bytes", "libc",