Add Extension Type / Metadata support for Scalar UDFs #15646

timsaucer · 2025-04-08T23:29:29Z

Which issue does this PR close?

Closes Change ReturnTypeInfo to return a Field rather than DataType #14247 but it is a different solution than requested

Rationale for this change

We have many users who wish to use extension data or other metadata when writing user defined functions. This PR enables two features for Scalar UDFs

When invoking the UDF, the field of each input column will be available, if it exists
The UDF can also specify output field that is attached to the schema of the record batch

What changes are included in this PR?

This is a fairly large change, but at a high level we add in a vector of argument fields to the ScalarFunctionArgs that gets passed when invoked. Additionally all Physical Expressions are now required to implement a new function to output their field. The rest of the work included is around plumbing these changes through the system, extracting the field from the input schema, and setting it as output.

Are these changes tested?

All existing unit tests pass.
Additional unit tests exercising this feature is added in datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

…ned an owned version of the metadata for simplicity

timsaucer · 2025-04-09T12:51:28Z

Some of the concerns I have:

Canonical extension types in arrow-rs have an implementation of TryFrom<&Field> which lends towards the original Issue's suggestion of passing Field rather than just the metadata.
I need to add unit tests to make sure other operations don't throw away the metadata. I don't know if there is a one size fits all solution for which should an which should not pass through metadata. For example, alias should definitely keep the metadata IMO but what about functions that take two inputs? And all of the other single input functions - it's not clear if we need to handle this one a case by case basis or not.
If we do switch over to Field then it isn't immediately obvious to me what the name should be for the field. I need to investigate this a little more because there might already be a solution.

paleolimbot · 2025-04-09T15:36:57Z

Just a link to my experiments starting from the Expr enum in case they are useful! #15036

…handle extension types directly

timsaucer · 2025-04-09T17:54:03Z

I have updated the PR to use Field instead of metadata HashMap<String, String>. In doing so we can now use extension types directly. I've added a second unit test that uses extension types, both canonical and a user defined.

timsaucer · 2025-04-09T19:14:50Z

I think Aggregate and Window UDFs should come as a separate PR. I did notice however that for Aggregates the input portion is already viable with this PR. Since AccumulatorArgs already passes in the input physical expression and input schema we would be able to compute the input fields. I've tested this locally with success. For Window functions we will want to add in the input schema.

paleolimbot · 2025-04-09T20:34:26Z

I'll take a look this evening!

It's mostly updating to Arrow 55, but #15663 (in particular e37ef60 ) is the change required to support Extension types if arrow-rs supported them (minus reading them from files, which is admittedly a big minus).

alamb · 2025-04-10T18:23:49Z

This sounds like a good idea to me, but I suggest we wait until DF 47 is shipped to merge it

alamb

TLDR while this will be a painful downstream change for other crates, I don't really see any other way to support User Defined types / extension types and thus I think we should do make the change

I feel like this particular PR still leaves the code in an inconsistent state -- data type is still used somewhat inconsistently. If we are going to make this change I think we should fully update the API to be in terms of Field instead of DataType

FYI @rluvaton as I think this is very related to

deprecating return_type in favor of return_type_from_args #15662

Shall we just change functions to return Field? I think that would be the cleanest (though disruptive) solution

alamb · 2025-04-10T18:25:10Z

datafusion/expr/src/udf.rs

    /// The evaluated arguments to the function
    pub args: Vec<ColumnarValue>,
+    /// Field associated with each arg, if it exists
+    pub arg_fields: Vec<Option<&'b Field>>,
    /// The number of rows in record batch being evaluated
    pub number_rows: usize,


Perhaps we should change the return type to also be pub return_field: &'a Field to be consistetn

I think we should also remove ReturnTypeInfo and return_type and have a single way to return values -- output_field (or return_field?)

tobixdev · 2025-04-10T19:09:52Z

Just chiming in here. I also encountered the similar issues with UDFs not having access to (and being able to return) fields during my experiments.

Thanks for working on this 🚀! I'd really love to help out here, as I am craving for better UDTs support in DF. Unfortunately, I am a bit swamped at the moment but I'll find a few hours if there is something I can help with.

I also believe that we should update the return API to Field as well.

paleolimbot

This is awesome! I have some questions primarily from the extension type/user-defined type angle, although I know that the primary motivation here is just metadata. It does seem like a thoughtful breaking change might result in a cleaner overall final design (but I'm very new here). In particular I wonder if a structure defined in datafusion (maybe ArgType or ExprType) with a to_field() would be less confusing/more flexible than a Field whose name and nullability is mostly ignored.

paleolimbot · 2025-04-10T19:03:59Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+    fn default() -> Self {
+        Self {
+            name: "canonical_extension_udf".to_string(),
+            signature: Signature::exact(vec![DataType::Int8], Volatility::Immutable),


Would the Signature also need additional options at some point such that it would have an opportunity to consider the metadata of the arguments?

paleolimbot · 2025-04-10T19:12:17Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+        let input_field = args.arg_fields[0].unwrap();
+
+        let output_as_bool = matches!(
+            CanonicalExtensionType::try_from(input_field),


One of the things that I've worried about (but haven't benchmarked) is whether the parsing of the extension metadata (usually small amounts of JSON) will add up here (my sense is that it won't but it may be worth checking).

paleolimbot · 2025-04-10T19:36:22Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+                let array_ref = Arc::new(StringArray::from(array_values)) as ArrayRef;
+                Ok(ColumnarValue::Array(array_ref))
+            }
+            ColumnarValue::Scalar(value) => {


I wonder how extension scalars/literals would fit in here (perhaps Literal will also need a lit_field() or ScalarValue::Extension would be needed)

paleolimbot · 2025-04-10T19:38:50Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+            Field::new("canonical_extension_udf", DataType::Utf8, true)
+                .with_extension_type(MyUserExtentionType {}),


I wonder what the name and/or nullability should be in these cases (this also came up in my adventures with the Expr enum). If we don't actually need them/their value is always ignored, I wonder if a dedicated structure would be better (e.g., closer to your original version that just used a HashMap).

rluvaton · 2025-04-11T09:09:12Z

FYI @rluvaton as I think this is very related to

deprecating return_type in favor of return_type_from_args #15662

Shall we just change functions to return Field? I think that would be the cleanest (though disruptive) solution

@alamb Not sure how this is solving the problem that I said in the issue

timsaucer · 2025-04-11T18:35:16Z

I need to take some time to review these comments and think more about it, likely next week. Also I'm dropping a note for myself that the current implementation isn't sufficient for my needs because I need the UDF to be able to compute the output field based on the arguments and my function signature for output_field will not enable that.

Case in point: GetFieldFunc should return the field with meta data for the subfield of the struct, so it needs to know which field it was called for.

alamb · 2025-04-14T20:40:20Z

FYI @rluvaton as I think this is very related to

deprecating return_type in favor of return_type_from_args #15662

Shall we just change functions to return Field? I think that would be the cleanest (though disruptive) solution

@alamb Not sure how this is solving the problem that I said in the issue

What I was thinking is that if we are going to be changing ScalarFunction again so that instead of returning a DataType for return_type_from_args it would return a Field

The Field includes nullability information (as well as extension type information)

timsaucer added 6 commits April 8, 2025 15:22

Add in plumbing to pass around metadata for physical expressions

949f54b

Adding argument metadata to scalar argument struct

5912b54

Since everywhere we use this we immediately clone, go ahead and retur…

706e91b

…ned an owned version of the metadata for simplicity

Cargo fmt

a36a45f

Benchmarks required args_metadata in tests

3bd1d4f

Clippy warnings

9e548fd

github-actions bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates core Core DataFusion crate proto Related to proto crate functions Changes to functions implementation ffi Changes to the ffi crate labels Apr 8, 2025

timsaucer mentioned this pull request Apr 8, 2025

Rerun PRs rerun-io/opensource#2

Open

Switching over to passing Field around instead of metadata so we can …

0070429

…handle extension types directly

timsaucer changed the title ~~Add metadata support for Scalar UDFs~~ Add Extension Type / Metadata support for Scalar UDFs Apr 9, 2025

timsaucer self-assigned this Apr 9, 2025

timsaucer marked this pull request as ready for review April 9, 2025 19:14

alamb reviewed Apr 10, 2025

View reviewed changes

alamb mentioned this pull request Apr 10, 2025

Support round tripping extension types in parquet apache/arrow-rs#7063

Open

paleolimbot reviewed Apr 10, 2025

View reviewed changes

timsaucer marked this pull request as draft April 11, 2025 18:35

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Extension Type / Metadata support for Scalar UDFs #15646

Add Extension Type / Metadata support for Scalar UDFs #15646

timsaucer commented Apr 8, 2025 •

edited

Loading

timsaucer commented Apr 9, 2025

paleolimbot commented Apr 9, 2025

timsaucer commented Apr 9, 2025

timsaucer commented Apr 9, 2025

paleolimbot commented Apr 9, 2025

alamb commented Apr 10, 2025

alamb left a comment

alamb Apr 10, 2025

alamb Apr 10, 2025

tobixdev commented Apr 10, 2025

paleolimbot left a comment

paleolimbot Apr 10, 2025

paleolimbot Apr 10, 2025

paleolimbot Apr 10, 2025

paleolimbot Apr 10, 2025

rluvaton commented Apr 11, 2025

timsaucer commented Apr 11, 2025

alamb commented Apr 14, 2025 •

edited

Loading

		Field::new("canonical_extension_udf", DataType::Utf8, true)
		.with_extension_type(MyUserExtentionType {}),

Add Extension Type / Metadata support for Scalar UDFs #15646

Are you sure you want to change the base?

Add Extension Type / Metadata support for Scalar UDFs #15646

Conversation

timsaucer commented Apr 8, 2025 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

timsaucer commented Apr 9, 2025

paleolimbot commented Apr 9, 2025

timsaucer commented Apr 9, 2025

timsaucer commented Apr 9, 2025

paleolimbot commented Apr 9, 2025

alamb commented Apr 10, 2025

alamb left a comment

Choose a reason for hiding this comment

alamb Apr 10, 2025

Choose a reason for hiding this comment

alamb Apr 10, 2025

Choose a reason for hiding this comment

tobixdev commented Apr 10, 2025

paleolimbot left a comment

Choose a reason for hiding this comment

paleolimbot Apr 10, 2025

Choose a reason for hiding this comment

paleolimbot Apr 10, 2025

Choose a reason for hiding this comment

paleolimbot Apr 10, 2025

Choose a reason for hiding this comment

paleolimbot Apr 10, 2025

Choose a reason for hiding this comment

rluvaton commented Apr 11, 2025

timsaucer commented Apr 11, 2025

alamb commented Apr 14, 2025 • edited Loading

timsaucer commented Apr 8, 2025 •

edited

Loading

alamb commented Apr 14, 2025 •

edited

Loading