feat: accept multiple on columns for merge insert on remote tables (#4102)

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) <noreply@anthropic.com>

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
 

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" },

nodejs/lancedb/table.ts

@@ -919,6 +919,16 @@ export abstract class Table {
   /** Return the table as an arrow table */
   abstract toArrow(): Promise<ArrowTable>;
 
+  /**
+   * 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

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
         --------

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")

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<S: HttpSend> BaseTable for RemoteTable<S> {
 
 #[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<String>,
     when_matched_update_all: bool,
     when_matched_update_all_filt: Option<String>,
     when_not_matched_insert_all: bool,
@@ -3667,6 +3672,17 @@ pub struct MergeInsertRequest {
     use_lsm: Option<bool>,
 }
 
+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<MergeInsertBuilder> 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<MergeInsertBuilder> 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<dyn RecordBatchReader + Send> = 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::<Vec<_>>();
+            assert_eq!(on, vec!["shard_key".to_string(), "id".to_string()]);
+
+            let params = request.url().query_pairs().collect::<HashMap<_, _>>();
+            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<dyn RecordBatchReader + Send> = 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(

rust/lancedb/src/remote/table/insert.rs

@@ -734,6 +734,7 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteWriteExec<S> {
                 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,

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
     ///