Skip to content

Commit edcb2a9

Browse files
Ian LaiIan Lai
Ian Lai
authored and
Ian Lai
committed
feat: enhance diagnostic error handling for function argument validation
1 parent 4d510a2 commit edcb2a9

File tree

3 files changed

+96
-127
lines changed

3 files changed

+96
-127
lines changed

‎datafusion/expr/src/expr_schema.rs

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,11 @@ impl ExprSchemable for Expr {
170170
"{} {}",
171171
match err {
172172
DataFusionError::Plan(msg) => msg,
173+
DataFusionError::Diagnostic(_, boxed_err) =>
174+
match *boxed_err {
175+
DataFusionError::Plan(msg) => msg,
176+
_ => boxed_err.to_string(),
177+
},
173178
err => err.to_string(),
174179
},
175180
utils::generate_signature_error_msg(
@@ -179,7 +184,15 @@ impl ExprSchemable for Expr {
179184
)
180185
);
181186

182-
if let Some(diagnostic) = diagnostic {
187+
if let Some(mut diagnostic) = diagnostic {
188+
diagnostic.add_help(
189+
utils::generate_signature_error_msg(
190+
func.name(),
191+
func.signature().clone(),
192+
&data_types,
193+
),
194+
None,
195+
);
183196
err.with_diagnostic(diagnostic)
184197
} else {
185198
err
@@ -435,7 +448,9 @@ impl ExprSchemable for Expr {
435448
// Verify that function is invoked with correct number and type of arguments as defined in `TypeSignature`
436449
let new_data_types = data_types_with_scalar_udf(&arg_types, func)
437450
.map_err(|err| {
438-
plan_datafusion_err!(
451+
let diagnostic = err.diagnostic().cloned();
452+
453+
let err = plan_datafusion_err!(
439454
"{} {}",
440455
match err {
441456
DataFusionError::Plan(msg) => msg,
@@ -451,7 +466,21 @@ impl ExprSchemable for Expr {
451466
func.signature().clone(),
452467
&arg_types,
453468
)
454-
)
469+
);
470+
471+
if let Some(mut diagnostic) = diagnostic {
472+
diagnostic.add_help(
473+
utils::generate_signature_error_msg(
474+
func.name(),
475+
func.signature().clone(),
476+
&arg_types,
477+
),
478+
None,
479+
);
480+
err.with_diagnostic(diagnostic)
481+
} else {
482+
err
483+
}
455484
})?;
456485

457486
let arguments = args

‎datafusion/expr/src/type_coercion/functions.rs

Lines changed: 63 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -248,158 +248,100 @@ pub fn check_function_length_with_diag(
248248
if current_types.is_empty() {
249249
if signature.supports_zero_argument() {
250250
return Ok(());
251-
} else if signature.used_to_support_zero_arguments() {
252-
// Special error to help during upgrade
253-
let base_error = plan_datafusion_err!(
254-
"'{}' does not support zero arguments. Use TypeSignature::Nullary for zero arguments",
255-
function_name
256-
);
257-
let mut diagnostic = Diagnostic::new_error(
258-
format!("Zero arguments not supported for {} function", function_name),
259-
function_call_site,
260-
);
261-
diagnostic.add_help(
262-
"Use TypeSignature::Nullary for functions that take no arguments",
263-
None,
264-
);
265-
return Err(base_error.with_diagnostic(diagnostic));
251+
}
252+
253+
let (error_message, note) = if signature.used_to_support_zero_arguments() {
254+
(
255+
format!("Zero arguments not supported for {} function. Use TypeSignature::Nullary for zero arguments", function_name),
256+
"Use TypeSignature::Nullary for functions that take no arguments".to_string()
257+
)
266258
} else {
267-
let base_error = plan_datafusion_err!(
268-
"'{}' does not support zero arguments",
269-
function_name
270-
);
271-
let mut diagnostic = Diagnostic::new_error(
259+
(
272260
format!("Zero arguments not supported for {} function", function_name),
273-
function_call_site,
274-
);
275-
diagnostic.add_note(
276-
format!("Function {} requires at least one argument", function_name),
277-
None,
278-
);
279-
return Err(base_error.with_diagnostic(diagnostic));
280-
}
261+
format!("Function {} requires at least one argument", function_name)
262+
)
263+
};
264+
265+
let mut diagnostic = Diagnostic::new_error(error_message.clone(), function_call_site);
266+
diagnostic.add_note(note, None);
267+
268+
return Err(plan_datafusion_err!("{}", error_message).with_diagnostic(diagnostic));
281269
}
282270

283271
// Helper closure to create and return an error with diagnostic information
284272
let create_error = |expected: &str, got: usize| {
285-
let base_error = plan_datafusion_err!(
273+
let error_message = format!(
286274
"Function '{}' {}, got {}",
287275
function_name,
288276
expected,
289-
got
277+
got,
290278
);
291279

292-
let mut diagnostic = Diagnostic::new_error(
293-
format!(
294-
"Wrong number of arguments for {} function call",
295-
function_name
296-
),
297-
function_call_site,
298-
);
299-
diagnostic.add_note(
300-
format!(
301-
"Function {} {}, but {} {} provided",
302-
function_name,
303-
expected,
304-
got,
305-
if got == 1 { "was" } else { "were" }
306-
),
307-
None,
308-
);
280+
let diagnostic = Diagnostic::new_error(error_message.clone(), function_call_site);
309281

310-
Err(base_error.with_diagnostic(diagnostic))
282+
Err(plan_datafusion_err!("{}", error_message).with_diagnostic(diagnostic))
311283
};
312284

313285
match signature {
314286
TypeSignature::Uniform(num, _) |
315287
TypeSignature::Numeric(num) |
316288
TypeSignature::String(num) |
317289
TypeSignature::Comparable(num) |
318-
TypeSignature::Any(num) => {
319-
if current_types.len() != *num {
320-
return create_error(&format!("expects {} arguments", num), current_types.len());
321-
}
290+
TypeSignature::Any(num) if current_types.len() != *num => {
291+
return create_error(&format!("expects {} arguments", num), current_types.len());
322292
},
323-
TypeSignature::Exact(types) => {
324-
if current_types.len() != types.len() {
325-
return create_error(&format!("expects {} arguments", types.len()), current_types.len());
326-
}
293+
// Length-based signature types
294+
TypeSignature::Exact(types) if current_types.len() != types.len() => {
295+
return create_error(&format!("expects {} arguments", types.len()), current_types.len());
327296
},
328-
TypeSignature::Coercible(types) => {
329-
if current_types.len() != types.len() {
330-
return create_error(&format!("expects {} arguments", types.len()), current_types.len());
331-
}
297+
TypeSignature::Coercible(types) if current_types.len() != types.len() => {
298+
return create_error(&format!("expects {} arguments", types.len()), current_types.len());
332299
},
333-
TypeSignature::Nullary => {
334-
if !current_types.is_empty() {
335-
return create_error("expects zero arguments", current_types.len());
336-
}
300+
301+
// Zero argument signature type
302+
TypeSignature::Nullary if !current_types.is_empty() => {
303+
return create_error("expects zero arguments", current_types.len());
337304
},
338-
TypeSignature::ArraySignature(array_signature) => {
339-
match array_signature {
340-
ArrayFunctionSignature::Array { arguments, .. } => {
341-
if current_types.len() != arguments.len() {
342-
return create_error(&format!("expects {} arguments", arguments.len()), current_types.len());
343-
}
344-
},
345-
ArrayFunctionSignature::RecursiveArray => {
346-
if current_types.len() != 1 {
347-
return create_error("expects exactly one array argument", current_types.len());
348-
}
349-
},
350-
ArrayFunctionSignature::MapArray => {
351-
if current_types.len() != 1 {
352-
return create_error("expects exactly one map argument", current_types.len());
353-
}
354-
},
355-
}
305+
306+
// Array signature types
307+
TypeSignature::ArraySignature(array_signature) => match array_signature {
308+
ArrayFunctionSignature::Array { arguments, .. } if current_types.len() != arguments.len() => {
309+
return create_error(&format!("expects {} arguments", arguments.len()), current_types.len());
310+
},
311+
ArrayFunctionSignature::RecursiveArray | ArrayFunctionSignature::MapArray if current_types.len() != 1 => {
312+
return create_error("expects exactly one array argument", current_types.len());
313+
},
314+
_ => {}
356315
},
316+
317+
// Multiple signature type
357318
TypeSignature::OneOf(signatures) => {
358-
// For OneOf, we'll consider it valid if it matches ANY of the signatures
359-
// We'll collect all errors to provide better diagnostics if nothing matches
360-
let mut all_errors = Vec::new();
361-
319+
// Try to match any signature
362320
for sig in signatures {
363-
match check_function_length_with_diag(function_name, sig, current_types, function_call_site) {
364-
Ok(()) => return Ok(()), // If any signature matches, return immediately
365-
Err(e) => all_errors.push(e),
321+
if check_function_length_with_diag(function_name, sig, current_types, function_call_site).is_ok() {
322+
return Ok(());
366323
}
367324
}
368325

369-
// none of the signatures matched
370-
if !all_errors.is_empty() {
371-
372-
let base_error = plan_datafusion_err!(
373-
"Function '{}' has no matching signature for {} arguments.",
374-
function_name,
375-
current_types.len()
376-
);
377-
378-
let mut diagnostic = Diagnostic::new_error(
379-
format!(
380-
"No matching signature for {} function with {} arguments",
381-
function_name,
382-
current_types.len()
383-
),
384-
function_call_site,
385-
);
386-
387-
diagnostic.add_note(
388-
format!("The function {} has multiple possible signatures", function_name),
389-
None,
390-
);
391-
392-
return Err(base_error.with_diagnostic(diagnostic));
393-
}
326+
// Create error for no matching signature
327+
let error_message = format!(
328+
"Function '{}' has no matching signature for {} arguments",
329+
function_name, current_types.len()
330+
);
331+
332+
let mut diagnostic = Diagnostic::new_error(error_message.clone(), function_call_site);
333+
diagnostic.add_note(
334+
format!("The function {} has multiple possible signatures", function_name),
335+
None,
336+
);
337+
338+
return Err(plan_datafusion_err!("{}", error_message).with_diagnostic(diagnostic));
394339
},
395-
// Signatures that accept variable numbers of arguments or are handled specially
396-
TypeSignature::Variadic(_) | TypeSignature::VariadicAny | TypeSignature::UserDefined => {
397-
// These cases are implicitly valid for any non-zero number of arguments:
398-
// - Variadic: accepts one or more arguments of specified types
399-
// - VariadicAny: accepts one or more arguments of any type
400-
// - UserDefined: custom validation handled by the UDF itself
401-
}
402340

341+
// All other cases:
342+
// Variadic signatures: they are specifically designed for variable argument counts, so length checking isn't necessary.
343+
// User-Defined signature: argument validation responsibility is delegated to the user-defined function implementation itself
344+
_ => {}
403345
}
404346

405347
Ok(())

‎datafusion/sql/tests/cases/diagnostic.rs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,14 @@ use datafusion_functions::string;
1919
use insta::assert_snapshot;
2020
use std::{collections::HashMap, sync::Arc};
2121

22+
use crate::{MockContextProvider, MockSessionState};
2223
use datafusion_common::{Diagnostic, Location, Result, Span};
2324
use datafusion_expr::test::function_stub::sum_udaf;
2425
use datafusion_sql::{
2526
parser::{DFParser, DFParserBuilder},
2627
planner::{ParserOptions, SqlToRel},
2728
};
2829
use regex::Regex;
29-
use sqlparser::{dialect::GenericDialect, parser::Parser};
30-
31-
use crate::{MockContextProvider, MockSessionState};
3230

3331
fn do_query(sql: &'static str) -> Diagnostic {
3432
let statement = DFParserBuilder::new(sql)

0 commit comments

Comments
 (0)