fix: scope the Function binding guard on schema evolution to the bound columns (#4204)

A table with one registered Function binding refused every add_columns,
alter_columns, drop_columns and field-metadata update, whatever column
they named. The hazard is narrower: a binding stores the exact Arrow
fields of its inputs, outputs and assignment column, so editing one of
those strands it and the table stops accepting rows. Any other column
was never at risk.

The guard now compares the columns a request names, a rename's target
included, against the set every binding depends on, and refuses only on
an intersection. Local and remote tables apply the same rule. The
blanket guard stays on update and merge insert, which cannot say what
they touch.

rust/lancedb/src/remote/table.rs

@@ -3314,9 +3314,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         _read_columns: Option<Vec<String>>,
     ) -> Result<AddColumnsResult> {
         self.check_mutable().await?;
-        crate::table::computed_columns::ensure_no_function_bindings_for_mutation(
+        crate::table::computed_columns::ensure_not_function_bound(
             self.schema().await?.as_ref(),
             "schema evolution",
+            crate::table::schema_evolution::new_column_names(&transforms),
         )?;
         match transforms {
             NewColumnTransform::SqlExpressions(expressions) => {
@@ -3369,9 +3370,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
 
     async fn add_computed_columns(&self, columns: &[(String, String)]) -> Result<AddColumnsResult> {
         self.check_mutable().await?;
-        crate::table::computed_columns::ensure_no_function_bindings_for_mutation(
+        crate::table::computed_columns::ensure_not_function_bound(
             self.schema().await?.as_ref(),
             "schema evolution",
+            columns.iter().map(|(name, _)| name),
         )?;
         // The server plans the declaration against its table schema, including
         // Blob v2 semantics inherited by a direct field projection.
@@ -7982,8 +7984,9 @@ mod tests {
         assert_eq!(result.version, 8);
     }
 
-    #[tokio::test]
-    async fn test_add_function_column_allows_an_existing_binding() {
+    /// The fixture binding's table: `title` and `body` bound as inputs, its
+    /// two outputs declared, plus an unbound `spare`.
+    fn fixture_bound_schema() -> Schema {
         let binding = crate::function::FunctionBinding::from_json(include_str!(
             "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json"
         ))
@@ -8010,13 +8013,78 @@ mod tests {
                 ),
             )
         }));
-        let schema = Schema::new_with_metadata(
+        fields.push(Field::new("spare", DataType::Int32, true));
+        Schema::new_with_metadata(
             fields,
             HashMap::from([(
                 crate::table::computed_columns::FUNCTION_BINDINGS_META_KEY.to_string(),
                 binding_metadata,
             )]),
-        );
+        )
+    }
+
+    /// Only a column the binding uses is refused, and it is refused before
+    /// any request goes out; the rest reach the server as usual.
+    #[tokio::test]
+    async fn test_add_columns_scopes_to_the_columns_a_binding_uses() {
+        let table = Table::new_with_handler("my_table", |request| match request.url().path() {
+            "/v1/table/my_table/describe/" => http::Response::builder()
+                .status(200)
+                .body(describe_response(&fixture_bound_schema()))
+                .unwrap(),
+            "/v1/table/my_table/add_columns/" => http::Response::builder()
+                .status(200)
+                .body(r#"{"version":10}"#.to_string())
+                .unwrap(),
+            path => panic!("Unexpected path: {path}"),
+        });
+        table
+            .add_columns()
+            .computed("doubled", "spare * 2")
+            .execute()
+            .await
+            .unwrap();
+        table
+            .add_columns()
+            .transform(NewColumnTransform::SqlExpressions(vec![(
+                "eager".into(),
+                "spare + 1".into(),
+            )]))
+            .execute()
+            .await
+            .unwrap();
+
+        let table = Table::new_with_handler("my_table", |request| match request.url().path() {
+            "/v1/table/my_table/describe/" => http::Response::builder()
+                .status(200)
+                .body(describe_response(&fixture_bound_schema()))
+                .unwrap(),
+            path => panic!("mutation request must not be sent: {path}"),
+        });
+        for name in ["title", "search_text"] {
+            let err = table
+                .add_columns()
+                .computed(name, "1")
+                .execute()
+                .await
+                .unwrap_err();
+            assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}");
+            let err = table
+                .add_columns()
+                .transform(NewColumnTransform::SqlExpressions(vec![(
+                    name.into(),
+                    "1".into(),
+                )]))
+                .execute()
+                .await
+                .unwrap_err();
+            assert!(matches!(err, Error::InvalidInput { .. }), "{err:?}");
+        }
+    }
+
+    #[tokio::test]
+    async fn test_add_function_column_allows_an_existing_binding() {
+        let schema = fixture_bound_schema();
         let table =
             Table::new_with_handler("my_table", move |request| match request.url().path() {
                 "/v1/table/my_table/describe/" => http::Response::builder()

rust/lancedb/src/table/computed_columns.rs

@@ -391,6 +391,9 @@ pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result
     Ok(())
 }
 
+/// Refuse `operation` outright on a table with a Function binding. For
+/// operations that cannot say which columns they touch; the others use
+/// [`ensure_not_function_bound`].
 pub(crate) fn ensure_no_function_bindings_for_mutation(
     schema: &ArrowSchema,
     operation: &str,
@@ -406,6 +409,49 @@ pub(crate) fn ensure_no_function_bindings_for_mutation(
     Ok(())
 }
 
+/// Refuse `operation` only when a path in `touched` names a column a Function
+/// binding depends on: an input's root, an output, or the assignment column.
+/// A binding stores those columns' exact Arrow fields, so editing one strands it.
+pub(crate) fn ensure_not_function_bound<S: AsRef<str>>(
+    schema: &ArrowSchema,
+    operation: &str,
+    touched: impl IntoIterator<Item = S>,
+) -> Result<()> {
+    ensure_supported_function_metadata(schema)?;
+    let mut protected = BTreeSet::new();
+    for binding in function_bindings(schema)? {
+        for input in binding.inputs() {
+            protected.insert(field_root(&input.field_path)?);
+        }
+        protected.extend(
+            binding
+                .outputs()
+                .iter()
+                .map(|output| output.output_name.clone()),
+        );
+        protected.extend(
+            binding
+                .assignment()
+                .map(|assignment| assignment.output_name.clone()),
+        );
+    }
+    if protected.is_empty() {
+        return Ok(());
+    }
+    for path in touched {
+        let column = field_root(path.as_ref())?;
+        if protected.contains(&column) {
+            return Err(Error::InvalidInput {
+                message: format!(
+                    "{operation} of '{column}' is not supported: a Function binding reads or \
+                     writes it"
+                ),
+            });
+        }
+    }
+    Ok(())
+}
+
 /// Read a field's computed-column declaration, if it carries one.
 ///
 /// A field flagged computed but carrying no kind, or a SQL one missing its
@@ -1330,7 +1376,7 @@ pub(crate) fn ensure_not_written<'a>(
         .map(|declaration| declaration.name)
         .collect();
     for name in written {
-        if declared.iter().any(|declared| declared == root(name)) {
+        if declared.iter().any(|declared| *declared == root(name)) {
             return Err(Error::InvalidInput {
                 message: format!(
                     "column '{}' is computed; its values come from refresh and cannot be \
@@ -1523,9 +1569,24 @@ pub(crate) fn ensure_not_retyped(schema: &ArrowSchema, paths: &[&str]) -> Result
     Ok(())
 }
 
-/// The top-level column a possibly nested input path reads.
-pub(crate) fn root(path: &str) -> &str {
-    path.split('.').next().unwrap_or(path)
+/// The top-level column a path addresses, by the grammar lance resolves it
+/// with, so a quoted spelling names the same column as a bare one.
+pub(crate) fn field_root(path: &str) -> Result<String> {
+    parse_field_path(path)
+        .map_err(|e| Error::InvalidInput {
+            message: format!("invalid column path '{path}': {e}"),
+        })?
+        .into_iter()
+        .next()
+        .ok_or_else(|| Error::InvalidInput {
+            message: format!("column path '{path}' is empty"),
+        })
+}
+
+/// [`field_root`], falling back to the text before the first dot for a
+/// spelling lance would not resolve.
+pub(crate) fn root(path: &str) -> String {
+    field_root(path).unwrap_or_else(|_| path.split('.').next().unwrap_or(path).to_string())
 }
 
 /// A declaration's expression bound to a schema, ready to evaluate.
@@ -1749,7 +1810,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result<
     let mut indices = Vec::with_capacity(inputs.len());
     for input in &inputs {
         let index = runtime_schema
-            .index_of(root(input))
+            .index_of(&root(input))
             .map_err(|_| invalid(format!("unknown column '{input}'")))?;
         if !indices.contains(&index) {
             indices.push(index);
@@ -1884,7 +1945,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result<Ve
 /// assert!(validate_declarations(schema, &[("c".into(), "random()".into())]).is_err());
 /// ```
 pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> {
-    ensure_no_function_bindings_for_mutation(schema.as_ref(), "schema evolution")?;
+    ensure_not_function_bound(
+        schema.as_ref(),
+        "schema evolution",
+        columns.iter().map(|(name, _)| name),
+    )?;
     plan(schema, columns).map(drop)
 }
 
@@ -3197,6 +3262,72 @@ mod tests {
         .unwrap();
     }
 
+    /// The scoped guard refuses exactly the columns a binding uses -- input
+    /// roots, outputs and the assignment column -- and nothing else.
+    #[test]
+    fn test_function_bound_columns_are_the_only_ones_refused() {
+        let mut raw_binding: Value = serde_json::from_str(include_str!(
+            "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json"
+        ))
+        .unwrap();
+        raw_binding["outputs"][0]["nullable"] = Value::Bool(true);
+        raw_binding["outputs"][1]["nullable"] = Value::Bool(true);
+        raw_binding["assignment"] = serde_json::json!({
+            "output_name": "__function_assignment_fb_01K3TEXT",
+            "output_field_id": -1,
+        });
+        raw_binding["output_schema"]["fields"]
+            .as_array_mut()
+            .unwrap()
+            .push(serde_json::json!({
+                "name": "__function_assignment_fb_01K3TEXT",
+                "nullable": true,
+                "type": {"type": "bool"},
+            }));
+        let binding: FunctionBinding = serde_json::from_value(raw_binding).unwrap();
+        let mut fields = valid_function_binding_schema(true, true, &binding)
+            .fields()
+            .to_vec();
+        fields.push(Arc::new(ArrowField::new("spare", DataType::Int32, true)));
+        let schema = ArrowSchema::new_with_metadata(
+            fields,
+            HashMap::from([(
+                FUNCTION_BINDINGS_META_KEY.to_string(),
+                function_bindings_metadata(std::slice::from_ref(&binding)).unwrap(),
+            )]),
+        );
+
+        ensure_not_function_bound(
+            &schema,
+            "schema evolution",
+            ["spare", "spare.nested", "new", "`spare`", "`spare.nested`"],
+        )
+        .unwrap();
+        for path in [
+            "title",
+            "body.nested",
+            "search_text",
+            "search_token_count",
+            "__function_assignment_fb_01K3TEXT",
+            "`title`",
+            "`body`.nested",
+            "`search_text`",
+        ] {
+            let err = ensure_not_function_bound(&schema, "schema evolution", [path]).unwrap_err();
+            assert!(
+                matches!(&err, Error::InvalidInput { message }
+                    if message.contains("a Function binding reads or writes it")),
+                "{path}: {err:?}"
+            );
+        }
+        // A spelling lance cannot resolve is refused rather than compared as text.
+        let err = ensure_not_function_bound(&schema, "schema evolution", ["`title"]).unwrap_err();
+        assert!(
+            matches!(&err, Error::InvalidInput { message } if message.contains("invalid column path")),
+            "{err:?}"
+        );
+    }
+
     #[test]
     fn test_binding_preserves_all_nullable_outputs_with_an_assignment_column() {
         let mut raw_binding: Value = serde_json::from_str(include_str!(

rust/lancedb/src/table/schema_evolution.rs

@@ -103,9 +103,10 @@ pub(crate) async fn execute_add_columns(
     transforms: NewColumnTransform,
     read_columns: Option<Vec<String>>,
 ) -> Result<AddColumnsResult> {
-    computed_columns::ensure_no_function_bindings_for_mutation(
+    computed_columns::ensure_not_function_bound(
         table.schema().await?.as_ref(),
         "schema evolution",
+        new_column_names(&transforms),
     )?;
     // Declarations are admitted only through [`execute_declare`].
     match &transforms {
@@ -131,9 +132,10 @@ pub(crate) async fn execute_declare(
     // An LSM write spec keeps visible rows in tiers refresh cannot reach;
     // checked against latest committed state, not this handle's snapshot.
     table.checkout_latest().await?;
-    computed_columns::ensure_no_function_bindings_for_mutation(
+    computed_columns::ensure_not_function_bound(
         table.schema().await?.as_ref(),
         "schema evolution",
+        columns.iter().map(|(name, _)| name),
     )?;
     // Unset drops the MemWAL index, so the spec alone stops describing a table
     // whose SSTables still hold rows. The shard directories outlive it and are
@@ -156,6 +158,26 @@ pub(crate) async fn execute_declare(
     commit_add_columns(table, transform, None).await
 }
 
+/// The top-level columns `transforms` adds.
+pub(crate) fn new_column_names(transforms: &NewColumnTransform) -> Vec<String> {
+    let names = |schema: &ArrowSchema| {
+        schema
+            .fields()
+            .iter()
+            .map(|field| field.name().clone())
+            .collect::<Vec<_>>()
+    };
+    match transforms {
+        NewColumnTransform::SqlExpressions(expressions) => {
+            expressions.iter().map(|(name, _)| name.clone()).collect()
+        }
+        NewColumnTransform::AllNulls(schema) => names(schema),
+        NewColumnTransform::BatchUDF(udf) => names(&udf.output_schema),
+        NewColumnTransform::Stream(stream) => names(&stream.schema()),
+        NewColumnTransform::Reader(reader) => names(&reader.schema()),
+    }
+}
+
 pub(crate) async fn commit_add_columns(
     table: &NativeTable,
     transforms: NewColumnTransform,
@@ -178,13 +200,18 @@ 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 schema = std::sync::Arc::new(ArrowSchema::from(dataset.schema()));
-    computed_columns::ensure_no_function_bindings_for_mutation(
+    // A Function binding stores its columns' exact fields, nullability
+    // included, so every alteration of one counts, and a rename's target too.
+    computed_columns::ensure_not_function_bound(
         schema.as_ref(),
         "schema evolution",
+        alterations.iter().flat_map(|alteration| {
+            std::iter::once(alteration.path.as_str()).chain(alteration.rename.as_deref())
+        }),
     )?;
+    // 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())
@@ -212,14 +239,9 @@ pub(crate) async fn execute_drop_columns(
 ) -> Result<DropColumnsResult> {
     table.dataset.ensure_mutable()?;
     let mut dataset = (*table.dataset.get().await?).clone();
-    computed_columns::ensure_no_function_bindings_for_mutation(
-        &ArrowSchema::from(dataset.schema()),
-        "schema evolution",
-    )?;
-    computed_columns::ensure_not_an_input(
-        &std::sync::Arc::new(ArrowSchema::from(dataset.schema())),
-        columns,
-    )?;
+    let schema = std::sync::Arc::new(ArrowSchema::from(dataset.schema()));
+    computed_columns::ensure_not_function_bound(schema.as_ref(), "schema evolution", columns)?;
+    computed_columns::ensure_not_an_input(&schema, columns)?;
     dataset.drop_columns(columns).await?;
     let version = dataset.version().version;
     table.dataset.update(dataset);
@@ -241,7 +263,11 @@ pub(crate) async fn execute_update_field_metadata(
     // binding out from under a refresh. A replace on a declared column would
     // silently erase it.
     let schema = ArrowSchema::from(dataset.schema());
-    computed_columns::ensure_no_function_bindings_for_mutation(&schema, "schema evolution")?;
+    computed_columns::ensure_not_function_bound(
+        &schema,
+        "field metadata update",
+        updates.iter().map(|update| update.path.as_str()),
+    )?;
     let declared: Vec<String> = computed_columns::computed_columns(&schema)
         .into_iter()
         .map(|declaration| declaration.name)
@@ -263,7 +289,7 @@ pub(crate) async fn execute_update_field_metadata(
         if update.replace
             && declared
                 .iter()
-                .any(|name| name == computed_columns::root(&update.path))
+                .any(|name| *name == computed_columns::root(&update.path))
         {
             return Err(Error::InvalidInput {
                 message: format!(
@@ -300,8 +326,202 @@ mod tests {
 
     use super::FieldMetadataUpdate;
     use crate::connect;
+    use crate::function::FunctionBinding;
     use crate::query::{ExecutableQuery, QueryBase, Select};
     use crate::table::NewColumnTransform;
+    use crate::table::computed_columns::{
+        FUNCTION_BINDINGS_META_KEY, ensure_supported_function_metadata, function_bindings,
+        function_bindings_metadata, function_computed_column_metadata,
+    };
+    use crate::{Error, Table};
+    use std::collections::HashMap;
+
+    /// A table carrying the fixture binding: `title` and `body` are its
+    /// inputs, `search_text` and `search_token_count` its outputs, `spare`
+    /// nobody's. Stamped the way the server does it, since no local path
+    /// declares a binding.
+    async fn bound_table() -> Table {
+        let conn = connect("memory://").execute().await.unwrap();
+        let batch = record_batch!(
+            ("title", Utf8, ["a"]),
+            ("body", Utf8, ["b"]),
+            ("search_text", Utf8, ["a b"]),
+            ("search_token_count", Int64, [2]),
+            ("spare", Int32, [1])
+        )
+        .unwrap();
+        let table = conn.create_table("bound", batch).execute().await.unwrap();
+        let binding = FunctionBinding::from_json(include_str!(
+            "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json"
+        ))
+        .unwrap();
+        let native = table.as_native().unwrap();
+        let mut dataset = native.dataset.get().await.unwrap().as_ref().clone();
+        dataset
+            .update_schema_metadata(vec![(
+                FUNCTION_BINDINGS_META_KEY.to_string(),
+                Some(function_bindings_metadata(std::slice::from_ref(&binding)).unwrap()),
+            )])
+            .await
+            .unwrap();
+        let inputs = ["title".to_string(), "body".to_string()];
+        let outputs = binding
+            .outputs()
+            .iter()
+            .map(|output| {
+                (
+                    dataset.schema().field(&output.output_name).unwrap().id as u32,
+                    function_computed_column_metadata(
+                        binding.binding_id(),
+                        output.output_ordinal,
+                        &inputs,
+                    ),
+                )
+            })
+            .collect::<Vec<_>>();
+        dataset.replace_field_metadata(outputs).await.unwrap();
+        native.dataset.update(dataset);
+        ensure_supported_function_metadata(&table.schema().await.unwrap()).unwrap();
+        table
+    }
+
+    fn metadata_update(path: &str) -> FieldMetadataUpdate {
+        FieldMetadataUpdate {
+            path: path.into(),
+            metadata: HashMap::from([("unit".to_string(), Some("label".to_string()))]),
+            replace: false,
+        }
+    }
+
+    /// Columns no binding uses evolve as on any table, and the binding is
+    /// still valid afterwards, which is what every later write checks.
+    #[tokio::test]
+    async fn test_schema_evolution_leaves_unbound_columns_free_on_a_bound_table() {
+        let table = bound_table().await;
+        table
+            .add_columns()
+            .transform(NewColumnTransform::SqlExpressions(vec![(
+                "eager".into(),
+                "1".into(),
+            )]))
+            .execute()
+            .await
+            .unwrap();
+        table
+            .add_columns()
+            .computed("derived", "spare * 2")
+            .execute()
+            .await
+            .unwrap();
+        table
+            .update_field_metadata(&[metadata_update("eager")])
+            .await
+            .unwrap();
+        table
+            .alter_columns(&[ColumnAlteration::new("eager".into()).rename("moved".into())])
+            .await
+            .unwrap();
+        table.drop_columns(&["moved"]).await.unwrap();
+
+        let schema = table.schema().await.unwrap();
+        ensure_supported_function_metadata(&schema).unwrap();
+        assert_eq!(function_bindings(&schema).unwrap().len(), 1);
+        assert!(schema.field_with_name("derived").is_ok());
+        assert!(schema.field_with_name("moved").is_err());
+    }
+
+    fn bound(err: Error) {
+        assert!(
+            matches!(&err, Error::InvalidInput { message }
+                if message.contains("a Function binding reads or writes it")),
+            "{err:?}"
+        );
+    }
+
+    /// Every schema-evolution door refuses a column a binding reads or
+    /// writes, including a rename onto one.
+    #[tokio::test]
+    async fn test_schema_evolution_refuses_the_columns_a_function_binding_uses() {
+        let table = bound_table().await;
+        let version = table.version().await.unwrap();
+        for column in ["title", "body", "search_text", "search_token_count"] {
+            bound(table.drop_columns(&[column]).await.unwrap_err());
+            bound(
+                table
+                    .alter_columns(&[ColumnAlteration::new(column.into()).rename("moved".into())])
+                    .await
+                    .unwrap_err(),
+            );
+            bound(
+                table
+                    .alter_columns(&[ColumnAlteration::new(column.into()).set_nullable(false)])
+                    .await
+                    .unwrap_err(),
+            );
+            bound(
+                table
+                    .update_field_metadata(&[metadata_update(column)])
+                    .await
+                    .unwrap_err(),
+            );
+            bound(
+                table
+                    .add_columns()
+                    .transform(NewColumnTransform::SqlExpressions(vec![(
+                        column.into(),
+                        "1".into(),
+                    )]))
+                    .execute()
+                    .await
+                    .unwrap_err(),
+            );
+            bound(
+                table
+                    .add_columns()
+                    .computed(column, "1")
+                    .execute()
+                    .await
+                    .unwrap_err(),
+            );
+        }
+        bound(
+            table
+                .alter_columns(&[ColumnAlteration::new("spare".into()).rename("title".into())])
+                .await
+                .unwrap_err(),
+        );
+        assert_eq!(table.version().await.unwrap(), version);
+    }
+
+    /// Lance resolves a quoted spelling to the same field as the bare one,
+    /// so the guard compares identities, not text.
+    #[tokio::test]
+    async fn quoted_function_output_path_is_still_refused() {
+        let table = bound_table().await;
+        let version = table.version().await.unwrap();
+        for path in ["`title`", "`search_text`", "`title`.nested"] {
+            bound(table.drop_columns(&[path]).await.unwrap_err());
+            bound(
+                table
+                    .alter_columns(&[ColumnAlteration::new(path.into()).set_nullable(false)])
+                    .await
+                    .unwrap_err(),
+            );
+            bound(
+                table
+                    .update_field_metadata(&[metadata_update(path)])
+                    .await
+                    .unwrap_err(),
+            );
+        }
+        bound(
+            table
+                .alter_columns(&[ColumnAlteration::new("spare".into()).rename("`title`".into())])
+                .await
+                .unwrap_err(),
+        );
+        assert_eq!(table.version().await.unwrap(), version);
+    }
 
     // Add Columns Tests