Handle one-element array return value in ScalarFunctionExpr (apache#12965) (#276)

This was done in apache#12922 only for math functions.
We now generalize this fallback to all scalar UDFs.

Author: joroKr21

datafusion/expr-common/src/columnar_value.rs

@@ -217,17 +217,6 @@ impl ColumnarValue {
             }
         }
     }
-
-    /// Converts an [`ArrayRef`] to a [`ColumnarValue`] based on the supplied arguments.
-    /// This is useful for scalar UDF implementations to fulfil their contract:
-    /// if all arguments are scalar values, the result should also be a scalar value.
-    pub fn from_args_and_result(args: &[Self], result: ArrayRef) -> Result<Self> {
-        if result.len() == 1 && args.iter().all(|arg| matches!(arg, Self::Scalar(_))) {
-            Ok(Self::Scalar(ScalarValue::try_from_array(&result, 0)?))
-        } else {
-            Ok(Self::Array(result))
-        }
-    }
 }
 
 #[cfg(test)]

datafusion/functions/src/macros.rs

@@ -208,8 +208,8 @@ macro_rules! make_math_unary_udf {
                     $EVALUATE_BOUNDS(inputs)
                 }
 
-                fn invoke(&self, col_args: &[ColumnarValue]) -> Result<ColumnarValue> {
-                    let args = ColumnarValue::values_to_arrays(col_args)?;
+                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(
                             args[0]
@@ -229,7 +229,7 @@ macro_rules! make_math_unary_udf {
                         }
                     };
 
-                    ColumnarValue::from_args_and_result(col_args, arr)
+                    Ok(ColumnarValue::Array(arr))
                 }
 
                 fn documentation(&self) -> Option<&Documentation> {
@@ -316,8 +316,8 @@ macro_rules! make_math_binary_udf {
                     $OUTPUT_ORDERING(input)
                 }
 
-                fn invoke(&self, col_args: &[ColumnarValue]) -> Result<ColumnarValue> {
-                    let args = ColumnarValue::values_to_arrays(col_args)?;
+                fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+                    let args = ColumnarValue::values_to_arrays(args)?;
                     let arr: ArrayRef = match args[0].data_type() {
                         DataType::Float64 => {
                             let y = args[0].as_primitive::<Float64Type>();
@@ -347,7 +347,7 @@ macro_rules! make_math_binary_udf {
                         }
                     };
 
-                    ColumnarValue::from_args_and_result(col_args, arr)
+                    Ok(ColumnarValue::Array(arr))
                 }
 
                 fn documentation(&self) -> Option<&Documentation> {