chore: move udf registration to better place (#1899)

Author: rluvaton

native/core/src/execution/jni_api.rs

@@ -51,7 +51,9 @@ use crate::{
 use datafusion::common::ScalarValue;
 use datafusion::execution::disk_manager::DiskManagerMode;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::logical_expr::ScalarUDF;
 use datafusion_comet_proto::spark_operator::Operator;
+use datafusion_spark::function::math::expm1::SparkExpm1;
 use futures::stream::StreamExt;
 use jni::objects::JByteBuffer;
 use jni::sys::JNI_FALSE;
@@ -284,6 +286,12 @@ fn prepare_datafusion_session_context(
 
     datafusion::functions_nested::register_all(&mut session_ctx)?;
 
+    // register UDFs from datafusion-spark crate
+    session_ctx.register_udf(ScalarUDF::new_from_impl(SparkExpm1::default()));
+
+    // Must be the last one to override existing functions with the same name
+    datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?;
+
     Ok(session_ctx)
 }
 

native/core/src/execution/planner.rs

@@ -65,8 +65,7 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_comet_spark_expr::{
-    create_comet_physical_fun, create_negate_expr, SparkBitwiseCount, SparkBitwiseNot,
-    SparkDateTrunc, SparkHour, SparkMinute, SparkSecond,
+    create_comet_physical_fun, create_negate_expr, SparkHour, SparkMinute, SparkSecond,
 };
 
 use crate::execution::operators::ExecutionError::GeneralError;
@@ -110,7 +109,6 @@ use datafusion_comet_spark_expr::{
     NormalizeNaNAndZero, RLike, SparkCastOptions, StartsWith, Stddev, StringSpaceExpr,
     SubstringExpr, SumDecimal, TimestampTruncExpr, ToJson, UnboundColumn, Variance,
 };
-use datafusion_spark::function::math::expm1::SparkExpm1;
 use itertools::Itertools;
 use jni::objects::GlobalRef;
 use num::{BigInt, ToPrimitive};
@@ -154,11 +152,6 @@ impl Default for PhysicalPlanner {
 
 impl PhysicalPlanner {
     pub fn new(session_ctx: Arc<SessionContext>) -> Self {
-        // register UDFs from datafusion-spark crate
-        session_ctx.register_udf(ScalarUDF::new_from_impl(SparkExpm1::default()));
-        session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default()));
-        session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseCount::default()));
-        session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateTrunc::default()));
         Self {
             exec_context_id: TEST_EXEC_CONTEXT_ID,
             session_ctx,

native/spark-expr/src/comet_scalar_funcs.rs

@@ -20,7 +20,7 @@ use crate::{
     spark_array_repeat, spark_ceil, spark_date_add, spark_date_sub, spark_decimal_div,
     spark_decimal_integral_div, spark_floor, spark_hex, spark_isnan, spark_make_decimal,
     spark_read_side_padding, spark_round, spark_rpad, spark_unhex, spark_unscaled_value,
-    SparkChrFunc,
+    SparkBitwiseCount, SparkBitwiseNot, SparkChrFunc, SparkDateTrunc,
 };
 use arrow::datatypes::DataType;
 use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -112,7 +112,6 @@ pub fn create_comet_physical_fun(
             let func = Arc::new(spark_xxhash64);
             make_comet_scalar_udf!("xxhash64", func, without data_type)
         }
-        "chr" => Ok(Arc::new(ScalarUDF::new_from_impl(SparkChrFunc::default()))),
         "isnan" => {
             let func = Arc::new(spark_isnan);
             make_comet_scalar_udf!("isnan", func, without data_type)
@@ -153,6 +152,25 @@ pub fn create_comet_physical_fun(
     }
 }
 
+fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> {
+    vec![
+        Arc::new(ScalarUDF::new_from_impl(SparkChrFunc::default())),
+        Arc::new(ScalarUDF::new_from_impl(SparkBitwiseNot::default())),
+        Arc::new(ScalarUDF::new_from_impl(SparkBitwiseCount::default())),
+        Arc::new(ScalarUDF::new_from_impl(SparkDateTrunc::default())),
+    ]
+}
+
+/// Registers all custom UDFs
+pub fn register_all_comet_functions(registry: &mut dyn FunctionRegistry) -> DataFusionResult<()> {
+    // This will override existing UDFs with the same name
+    all_scalar_functions()
+        .into_iter()
+        .try_for_each(|udf| registry.register_udf(udf).map(|_| ()))?;
+
+    Ok(())
+}
+
 struct CometScalarFunction {
     name: String,
     signature: Signature,

native/spark-expr/src/lib.rs

@@ -58,7 +58,7 @@ pub use bitwise_funcs::*;
 pub use conditional_funcs::*;
 pub use conversion_funcs::*;
 
-pub use comet_scalar_funcs::create_comet_physical_fun;
+pub use comet_scalar_funcs::{create_comet_physical_fun, register_all_comet_functions};
 pub use datetime_funcs::{
     spark_date_add, spark_date_sub, SparkDateTrunc, SparkHour, SparkMinute, SparkSecond,
     TimestampTruncExpr,