feat(table): route merge_insert through the MemWAL LSM write path (#3354)

## Summary

When an `LsmWriteSpec` is installed on a table (#3396), `merge_insert`
upsert
calls are dispatched through Lance's MemWAL `ShardWriter` (LSM-style
append)
instead of the standard merge path.

- **`use_lsm_write`** — a `merge_insert` builder option, default `true`;
set it
  `false` to use the standard path for a call even when a spec is set.
- **`assume_pre_sharded`** — a `merge_insert` builder option, default
`false`;
  skips the per-row shard check and routes by the first row only.
- **`close_lsm_writers`** — drains and closes the table's cached MemWAL
shard
  writers.
- The `merge_insert` **`on`** columns default to, and are validated
against,
  the table's unenforced primary key.
- Shard writers are cached alongside the dataset (in
  `DatasetConsistencyWrapper`) and reused for the session.
- `MergeResult` gains **`num_rows`** — on the LSM path the insert/update
  breakdown is unknown until compaction, so only the total is reported.

Routing covers all three sharding strategies — bucket (murmur3,
Iceberg-compatible), identity, and unsharded. Each `merge_insert` call
targets
a single shard; the whole input is collected and validated before a
single
atomic `ShardWriter::put`, so a validation failure leaves the MemWAL
untouched.

Bindings: Python (`merge_insert(...).use_lsm_write(...)` /
`.assume_pre_sharded(...)`, `Table.close_lsm_writers`) and TypeScript
(`mergeInsert(...).useLsmWrite(...)` / `.assumePreSharded(...)`,
`Table.closeLsmWriters`).

## Context

Reconstructed from the original #3354 branch onto current `main`: the
branch
predated the #3394 (unenforced primary key) / #3396 (`LsmWriteSpec`)
split and
has been rebuilt on that merged foundation. Depends on Lance
`v7.0.0-beta.13`.

The MemWAL read path (reading un-flushed shard data back into queries)
and
remote (LanceDB Cloud) LSM support are follow-ups.

---------

Co-authored-by: Jack Ye <yezhaoqin@gmail.com>

python/python/lancedb/_lancedb.pyi

@@ -220,6 +220,7 @@ class Table:
     async def set_unenforced_primary_key(self, columns: List[str]) -> None: ...
     async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
     async def unset_lsm_write_spec(self) -> None: ...
+    async def close_lsm_writers(self) -> None: ...
     @property
     def tags(self) -> Tags: ...
     def query(self) -> Query: ...
@@ -420,6 +421,7 @@ class MergeResult:
     num_inserted_rows: int
     num_deleted_rows: int
     num_attempts: int
+    num_rows: int
 
 class LsmWriteSpec:
     """Specification selecting Lance's MemWAL LSM-style write path for

python/python/lancedb/merge.py

@@ -34,6 +34,8 @@ class LanceMergeInsertBuilder(object):
         self._when_not_matched_by_source_condition = None
         self._timeout = None
         self._use_index = True
+        self._use_lsm_write = None
+        self._validate_single_shard = None
 
     def when_matched_update_all(
         self, *, where: Optional[str] = None
@@ -96,6 +98,46 @@ class LanceMergeInsertBuilder(object):
         self._use_index = use_index
         return self
 
+    def use_lsm_write(self, use_lsm_write: bool) -> LanceMergeInsertBuilder:
+        """
+        Controls whether the merge uses the MemWAL LSM write path.
+
+        By default (unset), a `merge_insert` on a table with an LSM write spec
+        is routed through Lance's MemWAL shard writer, and a table without one
+        uses the standard path. Pass `False` to force the standard path even
+        when a spec is set. Pass `True` to require a spec — `merge_insert`
+        raises an error if none is installed.
+
+        Parameters
+        ----------
+        use_lsm_write: bool
+            Whether to use the LSM write path.
+        """
+        self._use_lsm_write = use_lsm_write
+        return self
+
+    def validate_single_shard(
+        self, validate_single_shard: bool
+    ) -> LanceMergeInsertBuilder:
+        """
+        Controls how an LSM merge checks that its input targets a single shard.
+
+        When a table has an LSM write spec, every row in a `merge_insert` call
+        must route to the same shard. When `True` (the default), every row is
+        inspected to verify this. When `False`, only the first row is inspected
+        and the shard it routes to is used for the whole input — a faster path
+        for callers that have already pre-sharded their input.
+
+        Has no effect on tables without an LSM write spec.
+
+        Parameters
+        ----------
+        validate_single_shard: bool
+            Whether to check every row routes to one shard. Defaults to `True`.
+        """
+        self._validate_single_shard = validate_single_shard
+        return self
+
     def execute(
         self,
         new_data: DATA,

python/python/lancedb/remote/table.py

@@ -792,6 +792,10 @@ class RemoteTable(Table):
         """Not supported on LanceDB Cloud."""
         return LOOP.run(self._table.unset_lsm_write_spec())
 
+    def close_lsm_writers(self) -> None:
+        """No-op on LanceDB Cloud (no local shard writers)."""
+        return LOOP.run(self._table.close_lsm_writers())
+
     def drop_index(self, index_name: str):
         return LOOP.run(self._table.drop_index(index_name))
 

python/python/lancedb/table.py

@@ -1251,7 +1251,7 @@ class Table(ABC):
         ...      .when_not_matched_insert_all() \\
         ...      .execute(new_data)
         >>> res
-        MergeResult(version=2, num_updated_rows=2, num_inserted_rows=1, num_deleted_rows=0, num_attempts=1)
+        MergeResult(version=2, num_updated_rows=2, num_inserted_rows=1, num_deleted_rows=0, num_attempts=1, num_rows=3)
         >>> # The order of new rows is non-deterministic since we use
         >>> # a hash-join as part of this operation and so we sort here
         >>> table.to_arrow().sort_by("a").to_pandas()
@@ -3601,6 +3601,11 @@ class LanceTable(Table):
         [`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec]."""
         return LOOP.run(self._table.unset_lsm_write_spec())
 
+    def close_lsm_writers(self) -> None:
+        """Close cached MemWAL shard writers. See
+        [`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
+        return LOOP.run(self._table.close_lsm_writers())
+
     def uses_v2_manifest_paths(self) -> bool:
         """
         Check if the table is using the new v2 manifest paths.
@@ -4209,6 +4214,16 @@ class AsyncTable:
         """
         await self._inner.unset_lsm_write_spec()
 
+    async def close_lsm_writers(self) -> None:
+        """Drain and close any cached MemWAL shard writers for this table.
+
+        When an LSM write spec is installed, `merge_insert` opens MemWAL shard
+        writers and caches them for reuse across calls. This closes them,
+        flushing pending data; writers reopen lazily on the next
+        `merge_insert`. It is a no-op when no writers are cached.
+        """
+        await self._inner.close_lsm_writers()
+
     @property
     def name(self) -> str:
         """The name of the table."""
@@ -4659,7 +4674,7 @@ class AsyncTable:
         ...      .when_not_matched_insert_all() \\
         ...      .execute(new_data)
         >>> res
-        MergeResult(version=2, num_updated_rows=2, num_inserted_rows=1, num_deleted_rows=0, num_attempts=1)
+        MergeResult(version=2, num_updated_rows=2, num_inserted_rows=1, num_deleted_rows=0, num_attempts=1, num_rows=3)
         >>> # The order of new rows is non-deterministic since we use
         >>> # a hash-join as part of this operation and so we sort here
         >>> table.to_arrow().sort_by("a").to_pandas()
@@ -5039,6 +5054,8 @@ class AsyncTable:
                 when_not_matched_by_source_condition=merge._when_not_matched_by_source_condition,
                 timeout=merge._timeout,
                 use_index=merge._use_index,
+                use_lsm_write=merge._use_lsm_write,
+                validate_single_shard=merge._validate_single_shard,
             ),
         )
 

python/python/tests/docs/test_merge_insert.py

@@ -57,7 +57,7 @@ async def test_upsert_async(mem_db_async):
     await table.count_rows()  # 3
     res
     # MergeResult(version=2, num_updated_rows=1,
-    # num_inserted_rows=1, num_deleted_rows=0)
+    # num_inserted_rows=1, num_deleted_rows=0, num_rows=2)
     # --8<-- [end:upsert_basic_async]
     assert await table.count_rows() == 3
     assert res.version == 2
@@ -86,7 +86,7 @@ def test_insert_if_not_exists(mem_db):
     table.count_rows()  # 3
     res
     # MergeResult(version=2, num_updated_rows=0,
-    # num_inserted_rows=1, num_deleted_rows=0)
+    # num_inserted_rows=1, num_deleted_rows=0, num_rows=1)
     # --8<-- [end:insert_if_not_exists]
     assert table.count_rows() == 3
     assert res.version == 2
@@ -116,7 +116,7 @@ async def test_insert_if_not_exists_async(mem_db_async):
     await table.count_rows()  # 3
     res
     # MergeResult(version=2, num_updated_rows=0,
-    # num_inserted_rows=1, num_deleted_rows=0)
+    # num_inserted_rows=1, num_deleted_rows=0, num_rows=1)
     # --8<-- [end:insert_if_not_exists]
     assert await table.count_rows() == 3
     assert res.version == 2
@@ -150,7 +150,7 @@ def test_replace_range(mem_db):
     table.count_rows("doc_id = 1")  # 1
     res
     # MergeResult(version=2, num_updated_rows=1,
-    # num_inserted_rows=0, num_deleted_rows=1)
+    # num_inserted_rows=0, num_deleted_rows=1, num_rows=1)
     # --8<-- [end:insert_if_not_exists]
     assert table.count_rows("doc_id = 1") == 1
     assert res.version == 2
@@ -185,7 +185,7 @@ async def test_replace_range_async(mem_db_async):
     await table.count_rows("doc_id = 1")  # 1
     res
     # MergeResult(version=2, num_updated_rows=1,
-    # num_inserted_rows=0, num_deleted_rows=1)
+    # num_inserted_rows=0, num_deleted_rows=1, num_rows=1)
     # --8<-- [end:insert_if_not_exists]
     assert await table.count_rows("doc_id = 1") == 1
     assert res.version == 2

python/python/tests/test_merge_insert_lsm.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+"""Tests for the MemWAL LSM ``merge_insert`` dispatch."""
+
+from datetime import timedelta
+
+import lancedb
+import pyarrow as pa
+import pytest
+from lancedb._lancedb import LsmWriteSpec
+
+SCHEMA = pa.schema(
+    [
+        pa.field("id", pa.int64(), nullable=False),
+        pa.field("value", pa.int64(), nullable=False),
+    ]
+)
+
+REGION_SCHEMA = pa.schema(
+    [
+        pa.field("id", pa.int64(), nullable=False),
+        pa.field("region", pa.utf8(), nullable=False),
+    ]
+)
+
+
+def _reader(ids):
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array(ids, type=pa.int64()),
+            pa.array(list(range(len(ids))), type=pa.int64()),
+        ],
+        schema=SCHEMA,
+    )
+    return pa.RecordBatchReader.from_batches(SCHEMA, [batch])
+
+
+def _region_reader(rows):
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array([row[0] for row in rows], type=pa.int64()),
+            pa.array([row[1] for row in rows], type=pa.utf8()),
+        ],
+        schema=REGION_SCHEMA,
+    )
+    return pa.RecordBatchReader.from_batches(REGION_SCHEMA, [batch])
+
+
+def _bucket_table(tmp_path):
+    """A table with ``id`` as the primary key and a single-bucket LSM spec."""
+    db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
+    table = db.create_table("t", _reader([1, 2, 3]))
+    table.set_unenforced_primary_key("id")
+    # num_buckets = 1: every row routes to the single bucket.
+    table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1))
+    return table
+
+
+def test_lsm_merge_insert_bucket(tmp_path):
+    table = _bucket_table(tmp_path)
+    # Empty `on` defaults to the primary key.
+    result = (
+        table.merge_insert([])
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(_reader([3, 4, 5]))
+    )
+    # LSM path: rows go to the MemWAL, so only num_rows is populated.
+    assert result.num_rows == 3
+    assert result.version == 0
+    assert result.num_inserted_rows == 0
+    assert result.num_updated_rows == 0
+
+
+def test_lsm_merge_insert_unsharded(tmp_path):
+    db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
+    table = db.create_table("t", _reader([1, 2, 3]))
+    table.set_unenforced_primary_key("id")
+    table.set_lsm_write_spec(LsmWriteSpec.unsharded())
+    result = (
+        table.merge_insert("id")
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(_reader([10, 11, 12, 13]))
+    )
+    assert result.num_rows == 4
+
+
+def test_lsm_merge_insert_identity(tmp_path):
+    db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
+    table = db.create_table("t", _region_reader([(1, "us"), (2, "us")]))
+    table.set_unenforced_primary_key("id")
+    table.set_lsm_write_spec(LsmWriteSpec.identity("region"))
+    # All rows share one identity value, so they route to one shard.
+    result = (
+        table.merge_insert([])
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(_region_reader([(3, "us"), (4, "us")]))
+    )
+    assert result.num_rows == 2
+
+
+def test_lsm_merge_insert_use_lsm_write_false(tmp_path):
+    table = _bucket_table(tmp_path)  # rows id = 1, 2, 3
+    # use_lsm_write(False) opts out: the standard path runs and commits.
+    result = (
+        table.merge_insert("id")
+        .when_not_matched_insert_all()
+        .use_lsm_write(False)
+        .execute(_reader([3, 4, 5]))
+    )
+    assert result.num_inserted_rows == 2
+    assert table.count_rows() == 5
+
+
+def test_lsm_merge_insert_validate_single_shard_off(tmp_path):
+    table = _bucket_table(tmp_path)
+    result = (
+        table.merge_insert([])
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .validate_single_shard(False)
+        .execute(_reader([6, 7, 8]))
+    )
+    assert result.num_rows == 3
+
+
+def test_lsm_merge_insert_use_lsm_write_true_requires_spec(tmp_path):
+    # A table with a primary key but no LSM write spec installed.
+    db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(seconds=0))
+    table = db.create_table("t", _reader([1, 2, 3]))
+    table.set_unenforced_primary_key("id")
+    with pytest.raises(Exception, match="use_lsm_write"):
+        (
+            table.merge_insert("id")
+            .when_matched_update_all()
+            .when_not_matched_insert_all()
+            .use_lsm_write(True)
+            .execute(_reader([4]))
+        )
+
+
+def test_lsm_merge_insert_rejects_on_not_primary_key(tmp_path):
+    table = _bucket_table(tmp_path)
+    with pytest.raises(Exception, match="primary key"):
+        (
+            table.merge_insert("value")
+            .when_matched_update_all()
+            .when_not_matched_insert_all()
+            .execute(_reader([1]))
+        )
+
+
+def test_lsm_merge_insert_rejects_non_upsert(tmp_path):
+    table = _bucket_table(tmp_path)
+    # Insert-only (no when_matched_update_all) is not the upsert shape.
+    with pytest.raises(Exception, match="upsert"):
+        table.merge_insert([]).when_not_matched_insert_all().execute(_reader([4]))
+
+
+def test_lsm_close_writers(tmp_path):
+    table = _bucket_table(tmp_path)
+    (
+        table.merge_insert([])
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(_reader([7, 8]))
+    )
+    table.close_lsm_writers()
+    # The writer reopens lazily on the next merge_insert.
+    result = (
+        table.merge_insert([])
+        .when_matched_update_all()
+        .when_not_matched_insert_all()
+        .execute(_reader([9]))
+    )
+    assert result.num_rows == 1
+
+
+@pytest.mark.asyncio
+async def test_async_lsm_merge_insert(tmp_path):
+    db = await lancedb.connect_async(
+        tmp_path, read_consistency_interval=timedelta(seconds=0)
+    )
+    table = await db.create_table("t", _reader([1, 2, 3]))
+    await table.set_unenforced_primary_key("id")
+    await table.set_lsm_write_spec(LsmWriteSpec.bucket("id", 1))
+
+    builder = (
+        table.merge_insert([]).when_matched_update_all().when_not_matched_insert_all()
+    )
+    result = await builder.execute(_reader([3, 4, 5]))
+    assert result.num_rows == 3
+    await table.close_lsm_writers()

python/src/table.rs

@@ -143,18 +143,20 @@ pub struct MergeResult {
     pub num_inserted_rows: u64,
     pub num_deleted_rows: u64,
     pub num_attempts: u32,
+    pub num_rows: u64,
 }
 
 #[pymethods]
 impl MergeResult {
     pub fn __repr__(&self) -> String {
         format!(
-            "MergeResult(version={}, num_updated_rows={}, num_inserted_rows={}, num_deleted_rows={}, num_attempts={})",
+            "MergeResult(version={}, num_updated_rows={}, num_inserted_rows={}, num_deleted_rows={}, num_attempts={}, num_rows={})",
             self.version,
             self.num_updated_rows,
             self.num_inserted_rows,
             self.num_deleted_rows,
-            self.num_attempts
+            self.num_attempts,
+            self.num_rows
         )
     }
 }
@@ -167,6 +169,7 @@ impl From<lancedb::table::MergeResult> for MergeResult {
             num_inserted_rows: result.num_inserted_rows,
             num_deleted_rows: result.num_deleted_rows,
             num_attempts: result.num_attempts,
+            num_rows: result.num_rows,
         }
     }
 }
@@ -194,6 +197,12 @@ impl LsmWriteSpec {
     }
 
     /// Identity sharding — shard by the raw value of `column`.
+    ///
+    /// `column` must be a deterministic function of the unenforced primary
+    /// key: every row with a given primary key must always produce the same
+    /// `column` value, or upserts of that key can land in different shards
+    /// and a stale version can win. Typically `column` is the primary key
+    /// itself or a stable attribute of it.
     #[staticmethod]
     pub fn identity(column: String) -> Self {
         Self {
@@ -933,6 +942,12 @@ impl Table {
         if let Some(use_index) = parameters.use_index {
             builder.use_index(use_index);
         }
+        if let Some(use_lsm_write) = parameters.use_lsm_write {
+            builder.use_lsm_write(use_lsm_write);
+        }
+        if let Some(validate_single_shard) = parameters.validate_single_shard {
+            builder.validate_single_shard(validate_single_shard);
+        }
 
         future_into_py(self_.py(), async move {
             let res = builder.execute(Box::new(batches)).await.infer_error()?;
@@ -971,6 +986,13 @@ impl Table {
         })
     }
 
+    pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
+        let inner = self_.inner_ref()?.clone();
+        future_into_py(self_.py(), async move {
+            inner.close_lsm_writers().await.infer_error()
+        })
+    }
+
     pub fn uses_v2_manifest_paths(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
         let inner = self_.inner_ref()?.clone();
         future_into_py(self_.py(), async move {
@@ -1124,6 +1146,8 @@ pub struct MergeInsertParams {
     when_not_matched_by_source_condition: Option<String>,
     timeout: Option<std::time::Duration>,
     use_index: Option<bool>,
+    use_lsm_write: Option<bool>,
+    validate_single_shard: Option<bool>,
 }
 
 #[pyclass]