feat: {add|alter|drop}_columns APIs (#1015)

Initial work for #959. This exposes the basic functionality for each in
all of the APIs. Will add user guide documentation in a later PR.

nodejs/__test__/table.test.ts

@@ -17,7 +17,7 @@ import * as path from "path";
 import * as fs from "fs";
 
 import { connect } from "../dist";
-import { Schema, Field, Float32, Int32, FixedSizeList } from "apache-arrow";
+import { Schema, Field, Float32, Int32, FixedSizeList, Int64, Float64 } from "apache-arrow";
 import { makeArrowTable } from "../dist/arrow";
 
 describe("Test creating index", () => {
@@ -214,4 +214,69 @@ describe("Read consistency interval", () => {
       expect(await table2.countRows()).toEqual(2n);
     }
   });
+});
+
+
+describe('schema evolution', function () {
+  let tmpDir: string;
+  beforeEach(() => {
+    tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "schema-evolution-"));
+  });
+
+  // Create a new sample table
+  it('can add a new column to the schema', async function () {
+    const con = await connect(tmpDir)
+    const table = await con.createTable('vectors', [
+      { id: 1n, vector: [0.1, 0.2] }
+    ])
+
+    await table.addColumns([{ name: 'price', valueSql: 'cast(10.0 as float)' }])
+
+    const expectedSchema = new Schema([
+      new Field('id', new Int64(), true),
+      new Field('vector', new FixedSizeList(2, new Field('item', new Float32(), true)), true),
+      new Field('price', new Float32(), false)
+    ])
+    expect(await table.schema()).toEqual(expectedSchema)
+  });
+
+  it('can alter the columns in the schema', async function () {
+    const con = await connect(tmpDir)
+    const schema = new Schema([
+      new Field('id', new Int64(), true),
+      new Field('vector', new FixedSizeList(2, new Field('item', new Float32(), true)), true),
+      new Field('price', new Float64(), false)
+    ])
+    const table = await con.createTable('vectors', [
+      { id: 1n, vector: [0.1, 0.2] }
+    ])
+    // Can create a non-nullable column only through addColumns at the moment.
+    await table.addColumns([{ name: 'price', valueSql: 'cast(10.0 as double)' }])
+    expect(await table.schema()).toEqual(schema)
+
+    await table.alterColumns([
+      { path: 'id', rename: 'new_id' },
+      { path: 'price', nullable: true }
+    ])
+
+    const expectedSchema = new Schema([
+      new Field('new_id', new Int64(), true),
+      new Field('vector', new FixedSizeList(2, new Field('item', new Float32(), true)), true),
+      new Field('price', new Float64(), true)
+    ])
+    expect(await table.schema()).toEqual(expectedSchema)
+  });
+
+  it('can drop a column from the schema', async function () {
+    const con = await connect(tmpDir)
+    const table = await con.createTable('vectors', [
+      { id: 1n, vector: [0.1, 0.2] }
+    ])
+    await table.dropColumns(['vector'])
+
+    const expectedSchema = new Schema([
+      new Field('id', new Int64(), true)
+    ])
+    expect(await table.schema()).toEqual(expectedSchema)
+  });
 });

nodejs/lancedb/native.d.ts

@@ -12,6 +12,38 @@ export const enum MetricType {
   Cosine = 1,
   Dot = 2
 }
+/**
+ *  A definition of a column alteration. The alteration changes the column at
+ * `path` to have the new name `name`, to be nullable if `nullable` is true,
+ * and to have the data type `data_type`. At least one of `rename` or `nullable`
+ * must be provided.
+ */
+export interface ColumnAlteration {
+  /**
+   * The path to the column to alter. This is a dot-separated path to the column.
+   * If it is a top-level column then it is just the name of the column. If it is
+   * a nested column then it is the path to the column, e.g. "a.b.c" for a column
+   * `c` nested inside a column `b` nested inside a column `a`.
+   */
+  path: string
+  /**
+   * The new name of the column. If not provided then the name will not be changed.
+   * This must be distinct from the names of all other columns in the table.
+   */
+  rename?: string
+  /** Set the new nullability. Note that a nullable column cannot be made non-nullable. */
+  nullable?: boolean
+}
+/** A definition of a new column to add to a table. */
+export interface AddColumnsSql {
+  /** The name of the new column. */
+  name: string
+  /**
+   * The values to populate the new column with, as a SQL expression.
+   * The expression can reference other columns in the table.
+   */
+  valueSql: string
+}
 export interface ConnectionOptions {
   uri: string
   apiKey?: string
@@ -89,4 +121,7 @@ export class Table {
   delete(predicate: string): Promise<void>
   createIndex(): IndexBuilder
   query(): Query
+  addColumns(transforms: Array<AddColumnsSql>): Promise<void>
+  alterColumns(alterations: Array<ColumnAlteration>): Promise<void>
+  dropColumns(columns: Array<string>): Promise<void>
 }

nodejs/lancedb/table.ts

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 import { Schema, tableFromIPC } from "apache-arrow";
-import { Table as _NativeTable } from "./native";
+import { AddColumnsSql, ColumnAlteration, Table as _NativeTable } from "./native";
 import { toBuffer, Data } from "./arrow";
 import { Query } from "./query";
 import { IndexBuilder } from "./indexer";
@@ -150,4 +150,42 @@ export class Table {
     }
     return q;
   }
+
+  // TODO: Support BatchUDF
+  /**
+   * Add new columns with defined values.
+   *
+   * @param newColumnTransforms pairs of column names and the SQL expression to use
+   *                            to calculate the value of the new column. These
+   *                            expressions will be evaluated for each row in the
+   *                            table, and can reference existing columns in the table.
+   */
+  async addColumns(newColumnTransforms: AddColumnsSql[]): Promise<void> {
+    await this.inner.addColumns(newColumnTransforms);
+  }
+
+  /**
+   * Alter the name or nullability of columns.
+   *
+   * @param columnAlterations One or more alterations to apply to columns.
+   */
+  async alterColumns(columnAlterations: ColumnAlteration[]): Promise<void> {
+    await this.inner.alterColumns(columnAlterations);
+  }
+
+  /**
+   * Drop one or more columns from the dataset
+   *
+   * This is a metadata-only operation and does not remove the data from the
+   * underlying storage. In order to remove the data, you must subsequently
+   * call ``compact_files`` to rewrite the data without the removed columns and
+   * then call ``cleanup_files`` to remove the old files.
+   *
+   * @param columnNames The names of the columns to drop. These can be nested
+   *                    column references (e.g. "a.b.c") or top-level column
+   *                    names (e.g. "a").
+   */
+  async dropColumns(columnNames: string[]): Promise<void> {
+    await this.inner.dropColumns(columnNames);
+  }
 }

nodejs/src/table.rs

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 use arrow_ipc::writer::FileWriter;
-use lancedb::table::AddDataOptions;
-use lancedb::{ipc::ipc_file_to_batches, table::TableRef};
+use lance::dataset::ColumnAlteration as LanceColumnAlteration;
+use lancedb::{
+    ipc::ipc_file_to_batches,
+    table::{AddDataOptions, TableRef},
+};
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
 
@@ -93,4 +96,104 @@ impl Table {
     pub fn query(&self) -> Query {
         Query::new(self)
     }
+
+    #[napi]
+    pub async fn add_columns(&self, transforms: Vec<AddColumnsSql>) -> napi::Result<()> {
+        let transforms = transforms
+            .into_iter()
+            .map(|sql| (sql.name, sql.value_sql))
+            .collect::<Vec<_>>();
+        let transforms = lance::dataset::NewColumnTransform::SqlExpressions(transforms);
+        self.table
+            .add_columns(transforms, None)
+            .await
+            .map_err(|err| {
+                napi::Error::from_reason(format!(
+                    "Failed to add columns to table {}: {}",
+                    self.table, err
+                ))
+            })?;
+        Ok(())
+    }
+
+    #[napi]
+    pub async fn alter_columns(&self, alterations: Vec<ColumnAlteration>) -> napi::Result<()> {
+        for alteration in &alterations {
+            if alteration.rename.is_none() && alteration.nullable.is_none() {
+                return Err(napi::Error::from_reason(
+                    "Alteration must have a 'rename' or 'nullable' field.",
+                ));
+            }
+        }
+        let alterations = alterations
+            .into_iter()
+            .map(LanceColumnAlteration::from)
+            .collect::<Vec<_>>();
+
+        self.table
+            .alter_columns(&alterations)
+            .await
+            .map_err(|err| {
+                napi::Error::from_reason(format!(
+                    "Failed to alter columns in table {}: {}",
+                    self.table, err
+                ))
+            })?;
+        Ok(())
+    }
+
+    #[napi]
+    pub async fn drop_columns(&self, columns: Vec<String>) -> napi::Result<()> {
+        let col_refs = columns.iter().map(String::as_str).collect::<Vec<_>>();
+        self.table.drop_columns(&col_refs).await.map_err(|err| {
+            napi::Error::from_reason(format!(
+                "Failed to drop columns from table {}: {}",
+                self.table, err
+            ))
+        })?;
+        Ok(())
+    }
+}
+
+///  A definition of a column alteration. The alteration changes the column at
+/// `path` to have the new name `name`, to be nullable if `nullable` is true,
+/// and to have the data type `data_type`. At least one of `rename` or `nullable`
+/// must be provided.
+#[napi(object)]
+pub struct ColumnAlteration {
+    /// The path to the column to alter. This is a dot-separated path to the column.
+    /// If it is a top-level column then it is just the name of the column. If it is
+    /// a nested column then it is the path to the column, e.g. "a.b.c" for a column
+    /// `c` nested inside a column `b` nested inside a column `a`.
+    pub path: String,
+    /// The new name of the column. If not provided then the name will not be changed.
+    /// This must be distinct from the names of all other columns in the table.
+    pub rename: Option<String>,
+    /// Set the new nullability. Note that a nullable column cannot be made non-nullable.
+    pub nullable: Option<bool>,
+}
+
+impl From<ColumnAlteration> for LanceColumnAlteration {
+    fn from(js: ColumnAlteration) -> Self {
+        let ColumnAlteration {
+            path,
+            rename,
+            nullable,
+        } = js;
+        Self {
+            path,
+            rename,
+            nullable,
+        }
+    }
+}
+
+/// A definition of a new column to add to a table.
+#[napi(object)]
+pub struct AddColumnsSql {
+    /// The name of the new column.
+    pub name: String,
+    /// The values to populate the new column with, as a SQL expression.
+    /// The expression can reference other columns in the table.
+    pub value_sql: String,
 }