refactor: add macro for the binary math function in `datafusion-funct…

…ion` (apache#9889) * refactor: macro for the binary math function in datafusion-function * Update datafusion/functions/src/macros.rs --------- Co-authored-by: Andrew Lamb <[email protected]>
korowa · Apr 1, 2024 · b698e5f · b698e5f
1 parent c858455
commit b698e5f
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 143 deletions.
diff --git a/datafusion/functions/src/macros.rs b/datafusion/functions/src/macros.rs
@@ -251,7 +251,112 @@ macro_rules! make_math_unary_udf {
     };
 }
 
-#[macro_export]
+/// Macro to create a binary math UDF.
+///
+/// A binary math function takes two arguments of types Float32 or Float64,
+/// applies a binary floating function to the argument, and returns a value of the same type.
+///
+/// $UDF: the name of the UDF struct that implements `ScalarUDFImpl`
+/// $GNAME: a singleton instance of the UDF
+/// $NAME: the name of the function
+/// $BINARY_FUNC: the binary function to apply to the argument
+/// $MONOTONIC_FUNC: the monotonicity of the function
+macro_rules! make_math_binary_udf {
+    ($UDF:ident, $GNAME:ident, $NAME:ident, $BINARY_FUNC:ident, $MONOTONICITY:expr) => {
+        make_udf_function!($NAME::$UDF, $GNAME, $NAME);
+
+        mod $NAME {
+            use arrow::array::{ArrayRef, Float32Array, Float64Array};
+            use arrow::datatypes::DataType;
+            use datafusion_common::{exec_err, DataFusionError, Result};
+            use datafusion_expr::TypeSignature::*;
+            use datafusion_expr::{
+                ColumnarValue, FuncMonotonicity, ScalarUDFImpl, Signature, Volatility,
+            };
+            use std::any::Any;
+            use std::sync::Arc;
+
+            #[derive(Debug)]
+            pub struct $UDF {
+                signature: Signature,
+            }
+
+            impl $UDF {
+                pub fn new() -> Self {
+                    use DataType::*;
+                    Self {
+                        signature: Signature::one_of(
+                            vec![
+                                Exact(vec![Float32, Float32]),
+                                Exact(vec![Float64, Float64]),
+                            ],
+                            Volatility::Immutable,
+                        ),
+                    }
+                }
+            }
+
+            impl ScalarUDFImpl for $UDF {
+                fn as_any(&self) -> &dyn Any {
+                    self
+                }
+                fn name(&self) -> &str {
+                    stringify!($NAME)
+                }
+
+                fn signature(&self) -> &Signature {
+                    &self.signature
+                }
+
+                fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+                    let arg_type = &arg_types[0];
+
+                    match arg_type {
+                        DataType::Float32 => Ok(DataType::Float32),
+                        // For other types (possible values float64/null/int), use Float64
+                        _ => Ok(DataType::Float64),
+                    }
+                }
+
+                fn monotonicity(&self) -> Result<Option<FuncMonotonicity>> {
+                    Ok($MONOTONICITY)
+                }
+
+                fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+                    let args = ColumnarValue::values_to_arrays(args)?;
+
+                    let arr: ArrayRef = match args[0].data_type() {
+                        DataType::Float64 => Arc::new(make_function_inputs2!(
+                            &args[0],
+                            &args[1],
+                            "y",
+                            "x",
+                            Float64Array,
+                            { f64::$BINARY_FUNC }
+                        )),
+
+                        DataType::Float32 => Arc::new(make_function_inputs2!(
+                            &args[0],
+                            &args[1],
+                            "y",
+                            "x",
+                            Float32Array,
+                            { f32::$BINARY_FUNC }
+                        )),
+                        other => {
+                            return exec_err!(
+                                "Unsupported data type {other:?} for function {}",
+                                self.name()
+                            )
+                        }
+                    };
+                    Ok(ColumnarValue::Array(arr))
+                }
+            }
+        }
+    };
+}
+
 macro_rules! make_function_inputs2 {
     ($ARG1: expr, $ARG2: expr, $NAME1:expr, $NAME2: expr, $ARRAY_TYPE:ident, $FUNC: block) => {{
         let arg1 = downcast_arg!($ARG1, $NAME1, $ARRAY_TYPE);

diff --git a/datafusion/functions/src/math/atan2.rs b/datafusion/functions/src/math/atan2.rs
diff --git a/datafusion/functions/src/math/mod.rs b/datafusion/functions/src/math/mod.rs
@@ -18,13 +18,11 @@
 //! "math" DataFusion functions
 
 mod abs;
-mod atan2;
 mod nans;
 
 // Create UDFs
 make_udf_function!(nans::IsNanFunc, ISNAN, isnan);
 make_udf_function!(abs::AbsFunc, ABS, abs);
-make_udf_function!(atan2::Atan2, ATAN2, atan2);
 
 make_math_unary_udf!(Log2Func, LOG2, log2, log2, Some(vec![Some(true)]));
 make_math_unary_udf!(Log10Func, LOG10, log10, log10, Some(vec![Some(true)]));
@@ -39,6 +37,7 @@ make_math_unary_udf!(AtanhFunc, ATANH, atanh, atanh, Some(vec![Some(true)]));
 make_math_unary_udf!(AsinhFunc, ASINH, asinh, asinh, Some(vec![Some(true)]));
 make_math_unary_udf!(AcoshFunc, ACOSH, acosh, acosh, Some(vec![Some(true)]));
 make_math_unary_udf!(AtanFunc, ATAN, atan, atan, Some(vec![Some(true)]));
+make_math_binary_udf!(Atan2, ATAN2, atan2, atan2, Some(vec![Some(true)]));
 
 // Export the functions out of this package, both as expr_fn as well as a list of functions
 export_functions!(