Update in Node & Rust (#696)

Co-authored-by: Will Jones <willjones127@gmail.com>

node/src/index.ts

@@ -21,9 +21,10 @@ import type { EmbeddingFunction } from './embedding/embedding_function'
 import { RemoteConnection } from './remote'
 import { Query } from './query'
 import { isEmbeddingFunction } from './embedding/embedding_function'
+import { type Literal, toSQL } from './util'
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
-const { databaseNew, databaseTableNames, databaseOpenTable, databaseDropTable, tableCreate, tableAdd, tableCreateVectorIndex, tableCountRows, tableDelete, tableCleanupOldVersions, tableCompactFiles, tableListIndices, tableIndexStats } = require('../native.js')
+const { databaseNew, databaseTableNames, databaseOpenTable, databaseDropTable, tableCreate, tableAdd, tableCreateVectorIndex, tableCountRows, tableDelete, tableUpdate, tableCleanupOldVersions, tableCompactFiles, tableListIndices, tableIndexStats } = require('../native.js')
 
 export { Query }
 export type { EmbeddingFunction }
@@ -261,6 +262,39 @@ export interface Table<T = number[]> {
    */
   delete: (filter: string) => Promise<void>
 
+  /**
+   * Update rows in this table.
+   *
+   * This can be used to update a single row, many rows, all rows, or
+   * sometimes no rows (if your predicate matches nothing).
+   *
+   * @param args see {@link UpdateArgs} and {@link UpdateSqlArgs} for more details
+   *
+   * @examples
+   *
+   * ```ts
+   * const con = await lancedb.connect("./.lancedb")
+   * const data = [
+   *    {id: 1, vector: [3, 3], name: 'Ye'},
+   *    {id: 2, vector: [4, 4], name: 'Mike'},
+   * ];
+   * const tbl = await con.createTable("my_table", data)
+   *
+   * await tbl.update({
+   *   filter: "id = 2",
+   *   updates: { vector: [2, 2], name: "Michael" },
+   * })
+   *
+   * let results = await tbl.search([1, 1]).execute();
+   * // Returns [
+   * //   {id: 2, vector: [2, 2], name: 'Michael'}
+   * //   {id: 1, vector: [3, 3], name: 'Ye'}
+   * // ]
+   * ```
+   *
+   */
+  update: (args: UpdateArgs | UpdateSqlArgs) => Promise<void>
+
   /**
    * List the indicies on this table.
    */
@@ -272,6 +306,34 @@ export interface Table<T = number[]> {
   indexStats: (indexUuid: string) => Promise<IndexStats>
 }
 
+export interface UpdateArgs {
+  /**
+   * A filter in the same format used by a sql WHERE clause. The filter may be empty,
+   * in which case all rows will be updated.
+   */
+  where?: string
+
+  /**
+   * A key-value map of updates. The keys are the column names, and the values are the
+   * new values to set
+   */
+  values: Record<string, Literal>
+}
+
+export interface UpdateSqlArgs {
+  /**
+   * A filter in the same format used by a sql WHERE clause. The filter may be empty,
+   * in which case all rows will be updated.
+   */
+  where?: string
+
+  /**
+   * A key-value map of updates. The keys are the column names, and the values are the
+   * new values to set as SQL expressions.
+   */
+  valuesSql: Record<string, string>
+}
+
 export interface VectorIndex {
   columns: string[]
   name: string
@@ -481,6 +543,31 @@ export class LocalTable<T = number[]> implements Table<T> {
     return tableDelete.call(this._tbl, filter).then((newTable: any) => { this._tbl = newTable })
   }
 
+  /**
+   * Update rows in this table.
+   *
+   * @param args see {@link UpdateArgs} and {@link UpdateSqlArgs} for more details
+   *
+   * @returns
+   */
+  async update (args: UpdateArgs | UpdateSqlArgs): Promise<void> {
+    let filter: string | null
+    let updates: Record<string, string>
+
+    if ('valuesSql' in args) {
+      filter = args.where ?? null
+      updates = args.valuesSql
+    } else {
+      filter = args.where ?? null
+      updates = {}
+      for (const [key, value] of Object.entries(args.values)) {
+        updates[key] = toSQL(value)
+      }
+    }
+
+    return tableUpdate.call(this._tbl, filter, updates).then((newTable: any) => { this._tbl = newTable })
+  }
+
   /**
    * Clean up old versions of the table, freeing disk space.
    *

node/src/remote/index.ts

@@ -16,7 +16,8 @@ import {
   type EmbeddingFunction, type Table, type VectorIndexParams, type Connection,
   type ConnectionOptions, type CreateTableOptions, type VectorIndex,
   type WriteOptions,
-  type IndexStats
+  type IndexStats,
+  type UpdateArgs, type UpdateSqlArgs
 } from '../index'
 import { Query } from '../query'
 
@@ -246,6 +247,10 @@ export class RemoteTable<T = number[]> implements Table<T> {
     await this._client.post(`/v1/table/${this._name}/delete/`, { predicate: filter })
   }
 
+  async update (args: UpdateArgs | UpdateSqlArgs): Promise<void> {
+    throw new Error('Not implemented')
+  }
+
   async listIndices (): Promise<VectorIndex[]> {
     const results = await this._client.post(`/v1/table/${this._name}/index/list/`)
     return results.data.indexes?.map((index: any) => ({

node/src/test/test.ts

@@ -260,6 +260,46 @@ describe('LanceDB client', function () {
       assert.equal(await table.countRows(), 2)
     })
 
+    it('can update records in the table', async function () {
+      const uri = await createTestDB()
+      const con = await lancedb.connect(uri)
+
+      const table = await con.openTable('vectors')
+      assert.equal(await table.countRows(), 2)
+
+      await table.update({ where: 'price = 10', valuesSql: { price: '100' } })
+      const results = await table.search([0.1, 0.2]).execute()
+      assert.equal(results[0].price, 100)
+      assert.equal(results[1].price, 11)
+    })
+
+    it('can update the records using a literal value', async function () {
+      const uri = await createTestDB()
+      const con = await lancedb.connect(uri)
+
+      const table = await con.openTable('vectors')
+      assert.equal(await table.countRows(), 2)
+
+      await table.update({ where: 'price = 10', values: { price: 100 } })
+      const results = await table.search([0.1, 0.2]).execute()
+      assert.equal(results[0].price, 100)
+      assert.equal(results[1].price, 11)
+    })
+
+    it('can update every record in the table', async function () {
+      const uri = await createTestDB()
+      const con = await lancedb.connect(uri)
+
+      const table = await con.openTable('vectors')
+      assert.equal(await table.countRows(), 2)
+
+      await table.update({ valuesSql: { price: '100' } })
+      const results = await table.search([0.1, 0.2]).execute()
+
+      assert.equal(results[0].price, 100)
+      assert.equal(results[1].price, 100)
+    })
+
     it('can delete records from a table', async function () {
       const uri = await createTestDB()
       const con = await lancedb.connect(uri)
@@ -542,7 +582,7 @@ describe('Compact and cleanup', function () {
 
     // should have no effect, but this validates the arguments are parsed.
     await table.compactFiles({
-      targetRowsPerFragment: 1024 * 10,
+      targetRowsPerFragment: 102410,
       maxRowsPerGroup: 1024,
       materializeDeletions: true,
       materializeDeletionsThreshold: 0.5,

node/src/test/util.ts

@@ -0,0 +1,45 @@
+// Copyright 2023 LanceDB Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { toSQL } from '../util'
+import * as chai from 'chai'
+
+const expect = chai.expect
+
+describe('toSQL', function () {
+  it('should turn string to SQL expression', function () {
+    expect(toSQL('foo')).to.equal("'foo'")
+  })
+
+  it('should turn number to SQL expression', function () {
+    expect(toSQL(123)).to.equal('123')
+  })
+
+  it('should turn boolean to SQL expression', function () {
+    expect(toSQL(true)).to.equal('TRUE')
+  })
+
+  it('should turn null to SQL expression', function () {
+    expect(toSQL(null)).to.equal('NULL')
+  })
+
+  it('should turn Date to SQL expression', function () {
+    const date = new Date('05 October 2011 14:48 UTC')
+    expect(toSQL(date)).to.equal("'2011-10-05T14:48:00.000Z'")
+  })
+
+  it('should turn array to SQL expression', function () {
+    expect(toSQL(['foo', 'bar', true, 1])).to.equal("['foo', 'bar', TRUE, 1]")
+  })
+})

node/src/util.ts

@@ -0,0 +1,44 @@
+// Copyright 2023 LanceDB Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+export type Literal = string | number | boolean | null | Date | Literal[]
+
+export function toSQL (value: Literal): string {
+  if (typeof value === 'string') {
+    return `'${value}'`
+  }
+
+  if (typeof value === 'number') {
+    return value.toString()
+  }
+
+  if (typeof value === 'boolean') {
+    return value ? 'TRUE' : 'FALSE'
+  }
+
+  if (value === null) {
+    return 'NULL'
+  }
+
+  if (value instanceof Date) {
+    return `'${value.toISOString()}'`
+  }
+
+  if (Array.isArray(value)) {
+    return `[${value.map(toSQL).join(', ')}]`
+  }
+
+  // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
+  throw new Error(`Unsupported value type: ${typeof value} value: (${value})`)
+}