diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index f9c0662e682e8..e6fd3230326bc 100644 --- a/datafusion/expr/src/expr.rs +++ b/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, /// List of expressions to feed to the functions as arguments pub args: Vec, + /// 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 { + (&self.func, &self.args).partial_cmp(&(&other.func, &other.args)) + } +} + +impl Hash for ScalarFunction { + fn hash(&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, args: Vec) -> 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 { + Ok(DataType::Int32) + } + + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> Result { + 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; diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 518fe5e3eb6fb..e994b39c9c880 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/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::>>()?; - 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 { + Ok(DataType::Int32) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result { + 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 { + 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(); diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs index 941fd22ea179f..5ea461ea86ebe 100644 --- a/datafusion/expr/src/tree_node.rs +++ b/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) => { diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index 0f680469f6023..def2daf53fc4b 100644 --- a/datafusion/functions-nested/src/array_has.rs +++ b/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 diff --git a/datafusion/functions-nested/src/extract.rs b/datafusion/functions-nested/src/extract.rs index cb7a316b289a9..86ee7ea9b7639 100644 --- a/datafusion/functions-nested/src/extract.rs +++ b/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(), diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index 6ec874fb672d1..f960c455b9ba3 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -307,6 +307,7 @@ fn simplify_get_field_over_struct_constructor(args: &[Expr]) -> Option { 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::() { diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs index 732cfff6cf053..4d1764aa78a5d 100644 --- a/datafusion/functions/src/math/log.rs +++ b/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] diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index 30ac401b5ff7d..96a027e441e4a 100644 --- a/datafusion/functions/src/math/power.rs +++ b/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] diff --git a/datafusion/functions/src/string/concat.rs b/datafusion/functions/src/string/concat.rs index 1c1f6d640798a..6c8619ddaf818 100644 --- a/datafusion/functions/src/string/concat.rs +++ b/datafusion/functions/src/string/concat.rs @@ -394,6 +394,7 @@ pub(crate) fn simplify_concat(args: Vec) -> Result { ScalarFunction { func: concat(), args: new_args, + spans: Default::default(), }, ))) } else { diff --git a/datafusion/functions/src/string/concat_ws.rs b/datafusion/functions/src/string/concat_ws.rs index 8cb6869974813..adea732ea5811 100644 --- a/datafusion/functions/src/string/concat_ws.rs +++ b/datafusion/functions/src/string/concat_ws.rs @@ -403,6 +403,7 @@ fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result Result { let case = coerce_case_expression(case, self.schema)?; Ok(Transformed::yes(Expr::Case(case))) } - Expr::ScalarFunction(ScalarFunction { func, args }) => { + Expr::ScalarFunction(ScalarFunction { func, args, spans }) => { let new_expr = - coerce_arguments_for_signature(args, self.schema, func.as_ref())?; - Ok(Transformed::yes(Expr::ScalarFunction( - ScalarFunction::new_udf(func, new_expr), - ))) + coerce_arguments_for_signature(args, self.schema, func.as_ref()) + .map_err(|error| { + 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 + } + })?; + Ok(Transformed::yes(Expr::ScalarFunction(ScalarFunction { + func, + args: new_expr, + spans, + }))) } Expr::AggregateFunction(expr::AggregateFunction { func, @@ -1468,9 +1482,9 @@ mod test { }; use crate::assert_analyzed_plan_with_config_eq_snapshot; use datafusion_common::config::ConfigOptions; - use datafusion_common::tree_node::{TransformedResult, TreeNode}; + use datafusion_common::tree_node::{TransformedResult, TreeNode, TreeNodeRewriter}; use datafusion_common::{ - DFSchema, DFSchemaRef, Result, ScalarValue, Spans, TableReference, + DFSchema, DFSchemaRef, Location, Result, ScalarValue, Span, Spans, TableReference, }; use datafusion_expr::expr::{self, InSubquery, Like, ScalarFunction}; use datafusion_expr::logical_plan::{EmptyRelation, Projection, Sort}; @@ -2094,6 +2108,62 @@ mod test { ) } + #[test] + fn scalar_function_source_span_survives_type_coercion() -> Result<()> { + let schema = Arc::new(DFSchema::empty()); + let fun = ScalarUDF::new_from_impl(TestScalarUDF { + signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable), + }); + let span = Span::new( + Location { line: 2, column: 4 }, + Location { + line: 2, + column: 21, + }, + ); + let mut function = ScalarFunction::new_udf(Arc::new(fun), vec![lit(10i64)]); + function.spans.add_span(span); + + let result = Expr::ScalarFunction(function) + .rewrite(&mut TypeCoercionRewriter { schema: &schema }) + .data()?; + assert_eq!(result.spans().and_then(Spans::first), Some(span)); + Ok(()) + } + + #[test] + fn nested_scalar_function_error_keeps_inner_span() { + let schema = Arc::new(DFSchema::empty()); + let fun = Arc::new(ScalarUDF::new_from_impl(TestScalarUDF { + signature: Signature::uniform(1, vec![DataType::Float32], Volatility::Stable), + })); + let inner_span = Span::new( + Location { + line: 1, + column: 18, + }, + Location { + line: 1, + column: 35, + }, + ); + let mut inner = ScalarFunction::new_udf(Arc::clone(&fun), vec![lit("Apple")]); + inner.spans.add_span(inner_span); + let mut outer = ScalarFunction::new_udf(fun, vec![Expr::ScalarFunction(inner)]); + outer.spans.add_span(Span::new( + Location { line: 1, column: 8 }, + Location { + line: 1, + column: 42, + }, + )); + + let error = TypeCoercionRewriter { schema: &schema } + .f_up(Expr::ScalarFunction(outer)) + .unwrap_err(); + assert_eq!(error.diagnostic().and_then(|d| d.span), Some(inner_span)); + } + #[test] fn agg_udaf() -> Result<()> { let empty = empty(); diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index 5dfa9c61553ac..c89cba935b7c9 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -655,7 +655,7 @@ impl CSEController for ExprCSEController<'_> { // In case of `ScalarFunction`s and `HigherOrderFunction`s we don't know which children are surely // executed so start visiting all children conditionally and stop the // recursion with `TreeNodeRecursion::Jump`. - Expr::ScalarFunction(ScalarFunction { func, args }) => { + Expr::ScalarFunction(ScalarFunction { func, args, .. }) => { func.conditional_arguments(args) } Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => { diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 2b606687d47a3..4cb9830aa3185 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -1613,17 +1613,27 @@ impl TreeNodeRewriter for Simplifier<'_> { .not(), ) } - Expr::ScalarFunction(ScalarFunction { func: udf, args }) => { - match udf.simplify(args, info)? { - ExprSimplifyResult::Original(args) => { - Transformed::no(Expr::ScalarFunction(ScalarFunction { - func: udf, - args, - })) + Expr::ScalarFunction(ScalarFunction { + func: udf, + args, + spans, + }) => match udf.simplify(args, info)? { + ExprSimplifyResult::Original(args) => { + Transformed::no(Expr::ScalarFunction(ScalarFunction { + func: udf, + args, + spans, + })) + } + ExprSimplifyResult::Simplified(mut expr) => { + if let Expr::ScalarFunction(function) = &mut expr + && function.spans.get_spans().is_empty() + { + function.spans = spans; } - ExprSimplifyResult::Simplified(expr) => Transformed::yes(expr), + Transformed::yes(expr) } - } + }, Expr::AggregateFunction(datafusion_expr::expr::AggregateFunction { ref func, @@ -2113,7 +2123,7 @@ fn get_preimage( right_expr: &Expr, info: &SimplifyContext, ) -> Result { - let Expr::ScalarFunction(ScalarFunction { func, args }) = left_expr else { + let Expr::ScalarFunction(ScalarFunction { func, args, .. }) = left_expr else { return Ok(PreimageResult::None); }; if !is_literal_or_literal_cast(right_expr) { diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index a0902fffe6dab..aeb7545908788 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -452,7 +452,7 @@ pub fn create_physical_expr( execution_props, planning_ctx, )?), - Expr::ScalarFunction(ScalarFunction { func, args }) => { + Expr::ScalarFunction(ScalarFunction { func, args, .. }) => { let physical_args = create_physical_exprs( args, input_dfschema, diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index 16c3468465541..df3cd1c8c2011 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -258,7 +258,7 @@ pub fn serialize_expr( "Proto serialization error: Scalar Variable not supported".to_string(), )); } - Expr::ScalarFunction(ScalarFunction { func, args }) => { + Expr::ScalarFunction(ScalarFunction { func, args, .. }) => { let mut buf = Vec::new(); let _ = codec.try_encode_udf(func, &mut buf); protobuf::LogicalExprNode { diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index 783a2cfd4e476..1c4e89a4fe63e 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -393,7 +393,10 @@ impl SqlToRel<'_, S> { } let RawScalarExpr { func, args } = scalar_expr; - let inner = ScalarFunction::new_udf(func, args); + let mut inner = ScalarFunction::new_udf(func, args); + if let Some(span) = Span::try_from_sqlparser_span(sql_parser_span) { + inner.spans.add_span(span); + } if name.eq_ignore_ascii_case(inner.name()) { return Ok(Expr::ScalarFunction(inner)); diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index 9403e15406344..6c1fbf1671f06 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -170,7 +170,7 @@ impl Unparser<'_> { negated: *negated, }) } - Expr::ScalarFunction(ScalarFunction { func, args }) => { + Expr::ScalarFunction(ScalarFunction { func, args, .. }) => { let func_name = func.name(); if let Some(expr) = self @@ -3308,6 +3308,7 @@ mod tests { Expr::Cast(Cast::new(Box::new(col("a")), DataType::Float64)), Expr::Literal(ScalarValue::Int64(Some(2)), None), ], + spans: Default::default(), }); let ast = unparser.expr_to_sql(&expr)?; @@ -3487,6 +3488,7 @@ mod tests { let expr = Expr::ScalarFunction(ScalarFunction { func: Arc::new(ScalarUDF::from(FromUnixtimeFunc::new())), args: vec![col("date_col")], + spans: Default::default(), }); let ast = unparser.expr_to_sql(&expr)?; @@ -3571,6 +3573,7 @@ mod tests { Expr::Literal(ScalarValue::Utf8(Some(precision.to_string())), None), col("date_col"), ], + spans: Default::default(), }); let ast = unparser.expr_to_sql(&expr)?; @@ -3829,6 +3832,7 @@ mod tests { datafusion_functions::datetime::date_part::DatePartFunc::new(), )), args: vec![lit("YEAR"), col("date_col")], + spans: Default::default(), }); let actual = format!("{}", unparser.expr_to_sql(&expr)?); assert_eq!(actual, "EXTRACT(YEAR FROM `date_col`)"); diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index e824399de98d6..31611f8f47300 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -2412,6 +2412,61 @@ fn scalar_udf_keeps_precedence_over_same_arity_aggregate_udf() { ); } +#[test] +fn scalar_udf_type_error_identifies_the_failing_invocation() { + let udf = Arc::new(make_udf( + "exact_int", + vec![DataType::Int32], + DataType::Int32, + )); + for (sql, failing_call) in [ + ( + "SELECT exact_int(age), exact_int(first_name) FROM person", + "exact_int(first_name)", + ), + ( + "SELECT exact_int(first_name), exact_int(age) FROM person", + "exact_int(first_name)", + ), + ( + "SELECT exact_int(exact_int(first_name)) FROM person", + "exact_int(first_name)", + ), + ( + "SELECT exact_int(age),\n exact_int(first_name) FROM person", + "exact_int(first_name)", + ), + ] { + let state = mock_session_state().with_scalar_function(Arc::clone(&udf)); + let error = logical_plan_from_state( + sql, + &GenericDialect {}, + ParserOptions::default(), + state, + ) + .unwrap_err(); + let span = error + .diagnostic() + .and_then(|diagnostic| diagnostic.span) + .expect("UDF type error should carry its SQL source span"); + let prefix = &sql[..sql.find(failing_call).unwrap()]; + assert_eq!( + usize::try_from(span.start.line).unwrap_or(0), + prefix.matches('\n').count() + 1, + "{sql}" + ); + assert_eq!( + usize::try_from(span.start.column).unwrap_or(0), + prefix + .rsplit_once('\n') + .map_or(prefix, |(_, line)| line) + .len() + + 1, + "{sql}" + ); + } +} + #[test] fn aggregate_expr_planner_can_resolve_qualified_wildcard_from_schema() { let state =