feat: add merge_insert to the node and rust APIs (#915)

node/src/index.ts

@@ -37,6 +37,7 @@ const {
   tableCountRows,
   tableDelete,
   tableUpdate,
+  tableMergeInsert,
   tableCleanupOldVersions,
   tableCompactFiles,
   tableListIndices,
@@ -440,6 +441,38 @@ export interface Table<T = number[]> {
    */
   update: (args: UpdateArgs | UpdateSqlArgs) => Promise<void>
 
+  /**
+   * Runs a "merge insert" operation on the table
+   *
+   * This operation can add rows, update rows, and remove rows all in a single
+   * transaction. It is a very generic tool that can be used to create
+   * behaviors like "insert if not exists", "update or insert (i.e. upsert)",
+   * or even replace a portion of existing data with new data (e.g. replace
+   * all data where month="january")
+   *
+   * The merge insert operation works by combining new data from a
+   * **source table** with existing data in a **target table** by using a
+   * join.  There are three categories of records.
+   *
+   * "Matched" records are records that exist in both the source table and
+   * the target table. "Not matched" records exist only in the source table
+   * (e.g. these are new data) "Not matched by source" records exist only
+   * in the target table (this is old data)
+   *
+   * The MergeInsertArgs can be used to customize what should happen for
+   * each category of data.
+   *
+   * Please note that the data may appear to be reordered as part of this
+   * operation.  This is because updated rows will be deleted from the
+   * dataset and then reinserted at the end with the new values.
+   *
+   * @param on a column to join on.  This is how records from the source
+   *           table and target table are matched.
+   * @param data the new data to insert
+   * @param args parameters controlling how the operation should behave
+   */
+  mergeInsert: (on: string, data: Array<Record<string, unknown>> | ArrowTable, args: MergeInsertArgs) => Promise<void>
+
   /**
    * List the indicies on this table.
    */
@@ -483,6 +516,36 @@ export interface UpdateSqlArgs {
   valuesSql: Record<string, string>
 }
 
+export interface MergeInsertArgs {
+  /**
+   * If true then rows that exist in both the source table (new data) and
+   * the target table (old data) will be updated, replacing the old row
+   * with the corresponding matching row.
+   *
+   * If there are multiple matches then the behavior is undefined.
+   * Currently this causes multiple copies of the row to be created
+   * but that behavior is subject to change.
+   */
+  whenMatchedUpdateAll?: boolean
+  /**
+   * If true then rows that exist only in the source table (new data)
+   * will be inserted into the target table.
+   */
+  whenNotMatchedInsertAll?: boolean
+  /**
+   * If true then rows that exist only in the target table (old data)
+   * will be deleted.
+   *
+   * If this is a string then it will be treated as an SQL filter and
+   * only rows that both do not match any row in the source table and
+   * match the given filter will be deleted.
+   *
+   * This can be used to replace a selection of existing data with
+   * new data.
+   */
+  whenNotMatchedBySourceDelete?: string | boolean
+}
+
 export interface VectorIndex {
   columns: string[]
   name: string
@@ -821,6 +884,38 @@ export class LocalTable<T = number[]> implements Table<T> {
       })
   }
 
+  async mergeInsert (on: string, data: Array<Record<string, unknown>> | ArrowTable, args: MergeInsertArgs): Promise<void> {
+    const whenMatchedUpdateAll = args.whenMatchedUpdateAll ?? false
+    const whenNotMatchedInsertAll = args.whenNotMatchedInsertAll ?? false
+    let whenNotMatchedBySourceDelete = false
+    let whenNotMatchedBySourceDeleteFilt = null
+    if (args.whenNotMatchedBySourceDelete !== undefined && args.whenNotMatchedBySourceDelete !== null) {
+      whenNotMatchedBySourceDelete = true
+      if (args.whenNotMatchedBySourceDelete !== true) {
+        whenNotMatchedBySourceDeleteFilt = args.whenNotMatchedBySourceDelete
+      }
+    }
+
+    const schema = await this.schema
+    let tbl: ArrowTable
+    if (data instanceof ArrowTable) {
+      tbl = data
+    } else {
+      tbl = makeArrowTable(data, { schema })
+    }
+    const buffer = await fromTableToBuffer(tbl, this._embeddings, schema)
+
+    this._tbl = await tableMergeInsert.call(
+      this._tbl,
+      on,
+      whenMatchedUpdateAll,
+      whenNotMatchedInsertAll,
+      whenNotMatchedBySourceDelete,
+      whenNotMatchedBySourceDeleteFilt,
+      buffer
+    )
+  }
+
   /**
    * Clean up old versions of the table, freeing disk space.
    *

node/src/remote/index.ts

@@ -24,7 +24,8 @@ import {
   type IndexStats,
   type UpdateArgs,
   type UpdateSqlArgs,
-  makeArrowTable
+  makeArrowTable,
+  type MergeInsertArgs
 } from '../index'
 import { Query } from '../query'
 
@@ -274,6 +275,52 @@ export class RemoteTable<T = number[]> implements Table<T> {
     throw new Error('Not implemented')
   }
 
+  async mergeInsert (on: string, data: Array<Record<string, unknown>> | ArrowTable, args: MergeInsertArgs): Promise<void> {
+    let tbl: ArrowTable
+    if (data instanceof ArrowTable) {
+      tbl = data
+    } else {
+      tbl = makeArrowTable(data, await this.schema)
+    }
+
+    const queryParams: any = {
+      on
+    }
+    if (args.whenMatchedUpdateAll ?? false) {
+      queryParams.when_matched_update_all = 'true'
+    } else {
+      queryParams.when_matched_update_all = 'false'
+    }
+    if (args.whenNotMatchedInsertAll ?? false) {
+      queryParams.when_not_matched_insert_all = 'true'
+    } else {
+      queryParams.when_not_matched_insert_all = 'false'
+    }
+    if (args.whenNotMatchedBySourceDelete !== false && args.whenNotMatchedBySourceDelete !== null && args.whenNotMatchedBySourceDelete !== undefined) {
+      queryParams.when_not_matched_by_source_delete = 'true'
+      if (typeof args.whenNotMatchedBySourceDelete === 'string') {
+        queryParams.when_not_matched_by_source_delete_filt = args.whenNotMatchedBySourceDelete
+      }
+    } else {
+      queryParams.when_not_matched_by_source_delete = 'false'
+    }
+
+    const buffer = await fromTableToStreamBuffer(tbl, this._embeddings)
+    const res = await this._client.post(
+      `/v1/table/${this._name}/merge_insert/`,
+      buffer,
+      queryParams,
+      'application/vnd.apache.arrow.stream'
+    )
+    if (res.status !== 200) {
+      throw new Error(
+        `Server Error, status: ${res.status}, ` +
+          // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
+          `message: ${res.statusText}: ${res.data}`
+      )
+    }
+  }
+
   async add (data: Array<Record<string, unknown>> | ArrowTable): Promise<number> {
     let tbl: ArrowTable
     if (data instanceof ArrowTable) {

node/src/test/test.ts

@@ -531,6 +531,44 @@ describe('LanceDB client', function () {
       assert.equal(await table.countRows(), 2)
     })
 
+    it('can merge insert records into the table', async function () {
+      const dir = await track().mkdir('lancejs')
+      const con = await lancedb.connect(dir)
+
+      const data = [{ id: 1, age: 1 }, { id: 2, age: 1 }]
+      const table = await con.createTable('my_table', data)
+
+      let newData = [{ id: 2, age: 2 }, { id: 3, age: 2 }]
+      await table.mergeInsert('id', newData, {
+        whenNotMatchedInsertAll: true
+      })
+      assert.equal(await table.countRows(), 3)
+      assert.equal((await table.filter('age = 2').execute()).length, 1)
+
+      newData = [{ id: 3, age: 3 }, { id: 4, age: 3 }]
+      await table.mergeInsert('id', newData, {
+        whenNotMatchedInsertAll: true,
+        whenMatchedUpdateAll: true
+      })
+      assert.equal(await table.countRows(), 4)
+      assert.equal((await table.filter('age = 3').execute()).length, 2)
+
+      newData = [{ id: 5, age: 4 }]
+      await table.mergeInsert('id', newData, {
+        whenNotMatchedInsertAll: true,
+        whenMatchedUpdateAll: true,
+        whenNotMatchedBySourceDelete: 'age < 3'
+      })
+      assert.equal(await table.countRows(), 3)
+
+      await table.mergeInsert('id', newData, {
+        whenNotMatchedInsertAll: true,
+        whenMatchedUpdateAll: true,
+        whenNotMatchedBySourceDelete: true
+      })
+      assert.equal(await table.countRows(), 1)
+    })
+
     it('can update records in the table', async function () {
       const uri = await createTestDB()
       const con = await lancedb.connect(uri)