fix(rust): refuse schema changes that invalidate a computed column

A declaration records the columns its expression reads, but nothing consulted
them: renaming an input left an expression naming a column that no longer
exists, and the failure surfaced at refresh time as a plan error rather than
at the operation that caused it. Dropping or retyping an input did the same.

alter_columns and drop_columns now reject a change to a column some
declaration reads. Nullability is not part of what an expression resolves
against, so it stays allowed. A declaration does not read itself and so
travels with its own binding, and paths compare at their root, since a change
to `metadata.age` invalidates an expression reading `metadata` just as surely.

Binding to field ids instead would leave the expression text naming the old
column, so it would need rewriting stored SQL on every rename. Refusing the
operation is what a generated column does elsewhere.

rust/lancedb/src/table/computed_columns.rs

@@ -140,6 +140,40 @@ pub fn computed_columns(schema: &ArrowSchema) -> Vec<ComputedColumn> {
         .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());
+    }
 }

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<AlterColumnsResult> {
     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::<Vec<_>>();
+    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<DropColumnsResult> {
     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);