Remove ScalarFunctionDefinition (#10325)

* Remove ScalarFunctionDefinition * Fix test * rename func_def to func --------- Co-authored-by: Andrew Lamb <[email protected]>
apache · May 7, 2024 · 742e3c5 · 742e3c5
1 parent 5146f44
commit 742e3c5
Show file tree

Hide file tree

Showing 22 changed files with 124 additions and 231 deletions.
diff --git a/datafusion/core/src/datasource/listing/helpers.rs b/datafusion/core/src/datasource/listing/helpers.rs
@@ -38,7 +38,7 @@ use log::{debug, trace};
 
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{Column, DFSchema, DataFusionError};
-use datafusion_expr::{Expr, ScalarFunctionDefinition, Volatility};
+use datafusion_expr::{Expr, Volatility};
 use datafusion_physical_expr::create_physical_expr;
 use object_store::path::Path;
 use object_store::{ObjectMeta, ObjectStore};
@@ -89,16 +89,12 @@ pub fn expr_applicable_for_cols(col_names: &[String], expr: &Expr) -> bool {
             | Expr::Case { .. } => Ok(TreeNodeRecursion::Continue),
 
             Expr::ScalarFunction(scalar_function) => {
-                match &scalar_function.func_def {
-                    ScalarFunctionDefinition::UDF(fun) => {
-                        match fun.signature().volatility {
-                            Volatility::Immutable => Ok(TreeNodeRecursion::Continue),
-                            // TODO: Stable functions could be `applicable`, but that would require access to the context
-                            Volatility::Stable | Volatility::Volatile => {
-                                is_applicable = false;
-                                Ok(TreeNodeRecursion::Stop)
-                            }
-                        }
+                match scalar_function.func.signature().volatility {
+                    Volatility::Immutable => Ok(TreeNodeRecursion::Continue),
+                    // TODO: Stable functions could be `applicable`, but that would require access to the context
+                    Volatility::Stable | Volatility::Volatile => {
+                        is_applicable = false;
+                        Ok(TreeNodeRecursion::Stop)
                     }
                 }
             }

diff --git a/datafusion/core/src/physical_optimizer/projection_pushdown.rs b/datafusion/core/src/physical_optimizer/projection_pushdown.rs
@@ -1301,8 +1301,7 @@ mod tests {
     use datafusion_execution::object_store::ObjectStoreUrl;
     use datafusion_execution::{SendableRecordBatchStream, TaskContext};
     use datafusion_expr::{
-        ColumnarValue, Operator, ScalarFunctionDefinition, ScalarUDF, ScalarUDFImpl,
-        Signature, Volatility,
+        ColumnarValue, Operator, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
     };
     use datafusion_physical_expr::expressions::{
         BinaryExpr, CaseExpr, CastExpr, NegativeExpr,
@@ -1363,9 +1362,7 @@ mod tests {
             Arc::new(NegativeExpr::new(Arc::new(Column::new("f", 4)))),
             Arc::new(ScalarFunctionExpr::new(
                 "scalar_expr",
-                ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
-                    DummyUDF::new(),
-                ))),
+                Arc::new(ScalarUDF::new_from_impl(DummyUDF::new())),
                 vec![
                     Arc::new(BinaryExpr::new(
                         Arc::new(Column::new("b", 1)),
@@ -1431,9 +1428,7 @@ mod tests {
             Arc::new(NegativeExpr::new(Arc::new(Column::new("f", 5)))),
             Arc::new(ScalarFunctionExpr::new(
                 "scalar_expr",
-                ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
-                    DummyUDF::new(),
-                ))),
+                Arc::new(ScalarUDF::new_from_impl(DummyUDF::new())),
                 vec![
                     Arc::new(BinaryExpr::new(
                         Arc::new(Column::new("b", 1)),
@@ -1502,9 +1497,7 @@ mod tests {
             Arc::new(NegativeExpr::new(Arc::new(Column::new("f", 4)))),
             Arc::new(ScalarFunctionExpr::new(
                 "scalar_expr",
-                ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
-                    DummyUDF::new(),
-                ))),
+                Arc::new(ScalarUDF::new_from_impl(DummyUDF::new())),
                 vec![
                     Arc::new(BinaryExpr::new(
                         Arc::new(Column::new("b", 1)),
@@ -1570,9 +1563,7 @@ mod tests {
             Arc::new(NegativeExpr::new(Arc::new(Column::new("f_new", 5)))),
             Arc::new(ScalarFunctionExpr::new(
                 "scalar_expr",
-                ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
-                    DummyUDF::new(),
-                ))),
+                Arc::new(ScalarUDF::new_from_impl(DummyUDF::new())),
                 vec![
                     Arc::new(BinaryExpr::new(
                         Arc::new(Column::new("b_new", 1)),

diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
@@ -26,11 +26,11 @@ use std::sync::Arc;
 use crate::expr_fn::binary_expr;
 use crate::logical_plan::Subquery;
 use crate::utils::expr_to_columns;
-use crate::window_frame;
 use crate::{
     aggregate_function, built_in_window_function, udaf, ExprSchemable, Operator,
     Signature,
 };
+use crate::{window_frame, Volatility};
 
 use arrow::datatypes::{DataType, FieldRef};
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
@@ -399,60 +399,26 @@ impl Between {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-/// Defines which implementation of a function for DataFusion to call.
-pub enum ScalarFunctionDefinition {
-    /// Resolved to a user defined function
-    UDF(Arc<crate::ScalarUDF>),
-}
-
 /// ScalarFunction expression invokes a built-in scalar function
 #[derive(Clone, PartialEq, Eq, Hash, Debug)]
 pub struct ScalarFunction {
     /// The function
-    pub func_def: ScalarFunctionDefinition,
+    pub func: Arc<crate::ScalarUDF>,
     /// List of expressions to feed to the functions as arguments
     pub args: Vec<Expr>,
 }
 
 impl ScalarFunction {
     // return the Function's name
     pub fn name(&self) -> &str {
-        self.func_def.name()
-    }
-}
-
-impl ScalarFunctionDefinition {
-    /// Function's name for display
-    pub fn name(&self) -> &str {
-        match self {
-            ScalarFunctionDefinition::UDF(udf) => udf.name(),
-        }
-    }
-
-    /// Whether this function is volatile, i.e. whether it can return different results
-    /// when evaluated multiple times with the same input.
-    pub fn is_volatile(&self) -> Result<bool> {
-        match self {
-            ScalarFunctionDefinition::UDF(udf) => {
-                Ok(udf.signature().volatility == crate::Volatility::Volatile)
-            }
-        }
+        self.func.name()
     }
 }
 
 impl ScalarFunction {
     /// Create a new ScalarFunction expression with a user-defined function (UDF)
     pub fn new_udf(udf: Arc<crate::ScalarUDF>, args: Vec<Expr>) -> Self {
-        Self {
-            func_def: ScalarFunctionDefinition::UDF(udf),
-            args,
-        }
-    }
-
-    /// Create a new ScalarFunction expression with a user-defined function (UDF)
-    pub fn new_func_def(func_def: ScalarFunctionDefinition, args: Vec<Expr>) -> Self {
-        Self { func_def, args }
+        Self { func: udf, args }
     }
 }
 
@@ -1299,7 +1265,7 @@ impl Expr {
     /// results when evaluated multiple times with the same input.
     pub fn is_volatile(&self) -> Result<bool> {
         self.exists(|expr| {
-            Ok(matches!(expr, Expr::ScalarFunction(func) if func.func_def.is_volatile()?))
+            Ok(matches!(expr, Expr::ScalarFunction(func) if func.func.signature().volatility == Volatility::Volatile ))
         })
     }
 
@@ -1334,9 +1300,7 @@ impl Expr {
     /// and thus any side effects (like divide by zero) may not be encountered
     pub fn short_circuits(&self) -> bool {
         match self {
-            Expr::ScalarFunction(ScalarFunction { func_def, .. }) => {
-                matches!(func_def, ScalarFunctionDefinition::UDF(fun) if fun.short_circuits())
-            }
+            Expr::ScalarFunction(ScalarFunction { func, .. }) => func.short_circuits(),
             Expr::BinaryExpr(BinaryExpr { op, .. }) => {
                 matches!(op, Operator::And | Operator::Or)
             }
@@ -2071,7 +2035,7 @@ mod test {
     }
 
     #[test]
-    fn test_is_volatile_scalar_func_definition() {
+    fn test_is_volatile_scalar_func() {
         // UDF
         #[derive(Debug)]
         struct TestScalarUDF {
@@ -2100,7 +2064,7 @@ mod test {
         let udf = Arc::new(ScalarUDF::from(TestScalarUDF {
             signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable),
         }));
-        assert!(!ScalarFunctionDefinition::UDF(udf).is_volatile().unwrap());
+        assert_ne!(udf.signature().volatility, Volatility::Volatile);
 
         let udf = Arc::new(ScalarUDF::from(TestScalarUDF {
             signature: Signature::uniform(
@@ -2109,7 +2073,7 @@ mod test {
                 Volatility::Volatile,
             ),
         }));
-        assert!(ScalarFunctionDefinition::UDF(udf).is_volatile().unwrap());
+        assert_eq!(udf.signature().volatility, Volatility::Volatile);
     }
 
     use super::*;

diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs
@@ -19,7 +19,7 @@ use super::{Between, Expr, Like};
 use crate::expr::{
     AggregateFunction, AggregateFunctionDefinition, Alias, BinaryExpr, Cast,
     GetFieldAccess, GetIndexedField, InList, InSubquery, Placeholder, ScalarFunction,
-    ScalarFunctionDefinition, Sort, TryCast, Unnest, WindowFunction,
+    Sort, TryCast, Unnest, WindowFunction,
 };
 use crate::field_util::GetFieldAccessSchema;
 use crate::type_coercion::binary::get_result_type;
@@ -133,30 +133,26 @@ impl ExprSchemable for Expr {
                     }
                 }
             }
-            Expr::ScalarFunction(ScalarFunction { func_def, args }) => {
+            Expr::ScalarFunction(ScalarFunction { func, args }) => {
                 let arg_data_types = args
                     .iter()
                     .map(|e| e.get_type(schema))
                     .collect::<Result<Vec<_>>>()?;
-                match func_def {
-                    ScalarFunctionDefinition::UDF(fun) => {
                         // verify that function is invoked with correct number and type of arguments as defined in `TypeSignature`
-                        data_types(&arg_data_types, fun.signature()).map_err(|_| {
+                        data_types(&arg_data_types, func.signature()).map_err(|_| {
                             plan_datafusion_err!(
                                 "{}",
                                 utils::generate_signature_error_msg(
-                                    fun.name(),
-                                    fun.signature().clone(),
+                                    func.name(),
+                                    func.signature().clone(),
                                     &arg_data_types,
                                 )
                             )
                         })?;
 
                         // perform additional function arguments validation (due to limited
                         // expressiveness of `TypeSignature`), then infer return type
-                        Ok(fun.return_type_from_exprs(args, schema, &arg_data_types)?)
-                    }
-                }
+                        Ok(func.return_type_from_exprs(args, schema, &arg_data_types)?)
             }
             Expr::WindowFunction(WindowFunction { fun, args, .. }) => {
                 let data_types = args

diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs
@@ -63,7 +63,7 @@ pub use built_in_window_function::BuiltInWindowFunction;
 pub use columnar_value::ColumnarValue;
 pub use expr::{
     Between, BinaryExpr, Case, Cast, Expr, GetFieldAccess, GetIndexedField, GroupingSet,
-    Like, ScalarFunctionDefinition, TryCast, WindowFunctionDefinition,
+    Like, TryCast, WindowFunctionDefinition,
 };
 pub use expr_fn::*;
 pub use expr_schema::ExprSchemable;

diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs
@@ -20,7 +20,7 @@
 use crate::expr::{
     AggregateFunction, AggregateFunctionDefinition, Alias, Between, BinaryExpr, Case,
     Cast, GetIndexedField, GroupingSet, InList, InSubquery, Like, Placeholder,
-    ScalarFunction, ScalarFunctionDefinition, Sort, TryCast, Unnest, WindowFunction,
+    ScalarFunction, Sort, TryCast, Unnest, WindowFunction,
 };
 use crate::{Expr, GetFieldAccess};
 
@@ -281,11 +281,11 @@ impl TreeNode for Expr {
                 nulls_first,
             }) => transform_box(expr, &mut f)?
                 .update_data(|be| Expr::Sort(Sort::new(be, asc, nulls_first))),
-            Expr::ScalarFunction(ScalarFunction { func_def, args }) => {
-                transform_vec(args, &mut f)?.map_data(|new_args| match func_def {
-                    ScalarFunctionDefinition::UDF(fun) => {
-                        Ok(Expr::ScalarFunction(ScalarFunction::new_udf(fun, new_args)))
-                    }
+            Expr::ScalarFunction(ScalarFunction { func, args }) => {
+                transform_vec(args, &mut f)?.map_data(|new_args| {
+                    Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
+                        func, new_args,
+                    )))
                 })?
             }
             Expr::WindowFunction(WindowFunction {

diff --git a/datafusion/functions-array/src/rewrite.rs b/datafusion/functions-array/src/rewrite.rs
@@ -182,20 +182,20 @@ impl FunctionRewrite for ArrayFunctionRewriter {
 /// Returns true if expr is a function call to the specified named function.
 /// Returns false otherwise.
 fn is_func(expr: &Expr, func_name: &str) -> bool {
-    let Expr::ScalarFunction(ScalarFunction { func_def, args: _ }) = expr else {
+    let Expr::ScalarFunction(ScalarFunction { func, args: _ }) = expr else {
         return false;
     };
 
-    func_def.name() == func_name
+    func.name() == func_name
 }
 
 /// Returns true if expr is a function call with one of the specified names
 fn is_one_of_func(expr: &Expr, func_names: &[&str]) -> bool {
-    let Expr::ScalarFunction(ScalarFunction { func_def, args: _ }) = expr else {
+    let Expr::ScalarFunction(ScalarFunction { func, args: _ }) = expr else {
         return false;
     };
 
-    func_names.contains(&func_def.name())
+    func_names.contains(&func.name())
 }
 
 /// returns Some(col) if this is Expr::Column

diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs
@@ -24,9 +24,7 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
-use datafusion_expr::{
-    lit, ColumnarValue, Expr, FuncMonotonicity, ScalarFunctionDefinition,
-};
+use datafusion_expr::{lit, ColumnarValue, Expr, FuncMonotonicity, ScalarUDF};
 
 use arrow::array::{ArrayRef, Float32Array, Float64Array};
 use datafusion_expr::TypeSignature::*;
@@ -178,8 +176,8 @@ impl ScalarUDFImpl for LogFunc {
                     &info.get_data_type(&base)?,
                 )?)))
             }
-            Expr::ScalarFunction(ScalarFunction { func_def, mut args })
-                if is_pow(&func_def) && args.len() == 2 && base == args[0] =>
+            Expr::ScalarFunction(ScalarFunction { func, mut args })
+                if is_pow(&func) && args.len() == 2 && base == args[0] =>
             {
                 let b = args.pop().unwrap(); // length checked above
                 Ok(ExprSimplifyResult::Simplified(b))
@@ -207,15 +205,8 @@ impl ScalarUDFImpl for LogFunc {
 }
 
 /// Returns true if the function is `PowerFunc`
-fn is_pow(func_def: &ScalarFunctionDefinition) -> bool {
-    match func_def {
-        ScalarFunctionDefinition::UDF(fun) => fun
-            .as_ref()
-            .inner()
-            .as_any()
-            .downcast_ref::<PowerFunc>()
-            .is_some(),
-    }
+fn is_pow(func: &ScalarUDF) -> bool {
+    func.inner().as_any().downcast_ref::<PowerFunc>().is_some()
 }
 
 #[cfg(test)]

diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs
@@ -23,7 +23,7 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
-use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionDefinition};
+use datafusion_expr::{ColumnarValue, Expr, ScalarUDF};
 
 use arrow::array::{ArrayRef, Float64Array, Int64Array};
 use datafusion_expr::TypeSignature::*;
@@ -140,8 +140,8 @@ impl ScalarUDFImpl for PowerFunc {
             Expr::Literal(value) if value == ScalarValue::new_one(&exponent_type)? => {
                 Ok(ExprSimplifyResult::Simplified(base))
             }
-            Expr::ScalarFunction(ScalarFunction { func_def, mut args })
-                if is_log(&func_def) && args.len() == 2 && base == args[0] =>
+            Expr::ScalarFunction(ScalarFunction { func, mut args })
+                if is_log(&func) && args.len() == 2 && base == args[0] =>
             {
                 let b = args.pop().unwrap(); // length checked above
                 Ok(ExprSimplifyResult::Simplified(b))
@@ -152,15 +152,8 @@ impl ScalarUDFImpl for PowerFunc {
 }
 
 /// Return true if this function call is a call to `Log`
-fn is_log(func_def: &ScalarFunctionDefinition) -> bool {
-    match func_def {
-        ScalarFunctionDefinition::UDF(fun) => fun
-            .as_ref()
-            .inner()
-            .as_any()
-            .downcast_ref::<LogFunc>()
-            .is_some(),
-    }
+fn is_log(func: &ScalarUDF) -> bool {
+    func.inner().as_any().downcast_ref::<LogFunc>().is_some()
 }
 
 #[cfg(test)]