Allow UDFs to return custom Diagnostic

Is your feature request related to a problem or challenge?

When implementing ScalarUDFImpl, AggregateUDFImpl, or WindowUDFImpl, it would be very cool to have the chance to return a custom Diagnostic either during the planning phase or the execution phase. Currently, you can return errors e.g. in ScalarUDFImpl::return_type_from_args or ScalarUDFImpl::invoke, but no Spans are available there so there's no way to construct a Diagnostic that highlights the call expression, or a specific argument of the function, for instance.

Describe the solution you'd like

Add a FnCallSpans types that bundles:

The Span for the whole call expression
The Span for just the function identifier
A Vec<Span>, one for each argument

Add a diagnose trait function (defaults to doing nothing) which is called during logical planning and is given the types of the arguments for a particular callsite, as well as an instance of FnCallSpans. The implementor can here construct a Diagnostic and return an errors that contains it. The logical planner would bubble up the error.

Also add a FnCallSpans field to ScalarFunctionArgs (and equivalently for aggregate and window functions) so the implementor, in ScalarUDFImpl::invoke_with_args, can create Diagnostics for a particular execution as well. This shouldn't be too big of a breaking change, unless users are using struct destruction syntax with ScalarFunctionArgs.

Describe alternatives you've considered

No response

Additional context

No response

@alamb

What do you think @alamb? I could take on this and make a PR if you approve the solution :)

@eliaperantoni

Sounds like a good idea to me -- thanks @eliaperantoni

I think it would be most helpful to:

Use the new API in at least one of the built in functions so 1) the API is tested and 2) it provides value to DataFusion users

As in I think it would be less compelling if it was just an API with no use other than tests

@alamb

I think it would be most helpful to use the new API in at least one of the built in functions so 1) the API is tested and 2) it provides value to DataFusion users

That makes total sense, thank you @alamb!

Hey I've taken a cursory look at this ticket. So if I'm thinking about this correctly, the FnCallSpan implementation will build on top of the existing Span implementation and be something like:

pub struct FnCallSpans {
    pub fn_call: Span,   // Span covering the entire function call
    pub fn_name: Span,     // Span for just the function identifier
    pub args: Vec<Span>,   // Spans for each argument
}

Then, adding the diagnose treat would be something like:

pub trait ScalarUDFImpl: Debug + Send + Sync {
    fn diagnose(
        &self,
        arg_types: &[DataType],  // Argument types for this callsite
        spans: &FnCallSpans,     // The function call spans
    ) -> Result<(), Diagnostic> {
        // Default implementation: do nothing
        Ok(())
    }
}

And then the final part is essentially:

pub struct ScalarFunctionArgs<'a> {
    ...

    /// Spans for function call diagnostics (new field)
    pub spans: &'a FnCallSpans,
}

This will need to be done for the other UDFs too. So this essentially sets the groundwork so that a custom Diagnostic can be created for the UDFs, but does not actually implement them correct? I'm sure I'm missing a few things here, but hopefully I'm on the right track. Thanks for your help!

@jsai28

Hey @jsai28, that looks so good! ❤️

I have some points that I'd like to hear your opinion on:

I think FnCallSpans.args should have exactly as many elements as the arguments of the function. i.e. you should always be able to do FnCallSpans[i] for every arg_types[i]. What would then happen if a Span is not available for the i-th argument? For example, until sqlparser is updated, Spans are not available for literal values. So you wouldn't have a span for the 'foo' in concat(var1, 'foo', var2) I think perhaps we should use Vec<Option<Span>> instead?
What are the pros and cons of using Span instead of Spans (note the "s")? I think Spans is able to deal with a lack of Spans out of the box, so you would do args: Vec<Spans> and then to get the span for the i-th argument you'd do args[i].first(). Also Spans supports multiple Spans, not that we necessarily need it for function arguments, but who knows, maybe in the future. I also have this feeling that we should always use unless there's a good argument otherwise because it was created to be the global Spans container for DataFusion and will evolve to acomodate any future needs. What do you think?
I'm thinking maybe ScaladUDFImpl::diagnose should take arg_types: ReturnTypeArgs instead of arg_types: &[DataType] to be more general and futureproof. Just like ScalarUDFImpl::return_type_from_args is more general than ScalarUDFImpl::return_type.

So this essentially sets the groundwork so that a custom Diagnostic can be created for the UDFs, but does not actually implement them correct?

I'm not sure what you mean with "but does not actually implement them correct". @alamb suggested using this new feature by creating a few Diagnostic in some of the built-in functions you see in datafusion::functions* so that this API is tested and provides value :)

@alamb

Regarding your first two points, I do think that Vec<Spans> may be the way to do this. Mainly as it supports handling the case of literal values out of the box. If sqlparser is eventually updated to support literals, then wrapping the Span with Option would have been unnecessary extra steps. Also, if in the future we need multiple Spans, then we would probably need to make significant changes to this API with the Vec<Option<Span>> implementation. So that's why I'm leaning towards Vec<Spans>. I'm also thinking that if we wanted more sophisticated span handling in the future (like maybe the ability to prioritize certain errors?), then Vec<Spans> might make it easier. Although I haven't fully thought about how that could work.

For your third point, yes return_type_from_args makes more sense.

@alamb suggested using this new feature by creating a few Diagnostic in some of the built-in functions you see in datafusion::functions* so that this API is tested and provides value

Ah I see what you mean. So we will first implement Diagnostic for a few functions in datafusion::functions*, see how it works, and then eventually create another ticket to add it to the rest maybe?

@jsai28

@jsai28 Precisely. I agree with everything you said :)

@eliaperantoni

@eliaperantoni Awesome! Do you think there is enough here to make it into a GSoC project? Or do you think this could be one part of a project that includes adding diagnostic info to other areas within DataFusion as well?

@eliaperantoni Awesome! Do you think there is enough here to make it into a GSoC project? Or do you think this could be one part of a project that includes adding diagnostic info to other areas within DataFusion as well?

I think it's good but maybe there's space for a little more :) It would be cool if it was related to Diagnostic but it doesn't have to

Maybe some blog posts / documentation / tutorials on how Diagnostic works? Mind if I come up with a draft proposal and send to you (via email or something else?) ?

Maybe some blog posts / documentation / tutorials on how Diagnostic works? Mind if I come up with a draft proposal and send to you (via email or something else?) ?

That would be great!

If you want to make a blog, I think we could host it easily on the https://datafusion.apache.org/blog/

(make a PR to https://github.com/alamb/datafusion-site)

take

@eliaperantoni

@eliaperantoni I had a follow up question to this regarding the diagnose trait function. We want to call the diagnose function during logical planning and pass in ReturnTypeArgs and FnCallSpans so that the implementor can create a Diagnostic. In #14430, the Diagnostic is attached to an existing error (via the with_diagnostic method) and is then bubbled up. Should the diagnose trait function act similarly here? For example, would the diagnose function be called here for a Scalar UDF? And inside the diagnose function is where a Scalar UDF implementation would add notes to enrich the error message? Since we're passing in ReturnTypeArgs and FnCallSpans we can use ReturnTypeArgs to determine the erroneous argument and FnCallSpans to highlight where it is located. Is that right?

@jsai28

@eliaperantoni I had a follow up question to this regarding the diagnose trait function. We want to call the diagnose function during logical planning and pass in ReturnTypeArgs and FnCallSpans so that the implementor can create a Diagnostic. In #14430, the Diagnostic is attached to an existing error (via the with_diagnostic method) and is then bubbled up. Should the diagnose trait function act similarly here? For example, would the diagnose function be called here for a Scalar UDF? And inside the diagnose function is where a Scalar UDF implementation would add notes to enrich the error message? Since we're passing in ReturnTypeArgs and FnCallSpans we can use ReturnTypeArgs to determine the erroneous argument and FnCallSpans to highlight where it is located. Is that right?

@jsai28 I'm extremely sorry for the delay 🙏 I've been very busy with other things. But everything you said is absolutely correct. The only think I'd point out is that I think it would be the diagnose function itself that creates an error, because otherwise there's none. And then diagnose would attach one or more Diagnostic to that error.

This means that the programmer of a custom UDF has to build and return an error of some kind, even though they might just be interested in creating Diagnostics since that's what ends up being seen by the user.

Perhaps we could change the return type of diagnose to be Vec<Diagnostic>, and if a non empty vec is returned, then we create a standard error (eg. DataFusionError::Plan("udf is invalid")) and attach the diagnostics to it? What do you think?

@eliaperantoni

@eliaperantoni No worries! I had some follow-ups.

I think it would be the diagnose function itself that creates an error, because otherwise there's none

Isn't there already a DataFusion error if incorrect arguments are supplied to a UDF? This one. My understanding is that the data_types_with_scalar_udf function is called during the logical planning of a Scalar UDF and checks whether the right arguments were supplied (or whether they can be coerced). My thinking was that the diagnose trait function would be called when there is an error and supply more detailed diagnostic information.

Or is the diagnose trait function for errors that go beyond just incorrect function arguments and are more specific to the UDF? Maybe like a UDF that for example doesn't accept numbers larger than 50, then the diagnose function is where the implementor can produce that error and highlight the exact argument? If this is the case, then yes I agree with your other point that we should default to a standard error rather than have the implementor create their own DataFusionError. They should just return a Vec<Diagnostic which we can consume to make a DataFusionError for them, and attach the relevant diagnostics to it.

Navigation Menu

Repository navigation

Allow UDFs to return custom `Diagnostic` #15276

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Allow UDFs to return custom Diagnostic #15276

Description

Is your feature request related to a problem or challenge?

Describe the solution you'd like

Describe alternatives you've considered

Additional context

Activity

eliaperantoni commented on Mar 17, 2025

alamb commented on Mar 18, 2025

eliaperantoni commented on Mar 19, 2025

jsai28 commented on Mar 20, 2025

eliaperantoni commented on Mar 20, 2025

jsai28 commented on Mar 21, 2025

eliaperantoni commented on Mar 21, 2025

jsai28 commented on Mar 21, 2025

eliaperantoni commented on Mar 21, 2025

jsai28 commented on Mar 21, 2025

alamb commented on Mar 21, 2025

jsai28 commented on Mar 24, 2025

jsai28 commented on Mar 31, 2025

eliaperantoni commented on Apr 22, 2025

jsai28 commented on Apr 26, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Allow UDFs to return custom `Diagnostic` #15276