simplify LambdaUDF coerce_value_types

Author: gstvg

datafusion/expr/src/type_coercion/functions.rs

@@ -167,36 +167,42 @@ pub fn value_fields_with_lambda_udf<L: Clone>(
         LambdaTypeSignature::UserDefined => {
             let arg_types = current_fields
                 .iter()
-                .map(|p| match p {
-                    ValueOrLambda::Value(field) => {
-                        ValueOrLambda::Value(field.data_type().clone())
-                    }
-                    ValueOrLambda::Lambda(_) => ValueOrLambda::Lambda(()),
+                .filter_map(|p| match p {
+                    ValueOrLambda::Value(field) => Some(field.data_type().clone()),
+                    ValueOrLambda::Lambda(_) => None,
                 })
                 .collect::<Vec<_>>();
 
             let coerced_types = func.coerce_value_types(&arg_types)?;
 
-            std::iter::zip(current_fields, coerced_types)
-                .map(|(field, coerce_to)| match (field, coerce_to) {
-                    (ValueOrLambda::Value(field), Some(coerce_to)) => {
-                        Ok(ValueOrLambda::Value(Arc::new(
-                            field.as_ref().clone().with_data_type(coerce_to),
-                        )))
+            if coerced_types.len() != arg_types.len() {
+                return plan_err!(
+                    "{} coerce_value_types should have returned {} items but returned {}",
+                    func.name(),
+                    arg_types.len(),
+                    coerced_types.len()
+                );
+            }
+
+            let mut coerced_types = coerced_types.into_iter();
+
+            Ok(current_fields
+                .iter()
+                .map(|current_field| match current_field {
+                    ValueOrLambda::Value(field) => {
+                        let data_type = coerced_types
+                            .next()
+                            .expect("coerced_types len should have been checked above");
+
+                        ValueOrLambda::Value(Arc::new(
+                            field.as_ref().clone().with_data_type(data_type),
+                        ))
                     }
-                    (ValueOrLambda::Lambda(v), None) => {
-                        Ok(ValueOrLambda::Lambda(v.clone()))
+                    ValueOrLambda::Lambda(lambda) => {
+                        ValueOrLambda::Lambda(lambda.clone())
                     }
-                    (ValueOrLambda::Value(_), None) => plan_err!(
-                        "{} coerce_values_types returned None for a value",
-                        func.name()
-                    ),
-                    (ValueOrLambda::Lambda(_), Some(_)) => plan_err!(
-                        "{} coerce_values_types returned Some for a lambda",
-                        func.name()
-                    ),
                 })
-                .collect()
+                .collect())
         }
         LambdaTypeSignature::VariadicAny => Ok(current_fields.to_vec()),
         LambdaTypeSignature::Any(number) => {

datafusion/expr/src/udlf.rs

@@ -606,21 +606,21 @@ pub trait LambdaUDF: Debug + DynEq + DynHash + Send + Sync {
     /// See the [type coercion module](crate::type_coercion)
     /// documentation for more details on type coercion
     ///
-    /// For example, if your function requires a floating point arguments, but the user calls
-    /// it like `my_func(1::int)` (i.e. with `1` as an integer), coerce_types can return `[DataType::Float64]`
-    /// to ensure the argument is converted to `1::double`
+    /// For example, if your function requires a contiguous list argument, but the user calls
+    /// it like `my_func(c, v -> v+2)` (i.e. with `c` as a ListView), coerce_types can return `[DataType::List(..)]`
+    /// to ensure the argument is converted to a List
     ///
     /// # Parameters
-    /// * `arg_types`: The argument types of the arguments this function with
+    /// * `arg_types`: The argument types of the value arguments of this function, excluding lambdas
     ///
     /// # Return value
     /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call
     /// arguments to these specific types.
-    fn coerce_value_types(
-        &self,
-        _arg_types: &[ValueOrLambda<DataType, ()>],
-    ) -> Result<Vec<Option<DataType>>> {
-        not_impl_err!("Function {} does not implement coerce_types", self.name())
+    fn coerce_value_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        not_impl_err!(
+            "Function {} does not implement coerce_value_types",
+            self.name()
+        )
     }
 
     /// Returns the documentation for this Lambda UDF.

datafusion/functions-nested/src/array_transform.rs

@@ -100,11 +100,16 @@ impl LambdaUDF for ArrayTransform {
         &self.signature
     }
 
-    fn coerce_value_types(
-        &self,
-        arg_types: &[ValueOrLambda<DataType, ()>],
-    ) -> Result<Vec<Option<DataType>>> {
-        let (list, _lambda) = value_lambda_pair(self.name(), arg_types)?;
+    fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        let list = if arg_types.len() == 1 {
+            &arg_types[0]
+        } else {
+            return plan_err!(
+                "{} function requires 1 value arguments, got {}",
+                self.name(),
+                arg_types.len()
+            );
+        };
 
         let coerced = match list {
             DataType::List(_)
@@ -121,7 +126,7 @@ impl LambdaUDF for ArrayTransform {
             }
         };
 
-        Ok(vec![Some(coerced), None])
+        Ok(vec![coerced])
     }
 
     fn lambdas_parameters(

datafusion/sqllogictest/test_files/lambda.slt

@@ -245,13 +245,13 @@ select array_transform(arrow_cast(t.list, 'ListView(Int32)'), a -> a+1) from t;
 query error
 select array_transform();
 ----
-DataFusion error: Execution error: array_transform function requires 2 arguments, got 0
+DataFusion error: Error during planning: array_transform function requires 1 value arguments, got 0
 
 
 query error DataFusion error: Error during planning: array_transform expected a list as first argument, got Int64
 select array_transform(1, v -> v*2);
 
-query error DataFusion error: Error during planning: array_transform expects a value followed by a lambda, got Lambda\(\(\)\) and Value\(List\(Field \{ data_type: Int64, nullable: true \}\)\)
+query error DataFusion error: Error during planning: array_transform expects a value followed by a lambda, got Lambda\(\(\)\) and Value\(Field \{ name: "make_array\(Int64\(1\),Int64\(2\)\)", data_type: List\(Field \{ data_type: Int64, nullable: true \}\), nullable: true \}\)
 select array_transform(v -> v*2, [1, 2]);
 
 query error DataFusion error: Error during planning: lambda defined 3 params but UDF support only 2