Skip to content

feat: attach Diagnostic to "invalid function argument types" error #23063

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
choplin wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: attach Diagnostic to "invalid function argument types" error #23063

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions datafusion/core/tests/execution/logical_plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ async fn count_only_nulls() -> Result<()> {
order_by: vec![],
null_treatment: None,
},
spans: Spans::new(),
})],
)?);

Expand Down
41 changes: 35 additions & 6 deletions datafusion/expr/src/expr.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -930,21 +930,32 @@ pub struct ScalarFunction {
pub func: Arc<crate::ScalarUDF>,
/// List of expressions to feed to the functions as arguments
pub args: Vec<Expr>,
/// Original source code location, if known
pub spans: Spans,
}

impl ScalarFunction {
// return the Function's name
pub fn name(&self) -> &str {
self.func.name()
}

/// Returns a mutable reference to the spans
pub fn spans_mut(&mut self) -> &mut Spans {
&mut self.spans
}
}

impl ScalarFunction {
/// Create a new `ScalarFunction` from a [`ScalarUDF`]
///
/// [`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(),
}
}
}

Expand Down Expand Up @@ -1094,6 +1105,8 @@ pub struct AggregateFunction {
/// Name of the function
pub func: Arc<AggregateUDF>,
pub params: AggregateFunctionParams,
/// Original source code location, if known
pub spans: Spans,
}

#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
Expand Down Expand Up @@ -1127,8 +1140,14 @@ impl AggregateFunction {
order_by,
null_treatment,
},
spans: Spans::new(),
}
}

/// Returns a mutable reference to the spans
pub fn spans_mut(&mut self) -> &mut Spans {
&mut self.spans
}
}

/// A function used as a SQL window function
Expand Down Expand Up @@ -2287,6 +2306,8 @@ impl Expr {
match self {
Expr::Column(col) => Some(&col.spans),
Expr::Not(inner) | Expr::Negative(inner) => inner.spans(),
Expr::ScalarFunction(func) => Some(&func.spans),
Expr::AggregateFunction(func) => Some(&func.spans),
_ => None,
}
}
Expand Down Expand Up @@ -2482,10 +2503,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()
Expand All @@ -2506,6 +2529,7 @@ impl NormalizeEq for Expr {
order_by: self_order_by,
null_treatment: self_null_treatment,
},
..
}),
Expr::AggregateFunction(AggregateFunction {
func: other_func,
Expand All @@ -2517,6 +2541,7 @@ impl NormalizeEq for Expr {
order_by: other_order_by,
null_treatment: other_null_treatment,
},
..
}),
) => {
self_func.name() == other_func.name()
Expand Down Expand Up @@ -2797,7 +2822,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 {
Expand All @@ -2810,6 +2837,7 @@ impl HashNode for Expr {
order_by: _,
null_treatment,
},
..
}) => {
func.hash(state);
distinct.hash(state);
Expand Down Expand Up @@ -2952,7 +2980,7 @@ impl Display for SchemaDisplay<'_> {
| Expr::OuterReferenceColumn(..)
| Expr::Placeholder(_)
| Expr::Wildcard { .. } => write!(f, "{}", self.0),
Expr::AggregateFunction(AggregateFunction { func, params }) => {
Expr::AggregateFunction(AggregateFunction { func, params, .. }) => {
match func.schema_name(params) {
Ok(name) => {
write!(f, "{name}")
Expand Down Expand Up @@ -3109,7 +3137,7 @@ impl Display for SchemaDisplay<'_> {
Expr::Unnest(Unnest { expr }) => {
write!(f, "UNNEST({})", SchemaDisplay(expr))
}
Expr::ScalarFunction(ScalarFunction { func, args }) => {
Expr::ScalarFunction(ScalarFunction { func, args, .. }) => {
match func.schema_name(args) {
Ok(name) => {
write!(f, "{name}")
Expand Down Expand Up @@ -3408,7 +3436,7 @@ impl Display for SqlDisplay<'_> {

Ok(())
}
Expr::AggregateFunction(AggregateFunction { func, params }) => {
Expr::AggregateFunction(AggregateFunction { func, params, .. }) => {
match func.human_display(params) {
Ok(name) => {
write!(f, "{name}")
Expand Down Expand Up @@ -3639,7 +3667,7 @@ impl Display for Expr {
}
}
}
Expr::AggregateFunction(AggregateFunction { func, params }) => {
Expr::AggregateFunction(AggregateFunction { func, params, .. }) => {
match func.display_name(params) {
Ok(name) => {
write!(f, "{name}")
Expand Down Expand Up @@ -4399,6 +4427,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!()
Expand Down
60 changes: 42 additions & 18 deletions datafusion/expr/src/expr_schema.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,8 @@ use arrow::datatypes::FieldRef;
use arrow::datatypes::{DataType, Field};
use datafusion_common::datatype::FieldExt;
use datafusion_common::{
Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference,
not_impl_err, plan_datafusion_err, plan_err,
Column, DataFusionError, Diagnostic, ExprSchema, Result, ScalarValue, Span, Spans,
TableReference, not_impl_err, plan_datafusion_err, plan_err,
};
use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer;
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
Expand Down Expand Up @@ -549,13 +549,13 @@ impl ExprSchemable for Expr {
match fun {
WindowFunctionDefinition::AggregateUDF(udaf) => {
let new_fields =
verify_function_arguments(udaf.as_ref(), &fields)?;
verify_function_arguments(udaf.as_ref(), &fields, None)?;
let return_field = udaf.return_field(&new_fields)?;
Ok(return_field)
}
WindowFunctionDefinition::WindowUDF(udwf) => {
let new_fields =
verify_function_arguments(udwf.as_ref(), &fields)?;
verify_function_arguments(udwf.as_ref(), &fields, None)?;
let return_field = udwf
.field(WindowUDFFieldArgs::new(&new_fields, &schema_name))?;
Ok(return_field)
Expand All @@ -565,20 +565,23 @@ impl ExprSchemable for Expr {
Expr::AggregateFunction(AggregateFunction {
func,
params: AggregateFunctionParams { 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 new_fields =
verify_function_arguments(func.as_ref(), &fields, spans.first())?;
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 new_fields =
verify_function_arguments(func.as_ref(), &fields, spans.first())?;

let arguments = args
.iter()
Expand Down Expand Up @@ -720,25 +723,46 @@ impl ExprSchemable for Expr {
fn verify_function_arguments<F: UDFCoercionExt>(
function: &F,
input_fields: &[FieldRef],
func_span: Option<Span>,
) -> Result<Vec<FieldRef>> {
fields_with_udf(input_fields, function).map_err(|err| {
let data_types = input_fields
.iter()
.map(|f| f.data_type())
.cloned()
.collect::<Vec<_>>();
plan_datafusion_err!(
"{}. {}",
match err {
DataFusionError::Plan(msg) => msg,
err => err.to_string(),
},
utils::generate_signature_error_message(
function.name(),
function.signature(),
&data_types
)
let name = function.name();
let signature_msg = utils::generate_signature_error_message(
name,
function.signature(),
&data_types,
);
let err_msg = match err {
DataFusionError::Plan(msg) => msg,
err => err.to_string(),
};

let types_str = data_types
.iter()
.map(|dt| dt.to_string())
.collect::<Vec<_>>()
.join(", ");
let candidates = function
.signature()
.type_signature
.to_string_repr_with_names(function.signature().parameter_names.as_deref())
.iter()
.map(|args_str| format!("{name}({args_str})"))
.collect::<Vec<_>>()
.join(", ");
let diagnostic = Diagnostic::new_error(
format!("invalid argument type(s) for '{name}'"),
func_span,
)
.with_note(format!("called with argument type(s): {types_str}"), None)
.with_help(format!("candidate function(s): {candidates}"), None);

plan_datafusion_err!("{err_msg}. {signature_msg}").with_diagnostic(diagnostic)
})
}

Expand Down
28 changes: 17 additions & 11 deletions datafusion/expr/src/tree_node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) => {
Expand Down Expand Up @@ -304,16 +306,20 @@ impl TreeNode for Expr {
order_by,
null_treatment,
},
spans,
}) => (args, filter, order_by).map_elements(f)?.map_data(
|(new_args, new_filter, new_order_by)| {
Ok(Expr::AggregateFunction(AggregateFunction::new_udf(
Ok(Expr::AggregateFunction(AggregateFunction {
func,
new_args,
distinct,
new_filter,
new_order_by,
null_treatment,
)))
params: AggregateFunctionParams {
args: new_args,
distinct,
filter: new_filter,
order_by: new_order_by,
null_treatment,
},
spans,
}))
},
)?,
Expr::GroupingSet(grouping_set) => match grouping_set {
Expand Down
4 changes: 3 additions & 1 deletion datafusion/functions-aggregate/src/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

//! SQL planning extensions like [`AggregateFunctionPlanner`]

use datafusion_common::Result;
use datafusion_common::{Result, Spans};
use datafusion_expr::{
Expr,
expr::{AggregateFunction, AggregateFunctionParams},
Expand Down Expand Up @@ -52,6 +52,7 @@ impl ExprPlanner for AggregateFunctionPlanner {
order_by,
null_treatment,
},
spans: Spans::new(),
});

let saved_name = NamePreserver::new_for_projection().save(&origin_expr);
Expand All @@ -66,6 +67,7 @@ impl ExprPlanner for AggregateFunctionPlanner {
order_by,
null_treatment,
},
..
}) = origin_expr
else {
unreachable!("")
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions-nested/src/array_has.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,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
Expand Down
1 change: 1 addition & 0 deletions datafusion/functions-nested/src/extract.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1162,6 +1162,7 @@ mod tests {
Expr::Column(Column::new_unqualified("my_array")),
Expr::Column(Column::new_unqualified("my_index")),
],
spans: datafusion_common::Spans::new(),
});
assert_eq!(
ExprSchemable::get_type(&udf_expr, &schema).unwrap(),
Expand Down
1 change: 1 addition & 0 deletions datafusion/functions-nested/src/planner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,7 @@ impl ExprPlanner for FieldAccessPlanner {
order_by,
null_treatment,
},
..
}) if is_array_agg(&func) => Ok(PlannerResult::Planned(
Expr::AggregateFunction(AggregateFunction::new_udf(
nth_value_udaf(),
Expand Down
2 changes: 2 additions & 0 deletions datafusion/functions/src/core/getfield.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -607,6 +608,7 @@ impl ScalarUDFImpl for GetFieldFunc {
if let Expr::ScalarFunction(ScalarFunction {
func,
args: inner_args,
..
}) = current_expr
&& func.inner().is::<GetFieldFunc>()
{
Expand Down
2 changes: 1 addition & 1 deletion datafusion/functions/src/math/log.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -367,7 +367,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] =>
{
let b = args.pop().unwrap(); // length checked above
Expand Down
Loading