Support Serde for ScalarUDF in Physical Expressions

datafusion/core/src/physical_optimizer/projection_pushdown.rs

@@ -1287,6 +1287,7 @@ fn new_join_children(
 #[cfg(test)]
 mod tests {
  use super::*;
+ use std::any::Any;
  use std::sync::Arc;
 
  use crate::datasource::file_format::file_compression_type::FileCompressionType;
@@ -1313,7 +1314,10 @@ mod tests {
  use datafusion_common::{JoinSide, JoinType, Result, ScalarValue, Statistics};
  use datafusion_execution::object_store::ObjectStoreUrl;
  use datafusion_execution::{SendableRecordBatchStream, TaskContext};
- use datafusion_expr::{ColumnarValue, Operator};
+ use datafusion_expr::{
+ ColumnarValue, Operator, ScalarFunctionDefinition, ScalarUDF, ScalarUDFImpl,
+ Signature, Volatility,
+ };
  use datafusion_physical_expr::expressions::{
  BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr,
  };
@@ -1329,6 +1333,42 @@ mod tests {
 
  use itertools::Itertools;
 
+ /// Mocked UDF
+ #[derive(Debug)]
+ struct DummyUDF {
+ signature: Signature,
+ }
+
+ impl DummyUDF {
+ fn new() -> Self {
+ Self {
+ signature: Signature::variadic_any(Volatility::Immutable),
+ }
+ }
+ }
+
+ impl ScalarUDFImpl for DummyUDF {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "dummy_udf"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int32)
+ }
+
+ fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ unimplemented!("DummyUDF::invoke")
+ }
+ }
+
  #[test]
  fn test_update_matching_exprs() -> Result<()> {
  let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
@@ -1345,7 +1385,9 @@ mod tests {
  Arc::new(NegativeExpr::new(Arc::new(Column::new("f", 4)))),
  Arc::new(ScalarFunctionExpr::new(
  "scalar_expr",
- Arc::new(|_: &[ColumnarValue]| unimplemented!("not implemented")),
+ ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
+ DummyUDF::new(),
+ ))),
  vec![
  Arc::new(BinaryExpr::new(
  Arc::new(Column::new("b", 1)),
@@ -1412,7 +1454,9 @@ mod tests {
  Arc::new(NegativeExpr::new(Arc::new(Column::new("f", 5)))),
  Arc::new(ScalarFunctionExpr::new(
  "scalar_expr",
- Arc::new(|_: &[ColumnarValue]| unimplemented!("not implemented")),
+ ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
+ DummyUDF::new(),
+ ))),
  vec![
  Arc::new(BinaryExpr::new(
  Arc::new(Column::new("b", 1)),
@@ -1482,7 +1526,9 @@ mod tests {
  Arc::new(NegativeExpr::new(Arc::new(Column::new("f", 4)))),
  Arc::new(ScalarFunctionExpr::new(
  "scalar_expr",
- Arc::new(|_: &[ColumnarValue]| unimplemented!("not implemented")),
+ ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
+ DummyUDF::new(),
+ ))),
  vec![
  Arc::new(BinaryExpr::new(
  Arc::new(Column::new("b", 1)),
@@ -1549,7 +1595,9 @@ mod tests {
  Arc::new(NegativeExpr::new(Arc::new(Column::new("f_new", 5)))),
  Arc::new(ScalarFunctionExpr::new(
  "scalar_expr",
- Arc::new(|_: &[ColumnarValue]| unimplemented!("not implemented")),
+ ScalarFunctionDefinition::UDF(Arc::new(ScalarUDF::new_from_impl(
+ DummyUDF::new(),
+ ))),
  vec![
  Arc::new(BinaryExpr::new(
  Arc::new(Column::new("b_new", 1)),

datafusion/physical-expr/src/functions.rs

@@ -44,6 +44,7 @@ use arrow_array::Array;
 use datafusion_common::{exec_err, Result, ScalarValue};
 use datafusion_expr::execution_props::ExecutionProps;
 pub use datafusion_expr::FuncMonotonicity;
+use datafusion_expr::ScalarFunctionDefinition;
 use datafusion_expr::{
  type_coercion::functions::data_types, BuiltinScalarFunction, ColumnarValue,
  ScalarFunctionImplementation,
@@ -57,7 +58,7 @@ pub fn create_physical_expr(
  fun: &BuiltinScalarFunction,
  input_phy_exprs: &[Arc<dyn PhysicalExpr>],
  input_schema: &Schema,
-execution_props: &ExecutionProps,
+_execution_props: &ExecutionProps,
 ) -> Result<Arc<dyn PhysicalExpr>> {
  let input_expr_types = input_phy_exprs
  .iter()
@@ -69,14 +70,12 @@ pub fn create_physical_expr(
 
  let data_type = fun.return_type(&input_expr_types)?;
 
- let fun_expr: ScalarFunctionImplementation =
- create_physical_fun(fun, execution_props)?;
-
  let monotonicity = fun.monotonicity();
 
+ let fun_def = ScalarFunctionDefinition::BuiltIn(*fun);
  Ok(Arc::new(ScalarFunctionExpr::new(
  &format!("{fun}"),
-fun_expr,
+fun_def,
  input_phy_exprs.to_vec(),
  data_type,
  monotonicity,
@@ -195,7 +194,6 @@ where
 /// Create a physical scalar function.
 pub fn create_physical_fun(
  fun: &BuiltinScalarFunction,
- _execution_props: &ExecutionProps,
 ) -> Result<ScalarFunctionImplementation> {
  Ok(match fun {
  // math functions

datafusion/physical-expr/src/scalar_function.rs

@@ -34,22 +34,22 @@ use std::fmt::{self, Debug, Formatter};
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
-use crate::functions::out_ordering;
+use crate::functions::{create_physical_fun, out_ordering};
 use crate::physical_expr::{down_cast_any_ref, physical_exprs_equal};
 use crate::sort_properties::SortProperties;
 use crate::PhysicalExpr;
 
 use arrow::datatypes::{DataType, Schema};
 use arrow::record_batch::RecordBatch;
-use datafusion_common::Result;
+use datafusion_common::{internal_err, Result};
 use datafusion_expr::{
  expr_vec_fmt, BuiltinScalarFunction, ColumnarValue, FuncMonotonicity,
-ScalarFunctionImplementation,
+ScalarFunctionDefinition,
 };
 
 /// Physical expression of a scalar function
 pub struct ScalarFunctionExpr {
- fun: ScalarFunctionImplementation,
+ fun: ScalarFunctionDefinition,
  name: String,
  args: Vec<Arc<dyn PhysicalExpr>>,
  return_type: DataType,
@@ -79,7 +79,7 @@ impl ScalarFunctionExpr {
  /// Create a new Scalar function
  pub fn new(
  name: &str,
- fun: ScalarFunctionImplementation,
+ fun: ScalarFunctionDefinition,
  args: Vec<Arc<dyn PhysicalExpr>>,
  return_type: DataType,
  monotonicity: Option<FuncMonotonicity>,
@@ -96,7 +96,7 @@ impl ScalarFunctionExpr {
  }
 
  /// Get the scalar function implementation
- pub fn fun(&self) -> &ScalarFunctionImplementation {
+ pub fn fun(&self) -> &ScalarFunctionDefinition {
  &self.fun
  }
 
@@ -172,8 +172,18 @@ impl PhysicalExpr for ScalarFunctionExpr {
  };
 
  // evaluate the function
- let fun = self.fun.as_ref();
- (fun)(&inputs)
+ match self.fun {
+ ScalarFunctionDefinition::BuiltIn(ref fun) => {
+ let fun = create_physical_fun(fun)?;
+ (fun)(&inputs)
+ }
+ ScalarFunctionDefinition::UDF(ref fun) => fun.invoke(&inputs),
+ ScalarFunctionDefinition::Name(_) => {
+ internal_err!(
+ "Name function must be resolved to one of the other variants prior to physical planning"
+ )
+ }
+ }
  }
 
  fn children(&self) -> Vec<Arc<dyn PhysicalExpr>> {

datafusion/physical-expr/src/udf.rs

@@ -20,7 +20,9 @@ use crate::{PhysicalExpr, ScalarFunctionExpr};
 use arrow_schema::Schema;
 use datafusion_common::{DFSchema, Result};
 pub use datafusion_expr::ScalarUDF;
-use datafusion_expr::{type_coercion::functions::data_types, Expr};
+use datafusion_expr::{
+ type_coercion::functions::data_types, Expr, ScalarFunctionDefinition,
+};
 use std::sync::Arc;
 
 /// Create a physical expression of the UDF.
@@ -45,9 +47,10 @@ pub fn create_physical_expr(
  let return_type =
  fun.return_type_from_exprs(args, input_dfschema, &input_expr_types)?;
 
+ let fun_def = ScalarFunctionDefinition::UDF(Arc::new(fun.clone()));
  Ok(Arc::new(ScalarFunctionExpr::new(
  fun.name(),
-fun.fun(),
+fun_def,
  input_phy_exprs.to_vec(),
  return_type,
  fun.monotonicity()?,

datafusion/proto/proto/datafusion.proto

@@ -1458,6 +1458,7 @@ message PhysicalExprNode {
 message PhysicalScalarUdfNode {
  string name = 1;
  repeated PhysicalExprNode args = 2;
+ optional bytes fun_definition = 3;
  ArrowType return_type = 4;
 }