From d2ca0ce0abf89fbf378ba3eef4e1beb4f394925a Mon Sep 17 00:00:00 2001 From: Will Jones Date: Tue, 1 Sep 2026 16:32:04 -0700 Subject: [PATCH] feat: accept multiple `on` columns for merge insert on remote tables (#4102) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge insert has always taken a list of columns to match on, and local tables have always joined on all of them. Remote tables did not: any list longer than one was rejected with `MergeInsertBuilder only supports a single 'on' column`, so a composite-key upsert was impossible against LanceDB Cloud and Enterprise from Rust, Python or TypeScript. The remote request now carries `on` as a list and sends it as one repeated query parameter per column — `?on=shard_key&on=id`. That is how the lance-namespace spec encodes an array-valued `on`, so the server receives a composite key in the shape it expects. A single column still serializes to `?on=id`, exactly what clients sent before, so existing callers are unaffected. A column repeated within `on` is now rejected client-side rather than sent for the server to reject with a 400. No binding changes were needed: `Table.merge_insert` in Python and `Table.mergeInsert` in TypeScript already accepted a list, it just could not reach a remote table. Both gain a test for composite keys, and the doc comments now say what passing several columns means. Part of [ENT-2084](https://linear.app/lancedb/issue/ENT-2084/mergeinsertintotablerequest-support-multiple-columns-for-the). ## Example ```python table.merge_insert(["shard_key", "id"]) \ .when_matched_update_all() \ .when_not_matched_insert_all() \ .execute(new_data) ``` A row whose `id` matches an existing row but whose `shard_key` differs is an insert, not an update. ## Not included Java. Java callers reach merge insert through `org.lance.namespace.LanceNamespace`, whose `MergeInsertIntoTableRequest.on` is a single string until lance-namespace 0.12 ([lance-namespace#363](https://github.com/lance-format/lance-namespace/pull/363), [lance#8915](https://github.com/lance-format/lance/pull/8915)). There is nothing in this repo's Java SDK to change until the `lance-core` pin can move. Sending more than one column requires a server that accepts the repeated parameter ([sophon#7571](https://github.com/lancedb/sophon/pull/7571)); an older server returns a 400 rather than silently merging on one column. Co-authored-by: Claude Opus 5 (1M context) --- docs/src/js/classes/Table.md | 8 ++ nodejs/__test__/table.test.ts | 35 +++++++- nodejs/lancedb/table.ts | 10 +++ python/python/lancedb/table.py | 8 +- python/python/tests/test_table.py | 37 +++++++++ rust/lancedb/src/remote/table.rs | 103 ++++++++++++++++++++++-- rust/lancedb/src/remote/table/insert.rs | 3 +- rust/lancedb/src/table.rs | 4 +- 8 files changed, 196 insertions(+), 12 deletions(-) diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 159348450..894d7a464 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -676,9 +676,17 @@ List all the versions of the table abstract mergeInsert(on): MergeInsertBuilder ``` +Create a [MergeInsertBuilder](MergeInsertBuilder.md), which combines new data with the +existing table in a single transaction — inserting, updating and deleting +rows depending on how they match. + #### Parameters * **on**: `string` \| `string`[] + The column, or columns, to match source rows against target + rows on. Typically a key or id column. Several columns match on the + composite key: a source row updates a target row only when it agrees on + every one of them. #### Returns diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 554c7fcd3..6f80ca74e 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -737,11 +737,12 @@ it("should query documents with LangChain PDF metadata", async () => { describe("merge insert", () => { let tmpDir: tmp.DirResult; + let conn: Connection; let table: Table; beforeEach(async () => { tmpDir = tmp.dirSync({ unsafeCleanup: true }); - const conn = await connect(tmpDir.name); + conn = await connect(tmpDir.name); table = await conn.createTable("some_table", [ { a: 1, b: "a" }, @@ -779,6 +780,38 @@ describe("merge insert", () => { expect(result.map((row) => ({ ...row }))).toEqual(expected); }); + test("upsert on a composite key", async () => { + const composite = await conn.createTable("composite", [ + { shard: "a", id: 1, val: "x" }, + { shard: "a", id: 2, val: "y" }, + { shard: "b", id: 1, val: "z" }, + ]); + + // ("a", 1) matches an existing row and updates it. ("b", 2) agrees with an + // existing row on each key column separately but on neither pair, so it is + // an insert. + const mergeInsertRes = await composite + .mergeInsert(["shard", "id"]) + .whenMatchedUpdateAll() + .whenNotMatchedInsertAll() + .execute([ + { shard: "a", id: 1, val: "X" }, + { shard: "b", id: 2, val: "W" }, + ]); + expect(mergeInsertRes.numUpdatedRows).toBe(1); + expect(mergeInsertRes.numInsertedRows).toBe(1); + + const result = (await composite.toArrow()) + .toArray() + .sort((a, b) => a.shard.localeCompare(b.shard) || a.id - b.id); + + expect(result.map((row) => ({ ...row }))).toEqual([ + { shard: "a", id: 1, val: "X" }, + { shard: "a", id: 2, val: "y" }, + { shard: "b", id: 1, val: "z" }, + { shard: "b", id: 2, val: "W" }, + ]); + }); test("conditional update", async () => { const newData = [ { a: 2, b: "x" }, diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index dc062e337..28591f51c 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -919,6 +919,16 @@ export abstract class Table { /** Return the table as an arrow table */ abstract toArrow(): Promise; + /** + * Create a {@link MergeInsertBuilder}, which combines new data with the + * existing table in a single transaction — inserting, updating and deleting + * rows depending on how they match. + * + * @param on - The column, or columns, to match source rows against target + * rows on. Typically a key or id column. Several columns match on the + * composite key: a source row updates a target row only when it agrees on + * every one of them. + */ abstract mergeInsert(on: string | string[]): MergeInsertBuilder; /** List all the stats of a specified index diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 287dff1f6..e397272fc 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1547,7 +1547,9 @@ class Table(ABC): on: Union[str, Iterable[str]] A column (or columns) to join on. This is how records from the source table and target table are matched. Typically this is some - kind of key or id column. + kind of key or id column. Passing several columns matches on the + composite key: a source row updates a target row only when it + agrees on every one of them. Examples -------- @@ -5701,7 +5703,9 @@ class AsyncTable: on: Union[str, Iterable[str]] A column (or columns) to join on. This is how records from the source table and target table are matched. Typically this is some - kind of key or id column. + kind of key or id column. Passing several columns matches on the + composite key: a source row updates a target row only when it + agrees on every one of them. Examples -------- diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index fbdfac5d8..82ad045c8 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -2682,6 +2682,43 @@ def test_merge_insert(mem_db: DBConnection): ) +def test_merge_insert_composite_key(mem_db: DBConnection): + table = mem_db.create_table( + "my_table", + data=pa.table( + { + "shard": ["a", "a", "b"], + "id": [1, 2, 1], + "val": ["x", "y", "z"], + } + ), + ) + + # ("a", 1) matches an existing row and updates it. ("b", 2) agrees with an + # existing row on each key column separately but on neither pair, so it is + # an insert. + new_data = pa.table({"shard": ["a", "b"], "id": [1, 2], "val": ["X", "W"]}) + res = ( + table.merge_insert(["shard", "id"]) + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(new_data) + ) + assert res.num_updated_rows == 1 + assert res.num_inserted_rows == 1 + + expected = pa.table( + { + "shard": ["a", "a", "b", "b"], + "id": [1, 2, 1, 2], + "val": ["X", "y", "z", "W"], + } + ) + assert table.to_arrow().sort_by([("shard", "ascending"), ("id", "ascending")]) == ( + expected + ) + + def test_merge_insert_nullable_pandas_into_pydantic_schema(mem_db: DBConnection): # Regression test for https://github.com/lancedb/lancedb/issues/2366 pd = pytest.importorskip("pandas") diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index ae5338d65..5faffa8c7 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -72,7 +72,7 @@ use lance_datafusion::exec::{OneShotExec, execute_plan}; use reqwest::{RequestBuilder, Response}; use serde::{Deserialize, Serialize}; use serde_json::Number; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::io::Cursor; use std::pin::Pin; use std::sync::{Arc, Mutex}; @@ -3651,7 +3651,12 @@ impl BaseTable for RemoteTable { #[derive(Serialize, Clone, Debug)] pub struct MergeInsertRequest { - on: String, + // Sent as one repeated `on` query parameter per column, which is how the + // namespace spec encodes an array-valued `on`. serde_urlencoded (which + // reqwest's `query()` uses) cannot serialize a sequence nested in a struct, + // so this field is emitted separately by [`Self::on_query_params`]. + #[serde(skip_serializing)] + on: Vec, when_matched_update_all: bool, when_matched_update_all_filt: Option, when_not_matched_insert_all: bool, @@ -3667,6 +3672,17 @@ pub struct MergeInsertRequest { use_lsm: Option, } +impl MergeInsertRequest { + /// The `on` columns as repeated query parameters: `?on=a&on=b`. + /// + /// A single column serializes to `?on=a`, exactly what clients sent before + /// `on` became a list, so a server that predates composite keys sees no + /// change from a single-column caller. + pub(crate) fn on_query_params(&self) -> Vec<(&str, &str)> { + self.on.iter().map(|col| ("on", col.as_str())).collect() + } +} + fn is_true(b: &bool) -> bool { *b } @@ -3679,12 +3695,15 @@ impl TryFrom for MergeInsertRequest { return Err(Error::InvalidInput { message: "MergeInsertBuilder missing required 'on' field".into(), }); - } else if value.on.len() > 1 { - return Err(Error::NotSupported { - message: "MergeInsertBuilder only supports a single 'on' column".into(), + } + // The server rejects a repeated column with a 400; catching it here + // names the offending column and costs no round trip. + let mut seen = HashSet::with_capacity(value.on.len()); + if let Some(dup) = value.on.iter().find(|col| !seen.insert(*col)) { + return Err(Error::InvalidInput { + message: format!("MergeInsertBuilder 'on' column '{dup}' is repeated"), }); } - let on = value.on[0].clone(); let when_matched_update_all_filt = match value.when_matched_update_all_filt { Some(MergeFilter::Sql(sql)) => Some(sql), @@ -3708,7 +3727,7 @@ impl TryFrom for MergeInsertRequest { }; Ok(Self { - on, + on: value.on, when_matched_update_all: value.when_matched_update_all, when_matched_update_all_filt, when_not_matched_insert_all: value.when_not_matched_insert_all, @@ -4552,6 +4571,76 @@ mod tests { } } + #[tokio::test] + async fn test_merge_insert_composite_key() { + 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 data: Box = Box::new(RecordBatchIterator::new( + [Ok(batch.clone())], + batch.schema(), + )); + + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/merge_insert/"); + + // One repeated `on` per column, in the order the caller gave them. + let on = request + .url() + .query_pairs() + .filter(|(key, _)| key == "on") + .map(|(_, value)| value.into_owned()) + .collect::>(); + assert_eq!(on, vec!["shard_key".to_string(), "id".to_string()]); + + let params = request.url().query_pairs().collect::>(); + assert_eq!(params["when_matched_update_all"], "true"); + assert_eq!(params["when_not_matched_insert_all"], "true"); + + http::Response::builder() + .status(200) + .body(r#"{"version": 43, "num_deleted_rows": 0, "num_inserted_rows": 3, "num_updated_rows": 0}"#) + .unwrap() + }); + + let mut merge = table.merge_insert(&["shard_key", "id"]); + merge.when_matched_update_all(None); + merge.when_not_matched_insert_all(); + let result = table.base_table().merge_insert(merge, data).await.unwrap(); + + assert_eq!(result.num_inserted_rows, 3); + } + + #[tokio::test] + async fn test_merge_insert_rejects_repeated_on_column() { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let data: Box = Box::new(RecordBatchIterator::new( + [Ok(batch.clone())], + batch.schema(), + )); + + let table = Table::new_with_handler::<&str>("my_table", |request| { + panic!("Unexpected request: {}", request.url()); + }); + + let merge = table.merge_insert(&["id", "id"]); + let err = table + .base_table() + .merge_insert(merge, data) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("'id' is repeated")), + "unexpected error: {err}" + ); + } + #[tokio::test] async fn test_merge_insert_retries_on_409() { let batch = RecordBatch::try_new( diff --git a/rust/lancedb/src/remote/table/insert.rs b/rust/lancedb/src/remote/table/insert.rs index 4e0e0d666..eef1d8e42 100644 --- a/rust/lancedb/src/remote/table/insert.rs +++ b/rust/lancedb/src/remote/table/insert.rs @@ -734,6 +734,7 @@ impl ExecutionPlan for RemoteWriteExec { WriteOp::MergeInsert { query, timeout } => { let mut request = client .post(&format!("/v1/table/{}/merge_insert/", identifier)) + .query(&query.on_query_params()) .query(query) .header(CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE); if let Some(timeout) = timeout { @@ -1489,7 +1490,7 @@ mod tests { }); let query = MergeInsertRequest { - on: "id".to_string(), + on: vec!["id".to_string()], when_matched_update_all: false, when_matched_update_all_filt: None, when_not_matched_insert_all: false, diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 4602d35ed..33b6ea8ce 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1506,7 +1506,9 @@ impl Table { /// /// * `on` One or more columns to join on. This is how records from the /// source table and target table are matched. Typically this is some - /// kind of key or id column. + /// kind of key or id column. Several columns match on the composite + /// key: a source row updates a target row only when it agrees on every + /// one of them. /// /// # Examples ///