diff --git a/rust/lancedb/src/database/namespace.rs b/rust/lancedb/src/database/namespace.rs index d93f73c64..4e61710ff 100644 --- a/rust/lancedb/src/database/namespace.rs +++ b/rust/lancedb/src/database/namespace.rs @@ -336,6 +336,9 @@ impl Database for LanceNamespaceDatabase { // Refuse a bad declaration before the namespace records a table. crate::table::computed_columns::ensure_declarations_are_planned( &request.data.arrow_schema(), + self.session + .as_deref() + .and_then(lance::session::Session::datafusion_session), )?; let mut table_id = request.namespace_path.clone(); table_id.push(request.name.clone()); diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index c263dc001..bd7fbb942 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -958,7 +958,7 @@ fn is_declaration(key: &str) -> bool { || crate::table::computed_columns::is_declaration_key(key) } -fn without_declarations(field: &ArrowField) -> ArrowField { +pub(crate) fn without_declarations(field: &ArrowField) -> ArrowField { let metadata: HashMap = field .metadata() .iter() @@ -1123,6 +1123,10 @@ impl PreparedDeclaration { /// are validated over the assembled schema, and read only columns the /// view holds (see [`PreparedDeclaration::input_column`]). /// + /// `session` holds the DataFusion functions a declaration may call, as + /// for [`AddColumnsBuilder::computed`](crate::table::AddColumnsBuilder::computed); + /// `None` plans against Lance's built-in functions only. + /// /// ```no_run /// # #![recursion_limit = "256"] /// # use std::collections::HashMap; @@ -1148,7 +1152,7 @@ impl PreparedDeclaration { /// (INPUTS_META_KEY.into(), format!("[\"{text}\"]")), /// ])); /// let view = prepared - /// .with_computed_columns(vec![(1, length)], &[])? + /// .with_computed_columns(vec![(1, length)], &[], None)? /// .create("lengths") /// .await?; /// view.refresh().execute().await?; // rows land with `length` NULL @@ -1160,6 +1164,7 @@ impl PreparedDeclaration { mut self, columns: Vec<(usize, ArrowField)>, bindings: &[FunctionBinding], + session: Option<&datafusion::execution::context::SessionContext>, ) -> Result { let invalid = |message: String| Error::InvalidInput { message }; if columns.is_empty() { @@ -1223,7 +1228,7 @@ impl PreparedDeclaration { } rewrite_column_definitions(&mut metadata, self.schema.as_ref(), &fields)?; let schema = ArrowSchema::new_with_metadata(fields, metadata); - ensure_declarations_are_planned(&schema)?; + ensure_declarations_are_planned(&schema, session)?; self.schema = Arc::new(schema); Ok(self) } @@ -2685,7 +2690,7 @@ mod tests { (INPUTS_META_KEY.into(), "[\"x\"]".into()), ])); let view = prepared - .with_computed_columns(vec![(1, doubled)], &[]) + .with_computed_columns(vec![(1, doubled)], &[], None) .unwrap() .create("m") .await diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index 1072b7684..8851d071e 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/rust/lancedb/src/materialized_view/refresh.rs @@ -133,7 +133,7 @@ pub(crate) async fn execute_refresh( })?; let physical = Arc::new(ArrowSchema::from(view_ds.schema())); - ensure_declarations_are_planned(&physical)?; + ensure_declarations_are_planned(&physical, view_ds.session().datafusion_session())?; if let Some(field) = physical .fields() .iter() diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index cc7b820f9..a58567b2e 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -3407,7 +3407,11 @@ impl NativeTable { namespace_client: Option>, pushdown_operations: HashSet, ) -> Result { - let batches = computed_columns::admit_create_source(batches)?; + let session = params.as_ref().and_then(|p| p.session.as_deref()); + let batches = computed_columns::admit_create_source( + batches, + session.and_then(lance::session::Session::datafusion_session), + )?; // Default params uses format v1. let params = params.unwrap_or(WriteParams { ..Default::default() @@ -3508,7 +3512,12 @@ impl NativeTable { pushdown_operations: HashSet, session: Option>, ) -> Result { - let batches = computed_columns::admit_create_source(batches)?; + let batches = computed_columns::admit_create_source( + batches, + session + .as_deref() + .and_then(lance::session::Session::datafusion_session), + )?; // Build table_id from namespace + name for the storage options provider let mut table_id = namespace.clone(); table_id.push(name.to_string()); diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs index be319a8be..1b827bb37 100644 --- a/rust/lancedb/src/table/add_columns.rs +++ b/rust/lancedb/src/table/add_columns.rs @@ -89,6 +89,30 @@ impl AddColumnsBuilder { /// # Ok(()) /// # } /// ``` + /// + /// In LanceDB Embedded, an expression can also call functions registered + /// on a DataFusion `SessionContext` attached to the connection's Lance + /// session. They are not stored with the table, so attach them on every + /// connection that opens it. + /// + /// ```no_run + /// # use std::sync::Arc; + /// # use datafusion::execution::context::SessionContext; + /// # use datafusion::logical_expr::ScalarUDF; + /// # async fn declare(my_udf: ScalarUDF) -> Result<(), Box> { + /// let functions = SessionContext::new(); + /// functions.register_udf(my_udf); + /// let session = lance::session::Session::default().with_datafusion_session(functions); + /// let db = lancedb::connect("data/db") + /// .session(Arc::new(session)) + /// .execute() + /// .await?; + /// let table = db.open_table("t").execute().await?; + /// table.add_columns().computed("y", "my_udf(x)").execute().await?; + /// table.refresh_column("y").await?; + /// # Ok(()) + /// # } + /// ``` pub fn computed(mut self, name: impl Into, expression: impl Into) -> Self { self.computed.push((name.into(), expression.into())); self diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 0e02e7ce2..e6d7d432c 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -26,6 +26,7 @@ use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef}; +use datafusion::execution::context::SessionContext; use datafusion_common::{ScalarValue, tree_node::TreeNode}; use datafusion_expr::Expr; use datafusion_physical_plan::PhysicalExpr; @@ -40,6 +41,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::function::{FUNCTION_BLOB_V2_TYPE, FunctionApplication, FunctionBinding}; +use crate::materialized_view::without_declarations; use crate::utils::resolve_arrow_field_path; use crate::{Error, Result}; @@ -1640,8 +1642,12 @@ pub(crate) fn plan_function_application( /// /// 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: &SchemaRef, paths: &[&str]) -> Result<()> { - ensure_not_an_input_of(schema, paths, &[]) +pub(crate) fn ensure_not_an_input( + schema: &SchemaRef, + paths: &[&str], + session: Option<&SessionContext>, +) -> Result<()> { + ensure_not_an_input_of(schema, paths, &[], session) } /// [`ensure_not_an_input`] where `removed` names columns whose own @@ -1651,6 +1657,7 @@ pub(crate) fn ensure_not_an_input_of( schema: &SchemaRef, paths: &[&str], removed: &[&str], + session: Option<&SessionContext>, ) -> Result<()> { // By identity, not spelling: a quoted path names the same column, while // a path *into* one removes no declaration at all. @@ -1668,7 +1675,7 @@ pub(crate) fn ensure_not_an_input_of( // The expression, not stored inputs, is the source of truth; an // expression that no longer parses proves nothing, so refuse. let inputs = match &declaration.kind { - ComputedColumnKind::Sql { expression } => Planner::new(schema.clone()) + ComputedColumnKind::Sql { expression } => session_planner(schema.clone(), session) .parse_expr(expression) .map(|parsed| Planner::column_names_in_expr(&parsed)) .map_err(|e| Error::InvalidInput { @@ -1760,7 +1767,10 @@ pub(crate) fn ensure_batch_writes_no_computed_values( /// Function declaration satisfies the binding contract, and no SQL /// declaration reads another computed column. What passes here is what /// `refresh_column` can execute. -pub(crate) fn ensure_declarations_are_planned(schema: &ArrowSchema) -> Result<()> { +pub(crate) fn ensure_declarations_are_planned( + schema: &ArrowSchema, + session: Option<&SessionContext>, +) -> Result<()> { let invalid = |message: String| Error::InvalidInput { message }; // A field with any declaration key is a declaration; a partial one is // not "no declaration", it is a broken one. @@ -1794,7 +1804,12 @@ pub(crate) fn ensure_declarations_are_planned(schema: &ArrowSchema) -> Result<() .filter(|f| f.name() != &column.name) .map(|f| f.as_ref().clone()) .collect(); - let bound = bind(Arc::new(ArrowSchema::new(others)), &column.name, expression)?; + let bound = bind( + Arc::new(ArrowSchema::new(others)), + &column.name, + expression, + session, + )?; if let Some(input) = bound.roots.iter().find(|r| declared.contains(*r)) { return Err(invalid(format!( "computed column '{}' reads computed column '{input}'", @@ -1944,6 +1959,17 @@ pub(crate) struct BoundExpression { pub blob_paths: Vec, /// A directly projected Blob v2 field whose semantics the output inherits. projected_blob_field: Option, + /// Field metadata of the expression's return field, e.g. a function's + /// logical type. + return_metadata: HashMap, +} + +/// Plans against `session`'s functions when given, else Lance's defaults. +fn session_planner(schema: SchemaRef, session: Option<&SessionContext>) -> Planner { + match session { + Some(session) => Planner::new_with_session(schema, session.clone()), + None => Planner::new(schema), + } } fn is_direct_field_projection(expr: &Expr) -> bool { @@ -2089,7 +2115,12 @@ fn referenced_blob_paths(schema: &ArrowSchema, inputs: &[String]) -> Result 0`), /// and the guard protecting the stored SQL has to see every column the text /// names, not just the ones the simplified form still reads. -pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result { +pub(crate) fn bind( + schema: SchemaRef, + column: &str, + expression: &str, + session: Option<&SessionContext>, +) -> Result { let invalid = |message: String| Error::InvalidExpression { column: column.to_string(), message, @@ -2102,7 +2133,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< .into_iter() .collect::>(); let parsing_schema = blob_runtime_schema(schema.as_ref(), &all_blob_paths); - let planner = Planner::new(parsing_schema); + let planner = session_planner(parsing_schema, session); let parsed = planner .parse_expr(expression) .map_err(|e| invalid(e.to_string()))?; @@ -2171,22 +2202,22 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< .map(|field| field.name().clone()) .collect(); - let runtime_planner = Planner::new(runtime_schema); + let runtime_planner = session_planner(runtime_schema, session); let optimized = runtime_planner .optimize_expr(parsed) .map_err(|e| invalid(e.to_string()))?; - let physical = Planner::new(read_schema.clone()) + let physical = session_planner(read_schema.clone(), session) .create_physical_expr(&optimized) .map_err(|e| invalid(e.to_string()))?; - let data_type = physical - .data_type(read_schema.as_ref()) + let return_field = physical + .return_field(read_schema.as_ref()) .map_err(|e| invalid(e.to_string()))?; Ok(BoundExpression { inputs, roots, physical, - data_type, + data_type: return_field.data_type().clone(), blob_paths: blob_paths .iter() .map(|path| { @@ -2195,6 +2226,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< }) .collect(), projected_blob_field, + return_metadata: return_field.metadata().clone(), }) } @@ -2210,7 +2242,11 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< /// batch may declare `a` and then `b = a + 1` in one commit. Refresh order /// then matters, and refresh enforces it: `b` is refused while `a` still has /// unfilled rows. -fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result> { +fn plan_declarations( + schema: SchemaRef, + columns: &[(String, String)], + session: Option<&SessionContext>, +) -> Result> { if columns.is_empty() { return Err(Error::InvalidInput { message: "at least one computed column is required".into(), @@ -2225,23 +2261,22 @@ fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result< return Err(Error::ColumnAlreadyExists { name: name.clone() }); } - let bound = bind(schema.clone(), name, expression)?; + let bound = bind(schema.clone(), name, expression, session)?; // Declared columns start entirely null, so nullability is a property // of the declaration rather than of what the expression yields. let computed_metadata = computed_column_metadata(expression, &bound.inputs); - let field = match bound.projected_blob_field { - Some(source) => { - let mut metadata = source.metadata().clone(); - metadata.retain(|key, _| !is_declaration_key(key)); - metadata.extend(computed_metadata); - source - .with_name(name) - .with_nullable(true) - .with_metadata(metadata) + let projected = match bound.projected_blob_field { + Some(source) => source.with_name(name), + None => { + ArrowField::new(name, bound.data_type, true).with_metadata(bound.return_metadata) } - None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata), }; + // Keep the source's descriptive metadata, not its schema declarations. + let projected = without_declarations(&projected); + let mut metadata = projected.metadata().clone(); + metadata.extend(computed_metadata); + let field = projected.with_nullable(true).with_metadata(metadata); schema = Arc::new(ArrowSchema::new_with_metadata( schema .fields() @@ -2257,8 +2292,12 @@ fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result< Ok(fields) } -pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { - plan_declarations(schema, columns) +pub(crate) fn plan( + schema: SchemaRef, + columns: &[(String, String)], + session: Option<&SessionContext>, +) -> Result> { + plan_declarations(schema, columns, session) } /// Run the schema-level checks of @@ -2267,6 +2306,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result Result Result<()> { +pub fn validate_declarations( + schema: SchemaRef, + columns: &[(String, String)], + session: Option<&SessionContext>, +) -> Result<()> { ensure_not_function_bound( schema.as_ref(), "schema evolution", columns.iter().map(|(name, _)| name), )?; - plan(schema, columns).map(drop) + plan(schema, columns, session).map(drop) } /// Build the transform that declares `columns` against `schema`. @@ -2302,8 +2350,9 @@ pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> pub(crate) fn declare( schema: SchemaRef, columns: &[(String, String)], + session: Option<&SessionContext>, ) -> Result { - let fields = plan_declarations(schema, columns)?; + let fields = plan_declarations(schema, columns, session)?; Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( fields, )))) @@ -2334,9 +2383,10 @@ pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &st /// created. pub(crate) fn admit_create_source( batches: S, + session: Option<&SessionContext>, ) -> Result> { let schema = batches.arrow_schema(); - ensure_declarations_are_planned(&schema)?; + ensure_declarations_are_planned(&schema, session)?; let declared = computed_columns(&schema) .into_iter() .map(|c| c.name) @@ -2496,7 +2546,7 @@ mod tests { )]), )); let declarations = vec![("a".to_string(), "x + 1".to_string())]; - assert!(super::validate_declarations(schema, &declarations).is_err()); + assert!(super::validate_declarations(schema, &declarations, None).is_err()); } #[test] @@ -2588,7 +2638,7 @@ mod tests { } } - use arrow_array::record_batch; + use arrow_array::{RecordBatch, record_batch}; use arrow_schema::{DataType, TimeUnit}; use futures::TryStreamExt; use lance::dataset::ColumnAlteration; @@ -2618,6 +2668,137 @@ mod tests { computed_columns(table.schema().await.unwrap().as_ref()) } + /// Echoes an Int32 and tags its return field, like a function with a + /// logical type. + #[derive(Debug, PartialEq, Eq, Hash)] + struct Tagged(datafusion_expr::Signature); + + impl datafusion_expr::ScalarUDFImpl for Tagged { + fn name(&self) -> &str { + "tagged" + } + + fn signature(&self) -> &datafusion_expr::Signature { + &self.0 + } + + fn return_type(&self, _: &[DataType]) -> datafusion_common::Result { + Ok(DataType::Int32) + } + + fn return_field_from_args( + &self, + _: datafusion_expr::ReturnFieldArgs, + ) -> datafusion_common::Result { + Ok(Arc::new( + ArrowField::new("tagged", DataType::Int32, true) + .with_metadata(HashMap::from([("tag".into(), "yes".into())])), + )) + } + + fn invoke_with_args( + &self, + args: datafusion_expr::ScalarFunctionArgs, + ) -> datafusion_common::Result { + Ok(args.args[0].clone()) + } + } + + #[tokio::test] + #[cfg_attr( + not(feature = "experimental-functions"), + ignore = "requires experimental-functions" + )] + async fn test_declarations_plan_against_the_session_functions() { + let functions = SessionContext::new(); + functions.register_udf(datafusion_expr::ScalarUDF::new_from_impl(Tagged( + datafusion_expr::Signature::exact( + vec![DataType::Int32], + datafusion_expr::Volatility::Immutable, + ), + ))); + let session = + Arc::new(lance::session::Session::default().with_datafusion_session(functions)); + let conn = connect("memory://") + .session(session) + .execute() + .await + .unwrap(); + let batch = record_batch!(("x", Int32, [1, 2, 3]), ("y", Int32, [4, 5, 6])).unwrap(); + let table = conn + .create_table("session_functions", batch) + .execute() + .await + .unwrap(); + + add_computed(&table, &[("t".into(), "tagged(x)".into())]) + .await + .unwrap(); + let schema = table.schema().await.unwrap(); + assert_eq!( + schema.field_with_name("t").unwrap().metadata().get("tag"), + Some(&"yes".to_string()) + ); + table.refresh_column("t").await.unwrap(); + // Dropping a column re-plans every declaration on the table. + table.drop_columns(&["y"]).await.unwrap(); + + let batches = table + .query() + .select(Select::columns(&["t"])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let expected = record_batch!(("t", Int32, [1, 2, 3])).unwrap(); + assert_eq!(batches[0].column(0).as_ref(), expected.column(0).as_ref()); + + let plain = table_with_ints("no_session").await; + let err = add_computed(&plain, &[("t".into(), "tagged(x)".into())]) + .await + .unwrap_err(); + assert!(err.to_string().contains("tagged"), "{err}"); + } + + /// A projection keeps its source's descriptive metadata, not its schema + /// declarations: a copy of the primary key is not a primary key. + #[tokio::test] + #[cfg_attr( + not(feature = "experimental-functions"), + ignore = "requires experimental-functions" + )] + async fn test_projecting_the_primary_key_drops_its_declaration() { + let conn = connect("memory://").execute().await.unwrap(); + let ints = record_batch!(("x", Int32, [1, 2, 3])).unwrap(); + let schema = ArrowSchema::new(vec![ + ArrowField::new("x", DataType::Int32, false) + .with_metadata(HashMap::from([("note".into(), "kept".into())])), + ]); + let batch = RecordBatch::try_new(Arc::new(schema), ints.columns().to_vec()).unwrap(); + let table = conn + .create_table("key_projection", batch) + .execute() + .await + .unwrap(); + table.set_unenforced_primary_key(["x"]).await.unwrap(); + + add_computed(&table, &[("copy".into(), "x".into())]) + .await + .unwrap(); + let schema = table.schema().await.unwrap(); + let copy = schema.field_with_name("copy").unwrap(); + assert!( + !copy + .metadata() + .keys() + .any(|key| key.starts_with("lance-schema:")), + "{copy:?}" + ); + assert_eq!(copy.metadata().get("note"), Some(&"kept".to_string())); + } + #[tokio::test] #[cfg_attr( not(feature = "experimental-functions"), @@ -2658,6 +2839,7 @@ mod tests { ("first".to_string(), "image".to_string()), ("second".to_string(), "first".to_string()), ], + None, ) .unwrap(); @@ -2680,6 +2862,7 @@ mod tests { let fields = plan( schema, &[("payload".to_string(), "coalesce(image, image)".to_string())], + None, ) .unwrap(); @@ -2914,7 +3097,8 @@ mod tests { assert!( validate_declarations( table.schema().await.unwrap(), - &[("e".into(), "random()".into())] + &[("e".into(), "random()".into())], + None, ) .is_err() ); @@ -4803,7 +4987,7 @@ mod tests { let schema = declare_plan(&schema, &downstream, "fb_downstream"); ensure_supported_function_metadata(&schema).unwrap(); - ensure_declarations_are_planned(&schema).unwrap(); + ensure_declarations_are_planned(&schema, None).unwrap(); let bindings = function_bindings(&schema).unwrap(); assert!( computed_function_inputs(&schema, &bindings[0]) @@ -4868,7 +5052,7 @@ mod tests { let base = ArrowSchema::new(vec![ArrowField::new("x", DataType::Int32, true)]); let mut fields = base.fields().to_vec(); fields.extend( - plan(Arc::new(base), &[("doubled".into(), "x * 2".into())]) + plan(Arc::new(base), &[("doubled".into(), "x * 2".into())], None) .unwrap() .into_iter() .map(Arc::new), @@ -4887,7 +5071,7 @@ mod tests { .is_none_or(|metadata| !metadata.keys().any(|key| is_declaration_key(key))) ); let schema = Arc::new(declare_plan(&schema, &downstream, "fb_downstream")); - ensure_declarations_are_planned(&schema).unwrap(); + ensure_declarations_are_planned(&schema, None).unwrap(); let batch = RecordBatch::try_new( schema, vec![ diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index a8ae3c286..3c8ded1cf 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -308,6 +308,7 @@ async fn execute_refresh_column_with_source( schema.clone(), column, &expression, + dataset.session().datafusion_session(), )?); ensure_inputs_filled(&dataset, &schema, column, &bound).await?; let field = dataset @@ -490,7 +491,12 @@ async fn ensure_inputs_filled( ), }); }; - let input_bound = super::computed_columns::bind(schema.clone(), input, expression)?; + let input_bound = super::computed_columns::bind( + schema.clone(), + input, + expression, + dataset.session().datafusion_session(), + )?; let mut unfilled = 0u64; for fragment in dataset.get_fragments() { unfilled += count_fragment_gains(dataset, &fragment, &input_bound, input) diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index db4d6b91f..ab872c7fc 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -166,7 +166,9 @@ pub(crate) async fn execute_declare( .into(), }); } - let transform = computed_columns::declare(table.schema().await?, columns)?; + let session = table.dataset.get().await?.session(); + let transform = + computed_columns::declare(table.schema().await?, columns, session.datafusion_session())?; let NewColumnTransform::AllNulls(schema) = transform else { unreachable!("declarations are all-null") }; @@ -307,7 +309,11 @@ pub(crate) async fn execute_alter_columns( .filter(|alteration| alteration.rename.is_some() || alteration.data_type.is_some()) .map(|alteration| alteration.path.as_str()) .collect::>(); - computed_columns::ensure_not_an_input(&schema, &rebinding)?; + computed_columns::ensure_not_an_input( + &schema, + &rebinding, + dataset.session().datafusion_session(), + )?; let retyped = alterations .iter() .filter(|alteration| alteration.data_type.is_some()) @@ -344,7 +350,12 @@ pub(crate) async fn execute_drop_columns( let dropped = names.iter().map(String::as_str).collect::>(); ensure_not_embedding_column(&schema, "drop", dropped.iter().copied())?; computed_columns::ensure_not_bound_by(&unbinding.retained, "schema evolution", &dropped)?; - computed_columns::ensure_not_an_input_of(&schema, &dropped, &dropped)?; + computed_columns::ensure_not_an_input_of( + &schema, + &dropped, + &dropped, + dataset.session().datafusion_session(), + )?; if !unbinding.is_noop() { // The retirement commits before the projection, so the whole