diff --git a/Cargo.lock b/Cargo.lock index 25586e622..995424b20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7583,7 +7583,7 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "itertools 0.14.0", "log", "multimap", @@ -8498,9 +8498,9 @@ dependencies = [ [[package]] name = "rkyv" -version = "0.8.16" +version = "0.8.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73389e0c99e664f919275ab5b5b0471391fe9a8de61e1dff9b1eaf56a90f16e3" +checksum = "815cc8a37159a463064825246cadb07961e25cd9885908606f6d08a98d8f8874" dependencies = [ "bytecheck", "bytes", @@ -8517,9 +8517,9 @@ dependencies = [ [[package]] name = "rkyv_derive" -version = "0.8.16" +version = "0.8.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d2ed0b54125315fb36bd021e82d314d1c126548f871634b483f46b31d13cac6" +checksum = "c0ed1a78a1b19d184b0daa629dd9a024573173ec7d485b287cb369fb3607cc1c" dependencies = [ "proc-macro2", "quote", @@ -9277,7 +9277,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -9289,7 +9289,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40" dependencies = [ - "heck 0.5.0", + "heck 0.4.1", "proc-macro2", "quote", "syn 2.0.117", @@ -9730,7 +9730,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.2", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.61.2", diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 7dd902e19..2ac2fecb2 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -339,7 +339,9 @@ impl Table { let transforms = NewColumnTransform::SqlExpressions(transforms); let res = self .inner_ref()? - .add_columns(transforms, None) + .add_columns() + .transform(transforms) + .execute() .await .default_error()?; Ok(res.into()) @@ -356,7 +358,9 @@ impl Table { let transforms = NewColumnTransform::AllNulls(schema); let res = self .inner_ref()? - .add_columns(transforms, None) + .add_columns() + .transform(transforms) + .execute() .await .default_error()?; Ok(res.into()) diff --git a/python/src/table.rs b/python/src/table.rs index d69609093..f20bcf2bf 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1375,7 +1375,12 @@ impl Table { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { - let result = inner.add_columns(definitions, None).await.infer_error()?; + let result = inner + .add_columns() + .transform(definitions) + .execute() + .await + .infer_error()?; Ok(AddColumnsResult::from(result)) }) } @@ -1389,7 +1394,12 @@ impl Table { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { - let result = inner.add_columns(transform, None).await.infer_error()?; + let result = inner + .add_columns() + .transform(transform) + .execute() + .await + .infer_error()?; Ok(AddColumnsResult::from(result)) }) } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index de3bbc94a..5fefabeb2 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3089,10 +3089,12 @@ mod tests { Box::pin(table.delete("false").map_ok(|_| ())), Box::pin( table - .add_columns( - NewColumnTransform::SqlExpressions(vec![("x".into(), "y".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "x".into(), + "y".into(), + )])) + .execute() .map_ok(|_| ()), ), Box::pin(async { @@ -6388,13 +6390,12 @@ mod tests { }); let result = table - .add_columns( - NewColumnTransform::SqlExpressions(vec![ - ("b".into(), "a + 1".into()), - ("x".into(), "cast(NULL as int32)".into()), - ]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![ + ("b".into(), "a + 1".into()), + ("x".into(), "cast(NULL as int32)".into()), + ])) + .execute() .await .unwrap(); @@ -7119,10 +7120,12 @@ mod tests { } "add_columns" => { let _ = table - .add_columns( - NewColumnTransform::SqlExpressions(vec![("c".into(), "a + 1".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "c".into(), + "a + 1".into(), + )])) + .execute() .await; } "drop_columns" => { @@ -9880,10 +9883,12 @@ mod tests { .await .unwrap(); branch - .add_columns( - NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "b".into(), + "a + 1".into(), + )])) + .execute() .await .unwrap(); branch diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 6f8ecaf14..0d8a8e8b9 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -65,6 +65,7 @@ use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path}; use self::dataset::DatasetConsistencyWrapper; use self::merge::MergeInsertBuilder; +pub mod add_columns; mod add_data; pub mod branch_merge; mod create_index; @@ -79,6 +80,7 @@ pub mod schema_evolution; pub mod update; pub mod write_progress; use crate::index::waiter::wait_for_index; +pub use add_columns::AddColumnsBuilder; #[cfg(feature = "remote")] pub(crate) use add_data::PreprocessingOutput; pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior}; @@ -1620,12 +1622,8 @@ impl Table { } /// Add new columns to the table, providing values to fill in. - pub async fn add_columns( - &self, - transforms: NewColumnTransform, - read_columns: Option>, - ) -> Result { - self.inner.add_columns(transforms, read_columns).await + pub fn add_columns(&self) -> AddColumnsBuilder { + AddColumnsBuilder::new(self.inner.clone()) } /// Change a column's name or nullability. diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs new file mode 100644 index 000000000..0d410cd04 --- /dev/null +++ b/rust/lancedb/src/table/add_columns.rs @@ -0,0 +1,161 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Builder for adding columns to a table. + +use std::sync::Arc; + +use lance::dataset::NewColumnTransform; + +use super::BaseTable; +use super::schema_evolution::AddColumnsResult; +use crate::{Error, Result}; + +/// Adds columns to a table. See [`Table::add_columns`](super::Table::add_columns). +pub struct AddColumnsBuilder { + parent: Arc, + transform: Option, + read_columns: Option>, +} + +impl std::fmt::Debug for AddColumnsBuilder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AddColumnsBuilder") + .field("parent", &self.parent) + .field("has_transform", &self.transform.is_some()) + .field("read_columns", &self.read_columns) + .finish() + } +} + +impl AddColumnsBuilder { + pub(crate) fn new(parent: Arc) -> Self { + Self { + parent, + transform: None, + read_columns: None, + } + } + + /// Set how the new columns' values are produced. Required. + pub fn transform(mut self, transform: NewColumnTransform) -> Self { + self.transform = Some(transform); + self + } + + /// Limit which existing columns a [`NewColumnTransform::BatchUDF`] mapper + /// receives. Every other transform determines what it reads, so setting + /// this alongside one is an error rather than a silent no-op. + pub fn read_columns(mut self, columns: impl IntoIterator>) -> Self { + self.read_columns = Some(columns.into_iter().map(Into::into).collect()); + self + } + + /// Add the columns. + pub async fn execute(self) -> Result { + let Self { + parent, + transform, + read_columns, + } = self; + + let Some(transform) = transform else { + return Err(Error::InvalidInput { + message: "add_columns requires a transform".into(), + }); + }; + + if read_columns.is_some() && !matches!(transform, NewColumnTransform::BatchUDF(_)) { + return Err(Error::InvalidInput { + message: "read_columns applies only to a BatchUDF transform; \ + every other transform determines what it reads" + .into(), + }); + } + + parent.add_columns(transform, read_columns).await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow_array::{Int32Array, RecordBatch, record_batch}; + use arrow_schema::{DataType, Field, Schema}; + use lance::dataset::{BatchUDF, NewColumnTransform}; + + use crate::Table; + use crate::connect; + + async fn table_with_two_columns(name: &str) -> Table { + let conn = connect("memory://").execute().await.unwrap(); + let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [10, 20, 30])).unwrap(); + conn.create_table(name, batch).execute().await.unwrap() + } + + #[tokio::test] + async fn test_requires_a_transform() { + let table = table_with_two_columns("no_transform").await; + let err = table.add_columns().execute().await.unwrap_err(); + assert!( + err.to_string().contains("requires a transform"), + "got: {err}" + ); + } + + #[tokio::test] + async fn test_read_columns_with_sql_expressions_is_rejected() { + let table = table_with_two_columns("read_cols_sql").await; + let err = table + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "doubled".into(), + "x * 2".into(), + )])) + .read_columns(["x"]) + .execute() + .await + .unwrap_err(); + assert!(err.to_string().contains("BatchUDF"), "got: {err}"); + + let schema = table.schema().await.unwrap(); + assert!( + schema.field_with_name("doubled").is_err(), + "a rejected call must not commit" + ); + } + + #[tokio::test] + async fn test_read_columns_limits_what_a_batch_udf_sees() { + let table = table_with_two_columns("read_cols_udf").await; + + let output_schema = Arc::new(Schema::new(vec![Field::new("sum", DataType::Int32, true)])); + let mapper_schema = output_schema.clone(); + let udf = BatchUDF { + mapper: Box::new(move |batch: &RecordBatch| { + assert!(batch.column_by_name("x").is_some()); + assert!(batch.column_by_name("y").is_none(), "y was not requested"); + let x = batch["x"].as_any().downcast_ref::().unwrap(); + let doubled: Int32Array = x.iter().map(|v| v.map(|v| v * 2)).collect(); + Ok(RecordBatch::try_new( + mapper_schema.clone(), + vec![Arc::new(doubled)], + )?) + }), + output_schema, + result_checkpoint: None, + }; + + table + .add_columns() + .transform(NewColumnTransform::BatchUDF(udf)) + .read_columns(["x"]) + .execute() + .await + .unwrap(); + + let schema = table.schema().await.unwrap(); + assert!(schema.field_with_name("sum").is_ok()); + } +} diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index b5b2f5b19..11ba43dd6 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -576,10 +576,12 @@ mod tests { // Add a new physical column AFTER the embedding column. table - .add_columns( - NewColumnTransform::SqlExpressions(vec![("score".into(), "42.0".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "score".into(), + "42.0".into(), + )])) + .execute() .await .unwrap(); @@ -683,7 +685,9 @@ mod tests { true, )])); table - .add_columns(NewColumnTransform::AllNulls(nested_schema), None) + .add_columns() + .transform(NewColumnTransform::AllNulls(nested_schema)) + .execute() .await .unwrap(); diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index 52c8c191f..ce208111a 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -193,10 +193,12 @@ mod tests { // Add a computed column let result = table - .add_columns( - NewColumnTransform::SqlExpressions(vec![("doubled".into(), "id * 2".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "doubled".into(), + "id * 2".into(), + )])) + .execute() .await .unwrap(); @@ -251,13 +253,12 @@ mod tests { // Add multiple columns at once table - .add_columns( - NewColumnTransform::SqlExpressions(vec![ - ("y".into(), "x + 1".into()), - ("z".into(), "x * x".into()), - ]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![ + ("y".into(), "x + 1".into()), + ("z".into(), "x * x".into()), + ])) + .execute() .await .unwrap(); @@ -283,10 +284,12 @@ mod tests { // Add a column with a constant value table - .add_columns( - NewColumnTransform::SqlExpressions(vec![("constant".into(), "42".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "constant".into(), + "42".into(), + )])) + .execute() .await .unwrap(); @@ -659,10 +662,12 @@ mod tests { // Add column increments version let add_result = table - .add_columns( - NewColumnTransform::SqlExpressions(vec![("c".into(), "a + b".into())]), - None, - ) + .add_columns() + .transform(NewColumnTransform::SqlExpressions(vec![( + "c".into(), + "a + b".into(), + )])) + .execute() .await .unwrap(); assert!(add_result.version > v1);