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>

nodejs/__test__/table.test.ts

@@ -2625,3 +2625,97 @@ describe("setLsmWriteSpec / unsetLsmWriteSpec", () => {
     ).rejects.toThrow();
   });
 });
+
+describe("LSM merge insert", () => {
+  let tmpDir: tmp.DirResult;
+
+  beforeEach(() => {
+    tmpDir = tmp.dirSync({ unsafeCleanup: true });
+  });
+  afterEach(() => tmpDir.removeCallback());
+
+  async function bucketTable(conn: Connection): Promise<Table> {
+    // The primary key column must be non-nullable.
+    const table = await conn.createEmptyTable(
+      "t",
+      new arrow.Schema([
+        new arrow.Field("id", new arrow.Utf8(), false),
+        new arrow.Field("value", new arrow.Float64(), true),
+      ]),
+    );
+    await table.add([
+      { id: "a", value: 1 },
+      { id: "b", value: 2 },
+    ]);
+    await table.setUnenforcedPrimaryKey("id");
+    // numBuckets = 1: every row routes to the single bucket.
+    await table.setLsmWriteSpec({
+      specType: "bucket",
+      column: "id",
+      numBuckets: 1,
+    });
+    return table;
+  }
+
+  it("routes merge_insert through the shard writer", async () => {
+    const conn = await connect(tmpDir.name);
+    const table = await bucketTable(conn);
+
+    const res = await table
+      .mergeInsert("id")
+      .whenMatchedUpdateAll()
+      .whenNotMatchedInsertAll()
+      .execute([
+        { id: "c", value: 3 },
+        { id: "d", value: 4 },
+      ]);
+    // LSM path: rows go to the MemWAL, so only numRows is populated.
+    expect(res.numRows).toBe(2);
+    expect(res.version).toBe(0);
+    expect(res.numInsertedRows).toBe(0);
+
+    await table.closeLsmWriters();
+  });
+
+  it("falls back to the standard path with useLsmWrite(false)", async () => {
+    const conn = await connect(tmpDir.name);
+    const table = await bucketTable(conn);
+
+    const res = await table
+      .mergeInsert("id")
+      .whenNotMatchedInsertAll()
+      .useLsmWrite(false)
+      .execute([
+        { id: "b", value: 9 },
+        { id: "e", value: 5 },
+      ]);
+    // Standard path commits: id="e" inserted ("b" already exists).
+    expect(res.numInsertedRows).toBe(1);
+    expect(await table.countRows()).toBe(3);
+  });
+
+  it("supports validateSingleShard(false)", async () => {
+    const conn = await connect(tmpDir.name);
+    const table = await bucketTable(conn);
+
+    const res = await table
+      .mergeInsert("id")
+      .whenMatchedUpdateAll()
+      .whenNotMatchedInsertAll()
+      .validateSingleShard(false)
+      .execute([{ id: "f", value: 6 }]);
+    expect(res.numRows).toBe(1);
+  });
+
+  it("rejects a non-upsert merge under an LSM spec", async () => {
+    const conn = await connect(tmpDir.name);
+    const table = await bucketTable(conn);
+
+    await expect(
+      table
+        .mergeInsert("id")
+        .whenNotMatchedInsertAll()
+        .execute([{ id: "g", value: 7 }]),
+    ).rejects.toThrow();
+  });
+});

nodejs/lancedb/merge.ts

@@ -87,6 +87,41 @@ export class MergeInsertBuilder {
       this.#schema,
     );
   }
+  /**
+   * Controls whether the merge uses the MemWAL LSM write path.
+   *
+   * By default (unset), a `mergeInsert` 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 — `mergeInsert` rejects if none
+   * is installed.
+   *
+   * @param useLsmWrite - Whether to use the LSM write path.
+   */
+  useLsmWrite(useLsmWrite: boolean): MergeInsertBuilder {
+    return new MergeInsertBuilder(
+      this.#native.useLsmWrite(useLsmWrite),
+      this.#schema,
+    );
+  }
+  /**
+   * 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 `mergeInsert` 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.
+   *
+   * @param validateSingleShard - Whether to check every row routes to one shard. Defaults to `true`.
+   */
+  validateSingleShard(validateSingleShard: boolean): MergeInsertBuilder {
+    return new MergeInsertBuilder(
+      this.#native.validateSingleShard(validateSingleShard),
+      this.#schema,
+    );
+  }
   /**
    * Executes the merge insert operation
    *

nodejs/lancedb/table.ts

@@ -161,7 +161,10 @@ export interface Version {
  *
  * `specType` is `"bucket"`, `"identity"`, or `"unsharded"`. For `"bucket"`,
  * `column` and `numBuckets` are required; for `"identity"`, `column` is
- * required.
+ * required and 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).
  */
 export interface LsmWriteSpec {
   /** One of `"bucket"`, `"identity"`, or `"unsharded"`. */
@@ -567,6 +570,16 @@ export abstract class Table {
    * @returns {Promise<void>}
    */
   abstract unsetLsmWriteSpec(): Promise<void>;
+  /**
+   * Drain and close any cached MemWAL shard writers held for this table.
+   *
+   * When an {@link LsmWriteSpec} is installed, `mergeInsert` opens MemWAL
+   * shard writers and caches them for reuse across calls. This closes them,
+   * flushing pending data; writers reopen lazily on the next `mergeInsert`.
+   * It is a no-op when no writers are cached.
+   * @returns {Promise<void>}
+   */
+  abstract closeLsmWriters(): Promise<void>;
   /** Retrieve the version of the table */
 
   abstract version(): Promise<number>;
@@ -1041,6 +1054,10 @@ export class LocalTable extends Table {
     return await this.inner.unsetLsmWriteSpec();
   }
 
+  async closeLsmWriters(): Promise<void> {
+    return await this.inner.closeLsmWriters();
+  }
+
   async version(): Promise<number> {
     return await this.inner.version();
   }

nodejs/src/merge.rs

@@ -50,6 +50,20 @@ impl NativeMergeInsertBuilder {
         this
     }
 
+    #[napi]
+    pub fn use_lsm_write(&self, use_lsm_write: bool) -> Self {
+        let mut this = self.clone();
+        this.inner.use_lsm_write(use_lsm_write);
+        this
+    }
+
+    #[napi]
+    pub fn validate_single_shard(&self, validate_single_shard: bool) -> Self {
+        let mut this = self.clone();
+        this.inner.validate_single_shard(validate_single_shard);
+        this
+    }
+
     #[napi(catch_unwind)]
     pub async fn execute(&self, buf: Buffer) -> napi::Result<MergeResult> {
         let data = ipc_file_to_batches(buf.to_vec())

nodejs/src/table.rs

@@ -391,6 +391,11 @@ impl Table {
             .default_error()
     }
 
+    #[napi(catch_unwind)]
+    pub async fn close_lsm_writers(&self) -> napi::Result<()> {
+        self.inner_ref()?.close_lsm_writers().await.default_error()
+    }
+
     #[napi(catch_unwind)]
     pub async fn version(&self) -> napi::Result<i64> {
         self.inner_ref()?
@@ -940,6 +945,7 @@ pub struct MergeResult {
     pub num_updated_rows: i64,
     pub num_deleted_rows: i64,
     pub num_attempts: i64,
+    pub num_rows: i64,
 }
 
 impl From<lancedb::table::MergeResult> for MergeResult {
@@ -950,6 +956,7 @@ impl From<lancedb::table::MergeResult> for MergeResult {
             num_updated_rows: value.num_updated_rows as i64,
             num_deleted_rows: value.num_deleted_rows as i64,
             num_attempts: value.num_attempts as i64,
+            num_rows: value.num_rows as i64,
         }
     }
 }