From 2c2bb9a974fd10b8f75d62f09d170c19d1d1b07d Mon Sep 17 00:00:00 2001 From: Will Jones Date: Sat, 10 Oct 2026 17:12:13 +0000 Subject: [PATCH] feat(python): pass merge_insert sources as scannables Python's `execute_merge_insert` now takes a `PyScannable` (like `add()`), and the source collection threshold is exposed as a `source_collect_threshold_bytes` kwarg on `execute()` and a `merge_insert_source_collect_threshold_bytes` connection option. Node passes the decoded IPC batches as a `Vec`, so the merge join sees their exact row count. Co-Authored-By: Claude Opus 5.5 --- nodejs/__test__/table.test.ts | 10 +++++ nodejs/src/merge.rs | 20 ++++++++- python/python/lancedb/__init__.py | 19 +++++++++ python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/db.py | 9 ++++ python/python/lancedb/merge.py | 18 ++++++++ python/python/lancedb/table.py | 5 ++- python/python/tests/test_table.py | 67 ++++++++++++++++++++++++++++++ python/src/connection.rs | 6 ++- python/src/table.rs | 13 +++--- 10 files changed, 157 insertions(+), 11 deletions(-) diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index d2244b48c..fde773a43 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -1081,6 +1081,16 @@ describe("merge insert", () => { expect(res).toEqual(expected); }); + test("empty source with delete by source", async () => { + const empty = arrow.makeEmptyTable(await table.schema()); + const res = await table + .mergeInsert("a") + .whenNotMatchedBySourceDelete() + .execute(empty); + expect(res.numDeletedRows).toBe(3); + expect(await table.countRows()).toBe(0); + }); + test("timeout", async () => { const newData = [ { a: 2, b: "x" }, diff --git a/nodejs/src/merge.rs b/nodejs/src/merge.rs index ac8649419..bb0701a3d 100644 --- a/nodejs/src/merge.rs +++ b/nodejs/src/merge.rs @@ -3,6 +3,7 @@ use std::time::Duration; +use arrow_array::RecordBatch; use lancedb::{ipc::ipc_file_to_batches, table::merge::MergeInsertBuilder}; use napi::bindgen_prelude::*; use napi_derive::napi; @@ -66,9 +67,26 @@ impl NativeMergeInsertBuilder { #[napi(catch_unwind)] pub async fn execute(&self, buf: Buffer) -> napi::Result { - let data = ipc_file_to_batches(buf.to_vec()).map_err(|e| { + let reader = ipc_file_to_batches(buf.to_vec()).map_err(|e| { napi::Error::from_reason(format!("Failed to read IPC file: {}", convert_error(&e))) })?; + let schema = reader.schema(); + // Decoded batches are rescannable, so the merge join sees their exact + // row count. + let mut data = reader + .map(|batch| { + batch.map_err(|e| { + napi::Error::from_reason(format!( + "Failed to read record batch from IPC file: {}", + e + )) + }) + }) + .collect::>>()?; + // An empty Vec has no schema, so keep the source schema in an empty batch. + if data.is_empty() { + data.push(RecordBatch::new_empty(schema)); + } let this = self.clone(); diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 39353f270..50b704532 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -128,6 +128,7 @@ def connect( namespace_client_impl: Optional[str] = None, namespace_client_properties: Optional[Dict[str, str]] = None, namespace_client_pushdown_operations: Optional[List[str]] = None, + merge_insert_source_collect_threshold_bytes: Optional[int] = None, **kwargs: Any, ) -> DBConnection: """Connect to a LanceDB database. @@ -175,6 +176,11 @@ def connect( When true for local/native connections, use directory namespace manifests as the source of truth for table metadata. Existing directory-listed root tables are migrated into the manifest on access. + merge_insert_source_collect_threshold_bytes : int, optional + (For LanceDB OSS only) + Default for ``source_collect_threshold_bytes`` in + [`LanceMergeInsertBuilder.execute`][lancedb.merge.LanceMergeInsertBuilder.execute] + on tables from this connection. Defaults to 64MiB. session: Session, optional (For LanceDB OSS only) A session to use for this connection. Sessions allow you to configure @@ -311,6 +317,9 @@ def connect( session=session, manifest_enabled=manifest_enabled, namespace_client_properties=namespace_client_properties, + merge_insert_source_collect_threshold_bytes=( + merge_insert_source_collect_threshold_bytes + ), ) @@ -429,6 +438,9 @@ def deserialize_conn( storage_options=storage_options, manifest_enabled=parsed.get("manifest_enabled", False), namespace_client_properties=parsed.get("namespace_client_properties"), + merge_insert_source_collect_threshold_bytes=parsed.get( + "merge_insert_source_collect_threshold_bytes" + ), ) elif connection_type == "remote": return RemoteDBConnection( @@ -458,6 +470,7 @@ async def connect_async( manifest_enabled: bool = False, namespace_client_properties: Optional[Dict[str, str]] = None, oauth_config=None, + merge_insert_source_collect_threshold_bytes: Optional[int] = None, ) -> AsyncConnection: """Connect to a LanceDB database. @@ -514,6 +527,11 @@ async def connect_async( OAuth configuration for LanceDB Cloud/Enterprise. This is supported by ``connect_async`` only; synchronous ``connect`` uses API key authentication for ``db://`` URIs. + merge_insert_source_collect_threshold_bytes : int, optional + (For LanceDB OSS only) + Default for ``source_collect_threshold_bytes`` in + [`LanceMergeInsertBuilder.execute`][lancedb.merge.LanceMergeInsertBuilder.execute] + on tables from this connection. Defaults to 64MiB. Examples -------- @@ -571,6 +589,7 @@ async def connect_async( manifest_enabled, namespace_client_properties, oauth_config, + merge_insert_source_collect_threshold_bytes, ) ) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index cfaa8c68a..93be1f8c0 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -619,6 +619,7 @@ async def connect( manifest_enabled: bool = False, namespace_client_properties: Optional[Dict[str, str]] = None, oauth_config: Optional[Any] = None, + merge_insert_source_collect_threshold_bytes: Optional[int] = None, ) -> Connection: ... def connect_namespace( namespace_client_impl: str, diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 47d24b762..eca67ab52 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1283,11 +1283,15 @@ class LanceDBConnection(DBConnection): session: Optional[Session] = None, manifest_enabled: bool = False, namespace_client_properties: Optional[Dict[str, str]] = None, + merge_insert_source_collect_threshold_bytes: Optional[int] = None, _inner: Optional[LanceDbConnection] = None, ): self.storage_options = storage_options self._manifest_enabled = manifest_enabled self._namespace_client_properties = namespace_client_properties + self._merge_insert_source_collect_threshold_bytes = ( + merge_insert_source_collect_threshold_bytes + ) if _inner is not None: self._conn = _inner # Native-derived wrappers resolve this in their async reconstruction @@ -1335,6 +1339,8 @@ class LanceDBConnection(DBConnection): session, manifest_enabled, namespace_client_properties, + None, + merge_insert_source_collect_threshold_bytes, ) # TODO: It would be nice if we didn't store self.storage_options but it is @@ -1395,6 +1401,9 @@ class LanceDBConnection(DBConnection): "storage_options": self.storage_options, "manifest_enabled": self._manifest_enabled, "namespace_client_properties": self._namespace_client_properties, + "merge_insert_source_collect_threshold_bytes": ( + self._merge_insert_source_collect_threshold_bytes + ), "read_consistency_interval_seconds": ( rci.total_seconds() if rci else None ), diff --git a/python/python/lancedb/merge.py b/python/python/lancedb/merge.py index 115a47e2d..6a6372021 100644 --- a/python/python/lancedb/merge.py +++ b/python/python/lancedb/merge.py @@ -39,6 +39,7 @@ class LanceMergeInsertBuilder(object): self._use_index = True self._use_lsm = None self._validate_single_shard = None + self._source_collect_threshold_bytes = None def when_matched_update_all( self, *, where: Optional[str] = None @@ -161,6 +162,7 @@ class LanceMergeInsertBuilder(object): on_bad_vectors: str = "error", fill_value: float = 0.0, timeout: Optional[timedelta] = None, + source_collect_threshold_bytes: Optional[int] = None, ) -> MergeInsertResult: """ Executes the merge insert operation @@ -190,6 +192,20 @@ class LanceMergeInsertBuilder(object): When this is set, the timeout is enforced on all attempts, including the first. + source_collect_threshold_bytes: Optional[int], default None + How many bytes of a one-shot source (such as a + ``pyarrow.RecordBatchReader`` or an iterator) to read into memory + before streaming the rest. + + The merge join picks which side to build its hash table from using + the size each input reports. In-memory data (a table, batch, or + DataFrame) reports its exact row count. A one-shot source reports + nothing, so it is read up to this many bytes: if it ends first it is + merged from memory with an exact size; otherwise the rows read so + far are reported as a lower bound and the rest is streamed. + + Overrides the ``merge_insert_source_collect_threshold_bytes`` + connection option. Defaults to 64MiB. Only applies to local tables. Returns ------- @@ -198,4 +214,6 @@ class LanceMergeInsertBuilder(object): """ if timeout is not None: self._timeout = timeout + if source_collect_threshold_bytes is not None: + self._source_collect_threshold_bytes = source_collect_threshold_bytes return self._table._do_merge(self, new_data, on_bad_vectors, fill_value) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 21838ecef..86efc6dee 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -6536,8 +6536,8 @@ class AsyncTable: fill_value=fill_value, allow_subschema=True, ) - if isinstance(data, pa.Table): - data = pa.RecordBatchReader.from_batches(data.schema, data.to_batches()) + _register_optional_converters() + data = to_scannable(data) return await self._inner.execute_merge_insert( data, dict( @@ -6552,6 +6552,7 @@ class AsyncTable: use_index=merge._use_index, use_lsm=merge._use_lsm, validate_single_shard=merge._validate_single_shard, + source_collect_threshold_bytes=merge._source_collect_threshold_bytes, ), ) diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index a963a0d19..7faf735e9 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -2824,6 +2824,73 @@ def test_merge_insert(mem_db: DBConnection): ) +def _merge_source_reader(ids: List[int]) -> pa.RecordBatchReader: + batches = [ + pa.record_batch({"a": ids[i : i + 2], "b": ["new"] * len(ids[i : i + 2])}) + for i in range(0, len(ids), 2) + ] + return pa.RecordBatchReader.from_batches(batches[0].schema, batches) + + +@pytest.mark.parametrize("threshold", [0, 64 * 1024 * 1024]) +def test_merge_insert_one_shot_source(mem_db: DBConnection, threshold: int): + table = mem_db.create_table( + "my_table", data=pa.table({"a": [1, 2, 3], "b": ["old"] * 3}) + ) + # A threshold of 0 streams the source after its first batch; the default + # collects it into memory. + res = ( + table.merge_insert("a") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute( + _merge_source_reader([2, 3, 4, 5, 6]), + source_collect_threshold_bytes=threshold, + ) + ) + assert res.num_updated_rows == 2 + assert res.num_inserted_rows == 3 + assert table.to_arrow().sort_by("a") == pa.table( + {"a": [1, 2, 3, 4, 5, 6], "b": ["old", "new", "new", "new", "new", "new"]} + ) + + +def test_merge_insert_source_collect_threshold_connection_option(tmp_path): + db = lancedb.connect(tmp_path, merge_insert_source_collect_threshold_bytes=0) + table = db.create_table("my_table", data=pa.table({"a": [1, 2], "b": ["old"] * 2})) + res = ( + table.merge_insert("a") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(_merge_source_reader([2, 3, 4])) + ) + assert res.num_updated_rows == 1 + assert res.num_inserted_rows == 2 + + restored = lancedb.deserialize_conn(db.serialize()) + assert restored._merge_insert_source_collect_threshold_bytes == 0 + + +@pytest.mark.asyncio +async def test_merge_insert_source_collect_threshold_connection_option_async( + tmp_path, +): + db = await lancedb.connect_async( + tmp_path, merge_insert_source_collect_threshold_bytes=0 + ) + table = await db.create_table( + "my_table", data=pa.table({"a": [1, 2], "b": ["old"] * 2}) + ) + res = await ( + table.merge_insert("a") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(_merge_source_reader([2, 3, 4])) + ) + assert res.num_updated_rows == 1 + assert res.num_inserted_rows == 2 + + def test_merge_insert_composite_key(mem_db: DBConnection): table = mem_db.create_table( "my_table", diff --git a/python/src/connection.rs b/python/src/connection.rs index ec0768eb2..778b7bc19 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -1042,7 +1042,7 @@ impl Connection { } #[pyfunction] -#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, sql_host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None))] +#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, sql_host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None, merge_insert_source_collect_threshold_bytes=None))] #[allow(clippy::too_many_arguments)] pub fn connect( py: Python<'_>, @@ -1058,9 +1058,13 @@ pub fn connect( manifest_enabled: bool, namespace_client_properties: Option>, oauth_config: Option, + merge_insert_source_collect_threshold_bytes: Option, ) -> PyResult> { future_into_py(py, async move { let mut builder = lancedb::connect(&uri); + if let Some(bytes) = merge_insert_source_collect_threshold_bytes { + builder = builder.merge_insert_source_collect_threshold_bytes(bytes); + } if let Some(api_key) = api_key { builder = builder.api_key(&api_key); } diff --git a/python/src/table.rs b/python/src/table.rs index d8579f6eb..6a099544a 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -20,7 +20,6 @@ use crate::{ use arrow::{ array::{Array, LargeBinaryArray}, datatypes::{DataType, Schema}, - ffi_stream::ArrowArrayStreamReader, pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; use lancedb::blob::{BlobFile, BlobRangeRequest}; @@ -1566,10 +1565,9 @@ impl Table { pub fn execute_merge_insert<'a>( self_: PyRef<'a, Self>, - data: Bound<'a, PyAny>, + data: PyScannable, parameters: MergeInsertParams, ) -> PyResult> { - let batches: ArrowArrayStreamReader = ArrowArrayStreamReader::from_pyarrow_bound(&data)?; let on = parameters.on.iter().map(|s| s.as_str()).collect::>(); let mut builder = self_.inner_ref()?.merge_insert(&on); if parameters.when_matched_update_all { @@ -1599,12 +1597,12 @@ impl Table { if let Some(validate_single_shard) = parameters.validate_single_shard { builder.validate_single_shard(validate_single_shard); } + if let Some(bytes) = parameters.source_collect_threshold_bytes { + builder.source_collect_threshold_bytes(bytes); + } future_into_py(self_.py(), async move { - let res = builder - .execute(Box::new(batches) as Box) - .await - .infer_error()?; + let res = builder.execute(data).await.infer_error()?; Ok(MergeResult::from(res)) }) } @@ -2045,6 +2043,7 @@ pub struct MergeInsertParams { use_index: Option, use_lsm: Option, validate_single_shard: Option, + source_collect_threshold_bytes: Option, } #[pyclass]