Initial code for adding casts for UDFs in the logical plan

datafusion/optimizer/src/type_coercion.rs

@@ -24,8 +24,9 @@ use datafusion_expr::binary_rule::coerce_types;
 use datafusion_expr::expr_rewriter::{ExprRewritable, ExprRewriter, RewriteRecursion};
 use datafusion_expr::logical_plan::builder::build_join_schema;
 use datafusion_expr::logical_plan::JoinType;
+use datafusion_expr::type_coercion::data_types;
 use datafusion_expr::utils::from_plan;
-use datafusion_expr::{cast_if_needed, ExprSchemable};
+use datafusion_expr::{cast_if_needed, ExprSchemable, Signature};
 use datafusion_expr::{Expr, LogicalPlan};
 
 #[derive(Default)]
@@ -124,11 +125,49 @@ impl ExprRewriter for TypeCoercionRewriter {
                     }
                 }
             }
+            Expr::ScalarUDF { fun, args } => {
+                let new_expr = coerce_arguments_for_signature(
+                    args.as_slice(),
+                    &self.schema,
+                    &fun.signature,
+                )?;
+                Ok(Expr::ScalarUDF {
+                    fun: fun.clone(),
+                    args: new_expr,
+                })
+            }
             _ => Ok(expr),
         }
     }
 }
 
+/// Returns `expressions` coerced to types compatible with
+/// `signature`, if possible.
+///
+/// See the module level documentation for more detail on coercion.
+pub fn coerce_arguments_for_signature(
+    expressions: &[Expr],
+    schema: &DFSchema,
+    signature: &Signature,
+) -> Result<Vec<Expr>> {
+    if expressions.is_empty() {
+        return Ok(vec![]);
+    }
+
+    let current_types = expressions
+        .iter()
+        .map(|e| e.get_type(schema))
+        .collect::<Result<Vec<_>>>()?;
+
+    let new_types = data_types(&current_types, signature)?;
+
+    expressions
+        .iter()
+        .enumerate()
+        .map(|(i, expr)| expr.clone().cast_to(&new_types[i], schema))
+        .collect::<Result<Vec<_>>>()
+}
+
 #[cfg(test)]
 mod test {
     use crate::type_coercion::TypeCoercion;