From 8401e99524422592c92eb04334f5ac4fd7b6925a Mon Sep 17 00:00:00 2001 From: osipovartem Date: Wed, 21 Oct 2026 14:56:44 +0300 Subject: [PATCH] Instantiate stateful scalar UDF invocation per physical expression --- datafusion/expr/src/udf.rs | 24 +++ .../physical-expr/src/scalar_function.rs | 188 +++++++++++++++++- 2 files changed, 204 insertions(+), 8 deletions(-) diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 2de3be4c10fa4..a56c401fa3b48 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -634,6 +634,21 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { None } + /// Create an invocation instance for a newly constructed physical expression. + /// + /// The registered UDF remains responsible for type coercion, return-field + /// inference, and optimizer metadata. The returned UDF is used only for + /// [`Self::invoke_with_args`] and must accept the same arguments and produce + /// values matching the registered return field. Physical-expression rewrites + /// and plan clones share this instance; independent execution requires a + /// newly constructed physical expression. Foreign UDF wrappers do not + /// currently forward this hook across the FFI boundary. + /// + /// The default invokes the registered UDF directly. + fn create_physical_instance(&self) -> Option { + None + } + /// What type will be returned by this function, given the arguments? /// /// By default, this function calls [`Self::return_type`] with the @@ -1150,6 +1165,15 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { None } + fn create_physical_instance(&self) -> Option { + self.inner.create_physical_instance().map(|udf| { + ScalarUDF::new_from_impl(Self { + inner: Arc::clone(udf.inner()).into(), + aliases: self.aliases.clone(), + }) + }) + } + fn aliases(&self) -> &[String] { &self.aliases } diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 6a5ab219aa8dd..b46015ac033e3 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -51,6 +51,7 @@ use datafusion_expr::{ /// Physical expression of a scalar function pub struct ScalarFunctionExpr { fun: Arc, + eval_fun: Arc, name: String, args: Vec>, return_field: FieldRef, @@ -69,6 +70,12 @@ impl Debug for ScalarFunctionExpr { } impl ScalarFunctionExpr { + fn physical_instance(fun: &Arc) -> Arc { + fun.inner() + .create_physical_instance() + .map_or_else(|| Arc::clone(fun), Arc::new) + } + /// Create a new Scalar function pub fn new( name: &str, @@ -77,8 +84,10 @@ impl ScalarFunctionExpr { return_field: FieldRef, config_options: Arc, ) -> Self { + let eval_fun = Self::physical_instance(&fun); Self { fun, + eval_fun, name: name.to_owned(), args, return_field, @@ -111,8 +120,10 @@ impl ScalarFunctionExpr { scalar_arguments: &arguments, }; let return_field = fun.return_field_from_args(ret_args)?; + let eval_fun = Self::physical_instance(&fun); Ok(Self { fun, + eval_fun, name, args, return_field, @@ -187,12 +198,14 @@ impl PartialEq for ScalarFunctionExpr { } let Self { fun, + eval_fun, name, args, return_field, config_options, } = self; fun.eq(&o.fun) + && eval_fun.eq(&o.eval_fun) && name.eq(&o.name) && args.eq(&o.args) && return_field.eq(&o.return_field) @@ -206,12 +219,14 @@ impl Hash for ScalarFunctionExpr { fn hash(&self, state: &mut H) { let Self { fun, + eval_fun, name, args, return_field, config_options: _, // expensive to hash, and often equal } = self; fun.hash(state); + eval_fun.hash(state); name.hash(state); args.hash(state); return_field.hash(state); @@ -252,7 +267,7 @@ impl PhysicalExpr for ScalarFunctionExpr { .all(|arg| matches!(arg, ColumnarValue::Scalar(_))); // evaluate the function - let output = self.fun.invoke_with_args(ScalarFunctionArgs { + let output = self.eval_fun.invoke_with_args(ScalarFunctionArgs { args, arg_fields, number_rows: batch.num_rows(), @@ -292,13 +307,14 @@ impl PhysicalExpr for ScalarFunctionExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(ScalarFunctionExpr::new( - &self.name, - Arc::clone(&self.fun), - children, - Arc::clone(&self.return_field), - Arc::clone(&self.config_options), - ))) + Ok(Arc::new(Self { + fun: Arc::clone(&self.fun), + eval_fun: Arc::clone(&self.eval_fun), + name: self.name.clone(), + args: children, + return_field: Arc::clone(&self.return_field), + config_options: Arc::clone(&self.config_options), + })) } fn evaluate_bounds(&self, children: &[&Interval]) -> Result { @@ -357,9 +373,165 @@ impl PhysicalExpr for ScalarFunctionExpr { mod tests { use super::*; use crate::expressions::Column; + use arrow::array::{Int64Array, RecordBatchOptions}; use arrow::datatypes::Field; use datafusion_expr::{ScalarUDFImpl, Signature}; use datafusion_physical_expr_common::physical_expr::is_volatile; + use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; + + static NEXT_INSTANCE_ID: AtomicU64 = AtomicU64::new(0); + + #[derive(Debug, PartialEq, Eq, Hash)] + struct SequenceFactory { + signature: Signature, + } + + #[derive(Debug)] + struct SequenceInstance { + signature: Signature, + id: u64, + next: AtomicI64, + } + + impl PartialEq for SequenceInstance { + fn eq(&self, other: &Self) -> bool { + self.id == other.id + } + } + + impl Eq for SequenceInstance {} + + impl Hash for SequenceInstance { + fn hash(&self, state: &mut H) { + self.id.hash(state); + } + } + + impl ScalarUDFImpl for SequenceFactory { + fn name(&self) -> &str { + "test_sequence" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + internal_err!("factory must be replaced at physical expression creation") + } + + fn create_physical_instance(&self) -> Option { + Some(ScalarUDF::from(SequenceInstance { + signature: self.signature.clone(), + id: NEXT_INSTANCE_ID.fetch_add(1, Ordering::Relaxed), + next: AtomicI64::new(0), + })) + } + } + + impl ScalarUDFImpl for SequenceInstance { + fn name(&self) -> &str { + "test_sequence" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let count = i64::try_from(args.number_rows).map_err(|_| { + datafusion_common::DataFusionError::Execution("batch too large".into()) + })?; + let start = self.next.fetch_add(count, Ordering::Relaxed); + Ok(ColumnarValue::Array(Arc::new( + Int64Array::from_iter_values(start..start + count), + ))) + } + } + + fn empty_batch(rows: usize) -> Result { + RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(rows)), + ) + .map_err(Into::into) + } + + fn evaluate_sequence(expr: &ScalarFunctionExpr, rows: usize) -> Result> { + let batch = empty_batch(rows)?; + let values = expr.evaluate(&batch)?.into_array(rows)?; + let values = values + .as_any() + .downcast_ref::() + .ok_or_else(|| { + datafusion_common::DataFusionError::Internal("expected Int64".into()) + })?; + Ok(values.values().to_vec()) + } + + #[test] + fn physical_scalar_instances_are_independent_across_constructions() -> Result<()> { + let factory = Arc::new(ScalarUDF::from(SequenceFactory { + signature: Signature::exact(vec![], Volatility::Volatile), + })); + let config = Arc::new(ConfigOptions::default()); + let first = Arc::new(ScalarFunctionExpr::try_new( + Arc::clone(&factory), + vec![], + &Schema::empty(), + Arc::clone(&config), + )?); + assert!(first.fun().inner().is::()); + assert!( + ScalarFunctionExpr::try_downcast_func::(first.as_ref()) + .is_some() + ); + assert_eq!(evaluate_sequence(&first, 3)?, vec![0, 1, 2]); + assert_eq!(evaluate_sequence(&first, 2)?, vec![3, 4]); + + let rewritten = Arc::clone(&first).with_new_children(vec![])?; + let rewritten = + rewritten + .downcast_ref::() + .ok_or_else(|| { + datafusion_common::DataFusionError::Internal("expected UDF".into()) + })?; + assert_eq!(evaluate_sequence(rewritten, 1)?, vec![5]); + + let second = ScalarFunctionExpr::new( + "test_sequence", + Arc::clone(&factory), + vec![], + Arc::new(Field::new("test_sequence", DataType::Int64, false)), + config, + ); + assert_eq!(evaluate_sequence(&second, 2)?, vec![0, 1]); + + let aliased = Arc::new( + ScalarUDF::from(SequenceFactory { + signature: Signature::exact(vec![], Volatility::Volatile), + }) + .with_aliases(["seq_alias"]), + ); + let aliased = ScalarFunctionExpr::try_new( + aliased, + vec![], + &Schema::empty(), + Arc::new(ConfigOptions::default()), + )?; + assert_eq!(aliased.fun().aliases(), &["seq_alias"]); + assert_eq!(evaluate_sequence(&aliased, 2)?, vec![0, 1]); + Ok(()) + } /// Test helper to create a mock UDF with a specific volatility #[derive(Debug, PartialEq, Eq, Hash)]