Skip to content

Commit b36a165

Browse files
osipovartemosipovartem
authored
Merge pull request #150 from Embucket/df55-udf-diagnostic-spans
Preserve SQL source spans for scalar UDF type errors
2 parents c86343f + d412f43 commit b36a165

18 files changed

Lines changed: 375 additions & 41 deletions

File tree

‎datafusion/expr/src/expr.rs‎

Lines changed: 101 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -945,12 +945,44 @@ impl Between {
945945
/// Invoke a [`ScalarUDF`] with a set of arguments
946946
///
947947
/// [`ScalarUDF`]: crate::ScalarUDF
948-
#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
948+
#[derive(Clone)]
949949
pub struct ScalarFunction {
950950
/// The function
951951
pub func: Arc<crate::ScalarUDF>,
952952
/// List of expressions to feed to the functions as arguments
953953
pub args: Vec<Expr>,
954+
/// SQL source locations of this invocation, if known.
955+
pub spans: Spans,
956+
}
957+
958+
impl PartialEq for ScalarFunction {
959+
fn eq(&self, other: &Self) -> bool {
960+
self.func == other.func && self.args == other.args
961+
}
962+
}
963+
964+
impl Eq for ScalarFunction {}
965+
966+
impl PartialOrd for ScalarFunction {
967+
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
968+
(&self.func, &self.args).partial_cmp(&(&other.func, &other.args))
969+
}
970+
}
971+
972+
impl Hash for ScalarFunction {
973+
fn hash<H: Hasher>(&self, state: &mut H) {
974+
self.func.hash(state);
975+
self.args.hash(state);
976+
}
977+
}
978+
979+
impl fmt::Debug for ScalarFunction {
980+
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
981+
f.debug_struct("ScalarFunction")
982+
.field("func", &self.func)
983+
.field("args", &self.args)
984+
.finish()
985+
}
954986
}
955987

956988
impl ScalarFunction {
@@ -965,7 +997,11 @@ impl ScalarFunction {
965997
///
966998
/// [`ScalarUDF`]: crate::ScalarUDF
967999
pub fn new_udf(udf: Arc<crate::ScalarUDF>, args: Vec<Expr>) -> Self {
968-
Self { func: udf, args }
1000+
Self {
1001+
func: udf,
1002+
args,
1003+
spans: Spans::new(),
1004+
}
9691005
}
9701006
}
9711007

@@ -2304,6 +2340,7 @@ impl Expr {
23042340
pub fn spans(&self) -> Option<&Spans> {
23052341
match self {
23062342
Expr::Column(col) => Some(&col.spans),
2343+
Expr::ScalarFunction(function) => Some(&function.spans),
23072344
Expr::Not(inner) | Expr::Negative(inner) => inner.spans(),
23082345
_ => None,
23092346
}
@@ -2508,10 +2545,12 @@ impl NormalizeEq for Expr {
25082545
Expr::ScalarFunction(ScalarFunction {
25092546
func: self_func,
25102547
args: self_args,
2548+
..
25112549
}),
25122550
Expr::ScalarFunction(ScalarFunction {
25132551
func: other_func,
25142552
args: other_args,
2553+
..
25152554
}),
25162555
) => {
25172556
self_func.name() == other_func.name()
@@ -2823,7 +2862,9 @@ impl HashNode for Expr {
28232862
| Expr::TryCast(TryCast { expr: _expr, field }) => {
28242863
field.hash(state);
28252864
}
2826-
Expr::ScalarFunction(ScalarFunction { func, args: _args }) => {
2865+
Expr::ScalarFunction(ScalarFunction {
2866+
func, args: _args, ..
2867+
}) => {
28272868
func.hash(state);
28282869
}
28292870
Expr::AggregateFunction(AggregateFunction {
@@ -3158,7 +3199,7 @@ impl Display for SchemaDisplay<'_> {
31583199
let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" };
31593200
write!(f, "{name}({})", SchemaDisplay(expr))
31603201
}
3161-
Expr::ScalarFunction(ScalarFunction { func, args }) => {
3202+
Expr::ScalarFunction(ScalarFunction { func, args, .. }) => {
31623203
match func.schema_name(args) {
31633204
Ok(name) => {
31643205
write!(f, "{name}")
@@ -4555,6 +4596,7 @@ mod test {
45554596
let udf = Expr::ScalarFunction(ScalarFunction {
45564597
func: Arc::new(ScalarUDF::new_from_impl(TestUDF {})),
45574598
args: vec![expr(), expr()],
4599+
spans: Spans::new(),
45584600
});
45594601
let Expr::ScalarFunction(scalar) = &udf else {
45604602
unreachable!()
@@ -4591,6 +4633,61 @@ mod test {
45914633
}
45924634
}
45934635

4636+
#[test]
4637+
fn scalar_function_source_spans_do_not_change_identity() {
4638+
use datafusion_common::{Location, Span};
4639+
use std::collections::hash_map::DefaultHasher;
4640+
4641+
#[derive(Debug, PartialEq, Eq, Hash)]
4642+
struct TestUdf {
4643+
signature: Signature,
4644+
}
4645+
4646+
impl ScalarUDFImpl for TestUdf {
4647+
fn name(&self) -> &str {
4648+
"test_scalar_udf"
4649+
}
4650+
4651+
fn signature(&self) -> &Signature {
4652+
&self.signature
4653+
}
4654+
4655+
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
4656+
Ok(DataType::Int32)
4657+
}
4658+
4659+
fn invoke_with_args(
4660+
&self,
4661+
_args: ScalarFunctionArgs,
4662+
) -> Result<ColumnarValue> {
4663+
unimplemented!()
4664+
}
4665+
}
4666+
4667+
let udf = Arc::new(ScalarUDF::new_from_impl(TestUdf {
4668+
signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable),
4669+
}));
4670+
let plain = ScalarFunction::new_udf(Arc::clone(&udf), vec![lit(1)]);
4671+
let mut located = ScalarFunction::new_udf(udf, vec![lit(1)]);
4672+
located.spans.add_span(Span::new(
4673+
Location { line: 1, column: 8 },
4674+
Location {
4675+
line: 1,
4676+
column: 23,
4677+
},
4678+
));
4679+
4680+
assert_eq!(plain, located);
4681+
assert_eq!(plain.partial_cmp(&located), Some(Ordering::Equal));
4682+
assert_eq!(format!("{plain:?}"), format!("{located:?}"));
4683+
let hash = |function: &ScalarFunction| {
4684+
let mut hasher = DefaultHasher::new();
4685+
function.hash(&mut hasher);
4686+
hasher.finish()
4687+
};
4688+
assert_eq!(hash(&plain), hash(&located));
4689+
}
4690+
45944691
mod intersect_metadata_tests {
45954692
use super::super::intersect_metadata_for_union;
45964693
use std::collections::HashMap;

‎datafusion/expr/src/expr_schema.rs‎

Lines changed: 92 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_K
3535
use datafusion_common::datatype::FieldExt;
3636
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
3737
use datafusion_common::{
38-
Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference,
39-
internal_err, not_impl_err, plan_datafusion_err, plan_err,
38+
Column, DataFusionError, Diagnostic, ExprSchema, Result, ScalarValue, Spans,
39+
TableReference, internal_err, not_impl_err, plan_datafusion_err, plan_err,
4040
};
4141
use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer;
4242
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
@@ -811,12 +811,25 @@ impl ExprSchemable for Expr {
811811
let new_fields = verify_function_arguments(func.as_ref(), &fields)?;
812812
func.return_field(&new_fields)
813813
}
814-
Expr::ScalarFunction(ScalarFunction { func, args }) => {
814+
Expr::ScalarFunction(ScalarFunction { func, args, spans }) => {
815815
let fields = args
816816
.iter()
817817
.map(|e| e.to_field(schema).map(|(_, f)| f))
818818
.collect::<Result<Vec<_>>>()?;
819-
let new_fields = verify_function_arguments(func.as_ref(), &fields)?;
819+
let with_span = |error: DataFusionError| {
820+
if error.diagnostic().is_some() {
821+
error
822+
} else if let Some(span) = spans.first() {
823+
error.with_diagnostic(Diagnostic::new_error(
824+
format!("Invalid function '{}'", func.name()),
825+
Some(span),
826+
))
827+
} else {
828+
error
829+
}
830+
};
831+
let new_fields = verify_function_arguments(func.as_ref(), &fields)
832+
.map_err(&with_span)?;
820833

821834
let arguments = args
822835
.iter()
@@ -830,7 +843,7 @@ impl ExprSchemable for Expr {
830843
scalar_arguments: &arguments,
831844
};
832845

833-
func.return_field_from_args(args)
846+
func.return_field_from_args(args).map_err(with_span)
834847
}
835848
// _ => Ok((self.get_type(schema)?, self.nullable(schema)?)),
836849
Expr::Cast(Cast { expr, field }) => expr
@@ -1068,6 +1081,80 @@ mod tests {
10681081
use arrow::datatypes::Schema;
10691082
use datafusion_common::{DFSchema, assert_or_internal_err};
10701083

1084+
#[test]
1085+
fn scalar_function_preserves_udf_supplied_diagnostic() {
1086+
use datafusion_common::{Location, Span};
1087+
1088+
#[derive(Debug, PartialEq, Eq, Hash)]
1089+
struct DiagnosticUdf {
1090+
signature: crate::Signature,
1091+
}
1092+
1093+
impl crate::ScalarUDFImpl for DiagnosticUdf {
1094+
fn name(&self) -> &str {
1095+
"diagnostic_udf"
1096+
}
1097+
1098+
fn signature(&self) -> &crate::Signature {
1099+
&self.signature
1100+
}
1101+
1102+
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
1103+
Ok(DataType::Int32)
1104+
}
1105+
1106+
fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
1107+
Err(
1108+
DataFusionError::Plan("bad input".to_string()).with_diagnostic(
1109+
Diagnostic::new_error(
1110+
"UDF-specific error",
1111+
Some(Span::new(
1112+
Location {
1113+
line: 1,
1114+
column: 20,
1115+
},
1116+
Location {
1117+
line: 1,
1118+
column: 21,
1119+
},
1120+
)),
1121+
),
1122+
),
1123+
)
1124+
}
1125+
1126+
fn invoke_with_args(
1127+
&self,
1128+
_args: crate::ScalarFunctionArgs,
1129+
) -> Result<crate::ColumnarValue> {
1130+
unreachable!("planning must fail before invocation")
1131+
}
1132+
}
1133+
1134+
let udf = Arc::new(crate::ScalarUDF::new_from_impl(DiagnosticUdf {
1135+
signature: crate::Signature::exact(
1136+
vec![DataType::Int32],
1137+
crate::Volatility::Immutable,
1138+
),
1139+
}));
1140+
let mut function = ScalarFunction::new_udf(udf, vec![lit(1i32)]);
1141+
function.spans.add_span(Span::new(
1142+
Location { line: 1, column: 8 },
1143+
Location {
1144+
line: 1,
1145+
column: 22,
1146+
},
1147+
));
1148+
let error = Expr::ScalarFunction(function)
1149+
.to_field(&DFSchema::empty())
1150+
.unwrap_err();
1151+
let diagnostic = error
1152+
.diagnostic()
1153+
.expect("UDF diagnostic must be preserved");
1154+
assert_eq!(diagnostic.message, "UDF-specific error");
1155+
assert_eq!(diagnostic.span.map(|span| span.start.column), Some(20));
1156+
}
1157+
10711158
macro_rules! test_is_expr_nullable {
10721159
($EXPR_TYPE:ident) => {{
10731160
let expr = lit(ScalarValue::Null).$EXPR_TYPE();

‎datafusion/expr/src/tree_node.rs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -253,11 +253,13 @@ impl TreeNode for Expr {
253253
Expr::TryCast(TryCast { expr, field }) => expr
254254
.map_elements(f)?
255255
.update_data(|be| Expr::TryCast(TryCast::new_from_field(be, field))),
256-
Expr::ScalarFunction(ScalarFunction { func, args }) => {
256+
Expr::ScalarFunction(ScalarFunction { func, args, spans }) => {
257257
args.map_elements(f)?.map_data(|new_args| {
258-
Ok(Expr::ScalarFunction(ScalarFunction::new_udf(
259-
func, new_args,
260-
)))
258+
Ok(Expr::ScalarFunction(ScalarFunction {
259+
func,
260+
args: new_args,
261+
spans,
262+
}))
261263
})?
262264
}
263265
Expr::WindowFunction(window_fun) => {

‎datafusion/functions-nested/src/array_has.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -165,7 +165,7 @@ impl ScalarUDFImpl for ArrayHas {
165165
)));
166166
}
167167
}
168-
Expr::ScalarFunction(ScalarFunction { func, args })
168+
Expr::ScalarFunction(ScalarFunction { func, args, .. })
169169
if func == &make_array_udf() =>
170170
{
171171
// make_array has a static set of arguments, so we can pull the arguments out from it

‎datafusion/functions-nested/src/extract.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1181,6 +1181,7 @@ mod tests {
11811181
Expr::Column(Column::new_unqualified("my_array")),
11821182
Expr::Column(Column::new_unqualified("my_index")),
11831183
],
1184+
spans: Default::default(),
11841185
});
11851186
assert_eq!(
11861187
ExprSchemable::get_type(&udf_expr, &schema).unwrap(),

‎datafusion/functions/src/core/getfield.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,7 @@ fn simplify_get_field_over_struct_constructor(args: &[Expr]) -> Option<Expr> {
307307
let Expr::ScalarFunction(ScalarFunction {
308308
func,
309309
args: ctor_args,
310+
..
310311
}) = base
311312
else {
312313
return None;
@@ -607,6 +608,7 @@ impl ScalarUDFImpl for GetFieldFunc {
607608
if let Expr::ScalarFunction(ScalarFunction {
608609
func,
609610
args: inner_args,
611+
..
610612
}) = current_expr
611613
&& func.inner().is::<GetFieldFunc>()
612614
{

‎datafusion/functions/src/math/log.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -372,7 +372,7 @@ impl ScalarUDFImpl for LogFunc {
372372
&info.get_data_type(&base)?,
373373
)?)))
374374
}
375-
Expr::ScalarFunction(ScalarFunction { func, mut args })
375+
Expr::ScalarFunction(ScalarFunction { func, mut args, .. })
376376
if is_pow(&func)
377377
&& args.len() == 2
378378
&& base == args[0]

‎datafusion/functions/src/math/power.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ impl ScalarUDFImpl for PowerFunc {
179179
base, &base_type,
180180
)))
181181
}
182-
Expr::ScalarFunction(ScalarFunction { func, mut args })
182+
Expr::ScalarFunction(ScalarFunction { func, mut args, .. })
183183
if is_log(&func)
184184
&& args.len() == 2
185185
&& base == args[0]

‎datafusion/functions/src/string/concat.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,7 @@ pub(crate) fn simplify_concat(args: Vec<Expr>) -> Result<ExprSimplifyResult> {
394394
ScalarFunction {
395395
func: concat(),
396396
args: new_args,
397+
spans: Default::default(),
397398
},
398399
)))
399400
} else {

‎datafusion/functions/src/string/concat_ws.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,7 @@ fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result<ExprSimplifyRes
403403
ScalarFunction {
404404
func: concat(),
405405
args: args.to_vec(),
406+
spans: Default::default(),
406407
},
407408
)))
408409
}
@@ -460,6 +461,7 @@ fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result<ExprSimplifyRes
460461
ScalarFunction {
461462
func: concat_ws(),
462463
args: new_args,
464+
spans: Default::default(),
463465
},
464466
)))
465467
}

0 commit comments

Comments
 (0)