diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 2463cc18d..426d3796e 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -140,6 +140,40 @@ pub fn computed_columns(schema: &ArrowSchema) -> Vec { .collect() } +/// Reject a schema change to a column some declaration reads. +/// +/// A binding is SQL text naming its inputs, so renaming, retyping or dropping +/// one leaves an expression that no longer resolves. Refusing the change keeps +/// a declaration that survived [`plan`] evaluable for as long as it exists. +/// +/// Paths are compared at their root: a declaration reading `metadata` is +/// invalidated by a change to `metadata.age` just as surely. +pub(crate) fn ensure_not_an_input(schema: &ArrowSchema, paths: &[&str]) -> Result<()> { + let root = |path: &str| path.split('.').next().unwrap_or(path).to_string(); + for declaration in computed_columns(schema) { + for path in paths { + // A declaration does not read itself, so it is free to be dropped + // or renamed along with its binding. + if declaration.name == root(path) { + continue; + } + if declaration + .inputs + .iter() + .any(|input| root(input) == root(path)) + { + return Err(Error::InvalidInput { + message: format!( + "column '{}' is read by computed column '{}'; drop that column first", + path, declaration.name + ), + }); + } + } + } + Ok(()) +} + /// Resolve `(name, expression)` pairs against `schema` into fields carrying /// their bindings. /// @@ -270,6 +304,7 @@ mod tests { use arrow_array::record_batch; use arrow_schema::DataType; use futures::TryStreamExt; + use lance::dataset::ColumnAlteration; use super::*; use crate::connect; @@ -610,4 +645,61 @@ mod tests { vec!["a".to_string(), "b".to_string()] ); } + + #[tokio::test] + async fn test_dropping_an_input_is_refused() { + let table = table_with_ints("drop_input").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let err = table.drop_columns(&["x"]).await.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("doubled")), + "{err:?}" + ); + } + + #[tokio::test] + async fn test_renaming_an_input_is_refused() { + let table = table_with_ints("rename_input").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + let err = table + .alter_columns(&[ColumnAlteration::new("x".into()).rename("y".into())]) + .await + .unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("doubled")), + "{err:?}" + ); + } + + /// Nothing resolves against nullability, so it is not a rebinding. + #[tokio::test] + async fn test_altering_an_input_nullability_is_allowed() { + let table = table_with_ints("nullable_input").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + table + .alter_columns(&[ColumnAlteration::new("x".into()).set_nullable(true)]) + .await + .unwrap(); + } + + /// A declaration does not read itself, so it travels with its binding. + #[tokio::test] + async fn test_dropping_the_computed_column_is_allowed() { + let table = table_with_ints("drop_computed").await; + add_computed(&table, &[("doubled".into(), "x * 2".into())]) + .await + .unwrap(); + + table.drop_columns(&["doubled"]).await.unwrap(); + assert!(declared(&table).await.is_empty()); + } } diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index ce208111a..d7fc3dccd 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -8,11 +8,13 @@ //! - [`alter_columns`](execute_alter_columns): Rename columns, change types, or modify nullability //! - [`drop_columns`](execute_drop_columns): Remove columns from the table +use arrow_schema::Schema as ArrowSchema; use lance::dataset::{ColumnAlteration, NewColumnTransform}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use super::NativeTable; +use super::computed_columns; use crate::Result; /// The result of an add columns operation. @@ -116,6 +118,14 @@ pub(crate) async fn execute_alter_columns( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + // Nullability is not part of what an expression resolves against, so only + // a rename or a retype can invalidate a binding. + let rebinding = alterations + .iter() + .filter(|alteration| alteration.rename.is_some() || alteration.data_type.is_some()) + .map(|alteration| alteration.path.as_str()) + .collect::>(); + computed_columns::ensure_not_an_input(&ArrowSchema::from(dataset.schema()), &rebinding)?; dataset.alter_columns(alterations).await?; let version = dataset.version().version; table.dataset.update(dataset); @@ -131,6 +141,7 @@ pub(crate) async fn execute_drop_columns( ) -> Result { table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + computed_columns::ensure_not_an_input(&ArrowSchema::from(dataset.schema()), columns)?; dataset.drop_columns(columns).await?; let version = dataset.version().version; table.dataset.update(dataset);