Preserve SQL source spans for scalar UDF type errors

datafusion/expr/src/expr.rs

@@ -945,12 +945,44 @@ impl Between {
 /// Invoke a [`ScalarUDF`] with a set of arguments
 ///
 /// [`ScalarUDF`]: crate::ScalarUDF
-#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
+#[derive(Clone)]
 pub struct ScalarFunction {
     /// The function
     pub func: Arc<crate::ScalarUDF>,
     /// List of expressions to feed to the functions as arguments
     pub args: Vec<Expr>,
+    /// SQL source locations of this invocation, if known.
+    pub spans: Spans,
+}
+
+impl PartialEq for ScalarFunction {
+    fn eq(&self, other: &Self) -> bool {
+        self.func == other.func && self.args == other.args
+    }
+}
+
+impl Eq for ScalarFunction {}
+
+impl PartialOrd for ScalarFunction {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        (&self.func, &self.args).partial_cmp(&(&other.func, &other.args))
+    }
+}
+
+impl Hash for ScalarFunction {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.func.hash(state);
+        self.args.hash(state);
+    }
+}
+
+impl fmt::Debug for ScalarFunction {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ScalarFunction")
+            .field("func", &self.func)
+            .field("args", &self.args)
+            .finish()
+    }
 }
 
 impl ScalarFunction {
@@ -965,7 +997,11 @@ impl ScalarFunction {
     ///
     /// [`ScalarUDF`]: crate::ScalarUDF
     pub fn new_udf(udf: Arc<crate::ScalarUDF>, args: Vec<Expr>) -> Self {
-        Self { func: udf, args }
+        Self {
+            func: udf,
+            args,
+            spans: Spans::new(),
+        }
     }
 }
 
@@ -2304,6 +2340,7 @@ impl Expr {
     pub fn spans(&self) -> Option<&Spans> {
         match self {
             Expr::Column(col) => Some(&col.spans),
+            Expr::ScalarFunction(function) => Some(&function.spans),
             Expr::Not(inner) | Expr::Negative(inner) => inner.spans(),
             _ => None,
         }
@@ -2508,10 +2545,12 @@ impl NormalizeEq for Expr {
                 Expr::ScalarFunction(ScalarFunction {
                     func: self_func,
                     args: self_args,
+                    ..
                 }),
                 Expr::ScalarFunction(ScalarFunction {
                     func: other_func,
                     args: other_args,
+                    ..
                 }),
             ) => {
                 self_func.name() == other_func.name()
@@ -2823,7 +2862,9 @@ impl HashNode for Expr {
             | Expr::TryCast(TryCast { expr: _expr, field }) => {
                 field.hash(state);
             }
-            Expr::ScalarFunction(ScalarFunction { func, args: _args }) => {
+            Expr::ScalarFunction(ScalarFunction {
+                func, args: _args, ..
+            }) => {
                 func.hash(state);
             }
             Expr::AggregateFunction(AggregateFunction {
@@ -3158,7 +3199,7 @@ impl Display for SchemaDisplay<'_> {
                 let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" };
                 write!(f, "{name}({})", SchemaDisplay(expr))
             }
-            Expr::ScalarFunction(ScalarFunction { func, args }) => {
+            Expr::ScalarFunction(ScalarFunction { func, args, .. }) => {
                 match func.schema_name(args) {
                     Ok(name) => {
                         write!(f, "{name}")
@@ -4555,6 +4596,7 @@ mod test {
         let udf = Expr::ScalarFunction(ScalarFunction {
             func: Arc::new(ScalarUDF::new_from_impl(TestUDF {})),
             args: vec![expr(), expr()],
+            spans: Spans::new(),
         });
         let Expr::ScalarFunction(scalar) = &udf else {
             unreachable!()
@@ -4591,6 +4633,61 @@ mod test {
         }
     }
 
+    #[test]
+    fn scalar_function_source_spans_do_not_change_identity() {
+        use datafusion_common::{Location, Span};
+        use std::collections::hash_map::DefaultHasher;
+
+        #[derive(Debug, PartialEq, Eq, Hash)]
+        struct TestUdf {
+            signature: Signature,
+        }
+
+        impl ScalarUDFImpl for TestUdf {
+            fn name(&self) -> &str {
+                "test_scalar_udf"
+            }
+
+            fn signature(&self) -> &Signature {
+                &self.signature
+            }
+
+            fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+                Ok(DataType::Int32)
+            }
+
+            fn invoke_with_args(
+                &self,
+                _args: ScalarFunctionArgs,
+            ) -> Result<ColumnarValue> {
+                unimplemented!()
+            }
+        }
+
+        let udf = Arc::new(ScalarUDF::new_from_impl(TestUdf {
+            signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),
+        }));
+        let plain = ScalarFunction::new_udf(Arc::clone(&udf), vec![lit(1)]);
+        let mut located = ScalarFunction::new_udf(udf, vec![lit(1)]);
+        located.spans.add_span(Span::new(
+            Location { line: 1, column: 8 },
+            Location {
+                line: 1,
+                column: 23,
+            },
+        ));
+
+        assert_eq!(plain, located);
+        assert_eq!(plain.partial_cmp(&located), Some(Ordering::Equal));
+        assert_eq!(format!("{plain:?}"), format!("{located:?}"));
+        let hash = |function: &ScalarFunction| {
+            let mut hasher = DefaultHasher::new();
+            function.hash(&mut hasher);
+            hasher.finish()
+        };
+        assert_eq!(hash(&plain), hash(&located));
+    }
+
     mod intersect_metadata_tests {
         use super::super::intersect_metadata_for_union;
         use std::collections::HashMap;

datafusion/expr/src/expr_schema.rs

@@ -35,8 +35,8 @@ use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_K
 use datafusion_common::datatype::FieldExt;
 use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{
-    Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference,
-    internal_err, not_impl_err, plan_datafusion_err, plan_err,
+    Column, DataFusionError, Diagnostic, ExprSchema, Result, ScalarValue, Spans,
+    TableReference, internal_err, not_impl_err, plan_datafusion_err, plan_err,
 };
 use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer;
 use datafusion_functions_window_common::field::WindowUDFFieldArgs;
@@ -811,12 +811,25 @@ impl ExprSchemable for Expr {
                 let new_fields = verify_function_arguments(func.as_ref(), &fields)?;
                 func.return_field(&new_fields)
             }
-            Expr::ScalarFunction(ScalarFunction { func, args }) => {
+            Expr::ScalarFunction(ScalarFunction { func, args, spans }) => {
                 let fields = args
                     .iter()
                     .map(|e| e.to_field(schema).map(|(_, f)| f))
                     .collect::<Result<Vec<_>>>()?;
-                let new_fields = verify_function_arguments(func.as_ref(), &fields)?;
+                let with_span = |error: DataFusionError| {
+                    if error.diagnostic().is_some() {
+                        error
+                    } else if let Some(span) = spans.first() {
+                        error.with_diagnostic(Diagnostic::new_error(
+                            format!("Invalid function '{}'", func.name()),
+                            Some(span),
+                        ))
+                    } else {
+                        error
+                    }
+                };
+                let new_fields = verify_function_arguments(func.as_ref(), &fields)
+                    .map_err(&with_span)?;
 
                 let arguments = args
                     .iter()
@@ -830,7 +843,7 @@ impl ExprSchemable for Expr {
                     scalar_arguments: &arguments,
                 };
 
-                func.return_field_from_args(args)
+                func.return_field_from_args(args).map_err(with_span)
             }
             // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)),
             Expr::Cast(Cast { expr, field }) => expr
@@ -1068,6 +1081,80 @@ mod tests {
     use arrow::datatypes::Schema;
     use datafusion_common::{DFSchema, assert_or_internal_err};
 
+    #[test]
+    fn scalar_function_preserves_udf_supplied_diagnostic() {
+        use datafusion_common::{Location, Span};
+
+        #[derive(Debug, PartialEq, Eq, Hash)]
+        struct DiagnosticUdf {
+            signature: crate::Signature,
+        }
+
+        impl crate::ScalarUDFImpl for DiagnosticUdf {
+            fn name(&self) -> &str {
+                "diagnostic_udf"
+            }
+
+            fn signature(&self) -> &crate::Signature {
+                &self.signature
+            }
+
+            fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+                Ok(DataType::Int32)
+            }
+
+            fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
+                Err(
+                    DataFusionError::Plan("bad input".to_string()).with_diagnostic(
+                        Diagnostic::new_error(
+                            "UDF-specific error",
+                            Some(Span::new(
+                                Location {
+                                    line: 1,
+                                    column: 20,
+                                },
+                                Location {
+                                    line: 1,
+                                    column: 21,
+                                },
+                            )),
+                        ),
+                    ),
+                )
+            }
+
+            fn invoke_with_args(
+                &self,
+                _args: crate::ScalarFunctionArgs,
+            ) -> Result<crate::ColumnarValue> {
+                unreachable!("planning must fail before invocation")
+            }
+        }
+
+        let udf = Arc::new(crate::ScalarUDF::new_from_impl(DiagnosticUdf {
+            signature: crate::Signature::exact(
+                vec![DataType::Int32],
+                crate::Volatility::Immutable,
+            ),
+        }));
+        let mut function = ScalarFunction::new_udf(udf, vec![lit(1i32)]);
+        function.spans.add_span(Span::new(
+            Location { line: 1, column: 8 },
+            Location {
+                line: 1,
+                column: 22,
+            },
+        ));
+        let error = Expr::ScalarFunction(function)
+            .to_field(&DFSchema::empty())
+            .unwrap_err();
+        let diagnostic = error
+            .diagnostic()
+            .expect("UDF diagnostic must be preserved");
+        assert_eq!(diagnostic.message, "UDF-specific error");
+        assert_eq!(diagnostic.span.map(|span| span.start.column), Some(20));
+    }
+
     macro_rules! test_is_expr_nullable {
         ($EXPR_TYPE:ident) => {{
             let expr = lit(ScalarValue::Null).$EXPR_TYPE();

datafusion/expr/src/tree_node.rs

@@ -253,11 +253,13 @@ impl TreeNode for Expr {
             Expr::TryCast(TryCast { expr, field }) => expr
                 .map_elements(f)?
                 .update_data(|be| Expr::TryCast(TryCast::new_from_field(be, field))),
-            Expr::ScalarFunction(ScalarFunction { func, args }) => {
+            Expr::ScalarFunction(ScalarFunction { func, args, spans }) => {
                 args.map_elements(f)?.map_data(|new_args| {
-                    Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
-                        func, new_args,
-                    )))
+                    Ok(Expr::ScalarFunction(ScalarFunction {
+                        func,
+                        args: new_args,
+                        spans,
+                    }))
                 })?
             }
             Expr::WindowFunction(window_fun) => {

datafusion/functions-nested/src/array_has.rs

@@ -165,7 +165,7 @@ impl ScalarUDFImpl for ArrayHas {
                     )));
                 }
             }
-            Expr::ScalarFunction(ScalarFunction { func, args })
+            Expr::ScalarFunction(ScalarFunction { func, args, .. })
                 if func == &make_array_udf() =>
             {
                 // make_array has a static set of arguments, so we can pull the arguments out from it

datafusion/functions-nested/src/extract.rs

@@ -1181,6 +1181,7 @@ mod tests {
                 Expr::Column(Column::new_unqualified("my_array")),
                 Expr::Column(Column::new_unqualified("my_index")),
             ],
+            spans: Default::default(),
         });
         assert_eq!(
             ExprSchemable::get_type(&udf_expr, &schema).unwrap(),

datafusion/functions/src/core/getfield.rs

@@ -307,6 +307,7 @@ fn simplify_get_field_over_struct_constructor(args: &[Expr]) -> Option<Expr> {
     let Expr::ScalarFunction(ScalarFunction {
         func,
         args: ctor_args,
+        ..
     }) = base
     else {
         return None;
@@ -607,6 +608,7 @@ impl ScalarUDFImpl for GetFieldFunc {
             if let Expr::ScalarFunction(ScalarFunction {
                 func,
                 args: inner_args,
+                ..
             }) = current_expr
                 && func.inner().is::<GetFieldFunc>()
             {

datafusion/functions/src/math/log.rs

@@ -372,7 +372,7 @@ impl ScalarUDFImpl for LogFunc {
                     &info.get_data_type(&base)?,
                 )?)))
             }
-            Expr::ScalarFunction(ScalarFunction { func, mut args })
+            Expr::ScalarFunction(ScalarFunction { func, mut args, .. })
                 if is_pow(&func)
                     && args.len() == 2
                     && base == args[0]

datafusion/functions/src/math/power.rs

@@ -179,7 +179,7 @@ impl ScalarUDFImpl for PowerFunc {
                     base, &base_type,
                 )))
             }
-            Expr::ScalarFunction(ScalarFunction { func, mut args })
+            Expr::ScalarFunction(ScalarFunction { func, mut args, .. })
                 if is_log(&func)
                     && args.len() == 2
                     && base == args[0]

datafusion/functions/src/string/concat.rs

@@ -394,6 +394,7 @@ pub(crate) fn simplify_concat(args: Vec<Expr>) -> Result<ExprSimplifyResult> {
             ScalarFunction {
                 func: concat(),
                 args: new_args,
+                spans: Default::default(),
             },
         )))
     } else {

datafusion/functions/src/string/concat_ws.rs

@@ -403,6 +403,7 @@ fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result<ExprSimplifyRes
                                 ScalarFunction {
                                     func: concat(),
                                     args: args.to_vec(),
+                                    spans: Default::default(),
                                 },
                             )))
                         }
@@ -460,6 +461,7 @@ fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result<ExprSimplifyRes
                         ScalarFunction {
                             func: concat_ws(),
                             args: new_args,
+                            spans: Default::default(),
                         },
                     )))
                 }