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

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

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

datafusion/expr/src/udf.rs

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.

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

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)

Interesting. I don't know. As a a work around someone could make a light UDF that returns a scalar value but with metadata. Alternatively we could add methods to the literal expression to give the ability to add metadata to it. I don't have a clear use case in my head, though, so I'm mostly spitballing here.

It is a good call we'll need to add the metadata somewhere. I think we'll have to figure out the best place as a follow on PR (we can defer until we figure out the right place)

#15797 for tracking

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

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)

…ned an owned version of the metadata for simplicity

…s a lot of unit test updates, especially because the scalar arguments pass around borrowed values

timsaucer · 2025-04-21T19:14:44Z

I'm sorry the scope of this PR keeps growing, but I've tried to address the comments and I believe this is a cleaner interface. Now all physical expressions a required to implement return_field (which I renamed from output_field to be consistent). It meant updating a lot of unit tests.

timsaucer · 2025-04-21T19:28:06Z

@alamb Please let me know if you want a re-review or to merge in after CI finishes. I did make some decent sized changes per your suggestions. A lot of what changed are unit tests, just for the way we handle borrowed Field in the invocation arguments.

paleolimbot

Thank you again for working on this!

While I was taking a look the first time you fixed my main question (the relationship between return_type and output_field) 🙂 . I did my best to look through the whole change to make sure the functions I'm hoping to express can be expressed in this way and I believe that they can!

paleolimbot · 2025-04-21T19:28:38Z

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

+    fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
+        unimplemented!(
+            "this should never be called since return_field_from_args is implemented"
+        );
+    }
+
+    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<Field> {
+        Ok(Field::new(self.name(), DataType::UInt64, true)
+            .with_metadata(self.metadata.clone()))
+    }


Thank you for this!

@paleolimbot I have an end to end example working on my company's code, so if there are still gaps from what you need please let me know!

Thanks - I'll do testing as well on our suite when this merges and see if we can remove our internal workaround 🙂

berkaysynnada · 2025-04-22T14:40:50Z

Hi @timsaucer. I'd like to review this before getting merged. If it's not blocking anything, would it be okay to wait until weekend?

timsaucer · 2025-04-22T14:54:48Z

Hi @timsaucer. I'd like to review this before getting merged. If it's not blocking anything, would it be okay to wait until weekend?

I would greatly appreciate the review, and I think it is worth waiting to get another look.

crystalxyz

The changes are impressive and it should open up more opportunities for users to customize their usage! Hopefully we don't need to make these breaking changes too often as all the projects depending on datafusion will have a hard time updating their code :)
I can imagine that we will need to update scalar udf in datafusion-python once datafusion v48 gets released, and it's exciting!

crystalxyz · 2025-04-24T12:11:48Z

docs/source/library-user-guide/upgrading.md

+repackaging of the output data into a `Field`. The name you specify on the
+field is not important. It will be overwritten during planning. `ReturnInfo`


Just wondering, it seems a little weird to me here that we ask users to provide a name value but then overwrite it. Is it considered acceptable (because it is part of arrow's Field class), or could we maybe hide this field from the user?

I think it's more work to ask them to pass in all of the other parts rather than reuse the Field

Yeah that's true. Sounds good to me

crystalxyz · 2025-04-24T12:16:12Z

datafusion/physical-expr-common/src/physical_expr.rs

+    /// The output field associated with this expression
+    fn return_field(&self, input_schema: &Schema) -> Result<Field> {
+        Ok(Field::new(
+            format!("{self}"),


If I understand correctly, this is one of the places where we overwrite the name of the field?

I believe the place it is computed is here and then used later down in the same function.

datafusion/datafusion/expr/src/expr_schema.rs

Line 375 in e4d5846

let (relation, schema_name) = self.qualified_name();

timsaucer · 2025-04-28T17:07:25Z

@berkaysynnada do you think you'll be able to review soon? I know we wanted to get this in earlier in the 48 cycle to shake out any bugs since it is a big change

berkaysynnada · 2025-04-29T07:21:07Z

@berkaysynnada do you think you'll be able to review soon? I know we wanted to get this in earlier in the 48 cycle to shake out any bugs since it is a big change

Sorry, I wasn't able to finish it yet, and it doesn't look like I can do in the next few days either. I don't want to block progress, so I'll do a post-merge review.

timsaucer · 2025-04-29T13:33:05Z

Thank you, everyone, for the thoughtful discussions and reviews.

alamb · 2025-04-29T20:54:58Z

The march to supporting user defined types has begun!

* Add in plumbing to pass around metadata for physical expressions * Adding argument metadata to scalar argument struct * Since everywhere we use this we immediately clone, go ahead and returned an owned version of the metadata for simplicity * Cargo fmt * Benchmarks required args_metadata in tests * Clippy warnings * Switching over to passing Field around instead of metadata so we can handle extension types directly * Switching return_type_from_args to return_field_from_args * Updates to unit tests for switching to field instead of data_type * Resolve unit test issues * Update after rebase on main * GetFieldFunc should return the field it finds instead of creating a new one * Get metadata from scalar functions * Change expr_schema to use to_field primarily instead of individual calls for getting data type, nullability, and schema * Scalar function arguments should take return field instead of return data type now * subquery should just get the field from below and not lose potential metadata * Update comment * Remove output_field now that we've determined it using return_field_from_args * Change name to_field to field_from_column to be more consistent with the usage and prevent misconception about if we are doing some conversion * Minor moving around of the explicit lifetimes in the struct definition * Change physical expression to require to output a field which requires a lot of unit test updates, especially because the scalar arguments pass around borrowed values * Change name from output_field to return_field to be more consistent * Update migration guide for DF48 with user defined functions * Whitespace * Docstring correction

Working `ST_AsBinary` even persisting CRS! Enabled by apache/datafusion#15646 ### Change list - Restore `geodatafusion` crate and restore `datafusion` dependency (pinned to upstream main) - Add `with_metadata` to all concrete arrays and to the `GeoArrowArray` trait so that we can replace the metadata (e.g. CRS) of an existing array. - Remove `metadata` field from `MixedGeometryArray` because the inner array of the GeometryCollectionArray doesn't need to store GeoArrow metadata.

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

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

datafusion/expr/src/udf.rs Show resolved Hide resolved

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

timsaucer added 6 commits April 16, 2025 07:43

Add in plumbing to pass around metadata for physical expressions

39d1a72

Adding argument metadata to scalar argument struct

9bcd4b5

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

ea561b0

…ned an owned version of the metadata for simplicity

Cargo fmt

8daa356

Benchmarks required args_metadata in tests

a2d5f9e

Clippy warnings

a3514de

timsaucer added 4 commits April 21, 2025 09:11

Minor moving around of the explicit lifetimes in the struct definition

c439224

Change physical expression to require to output a field which require…

d69deee

…s a lot of unit test updates, especially because the scalar arguments pass around borrowed values

Change name from output_field to return_field to be more consistent

317ed22

Update migration guide for DF48 with user defined functions

6481b3e

github-actions bot added the documentation Improvements or additions to documentation label Apr 21, 2025

This was referenced Apr 21, 2025

Support metadata on literal values #15797

Open

Deprecate ExprSchemable functions #15798

Open

Merge branch 'main' into feat/physical_expr_evaluate_metadata

ef6c58c

Whitespace

cddb52a

paleolimbot reviewed Apr 21, 2025

View reviewed changes

Docstring correction

e4d5846

crystalxyz approved these changes Apr 24, 2025

View reviewed changes

timsaucer merged commit 2d80194 into apache:main Apr 29, 2025
31 checks passed

timsaucer deleted the feat/physical_expr_evaluate_metadata branch April 29, 2025 13:33

timsaucer mentioned this pull request Apr 30, 2025

Add metadata support for Aggregate and Window Functions #15902

Open

kylebarron mentioned this pull request Apr 30, 2025

Prototype Datafusion scalar UDFs with new upstream support geoarrow/geoarrow-rs#1091

Closed

timsaucer mentioned this pull request May 1, 2025

feat: metadata handling for aggregates and window functions #15911

Open

3 tasks

andygrove mentioned this pull request May 2, 2025

Investigate impact of breaking changes in DataFusion since 47.0.0 release apache/datafusion-comet#1709

Closed

kylebarron mentioned this pull request May 6, 2025

DataFusion extension types working proof of concept! geoarrow/geoarrow-rs#1093

Merged

friendlymatthew mentioned this pull request May 8, 2025

DNM: Point dependency to the _latest_ version datafusion-contrib/datafusion-functions-json#84

Open

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

tobixdev commented Apr 10, 2025

paleolimbot left a comment

paleolimbot Apr 10, 2025

timsaucer Apr 18, 2025

alamb Apr 19, 2025

timsaucer Apr 21, 2025

rluvaton commented Apr 11, 2025

timsaucer commented Apr 11, 2025

alamb commented Apr 14, 2025 •

edited

Loading

timsaucer commented Apr 21, 2025

timsaucer commented Apr 21, 2025

paleolimbot left a comment

paleolimbot Apr 21, 2025

timsaucer Apr 21, 2025

paleolimbot Apr 22, 2025

berkaysynnada commented Apr 22, 2025

timsaucer commented Apr 22, 2025

crystalxyz left a comment

crystalxyz Apr 24, 2025

timsaucer Apr 24, 2025

crystalxyz Apr 25, 2025

crystalxyz Apr 24, 2025

timsaucer Apr 24, 2025 •

edited

Loading

timsaucer commented Apr 28, 2025

berkaysynnada commented Apr 29, 2025

timsaucer commented Apr 29, 2025

alamb commented Apr 29, 2025

		repackaging of the output data into a `Field`. The name you specify on the
		field is not important. It will be overwritten during planning. `ReturnInfo`

Add Extension Type / Metadata support for Scalar UDFs #15646

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

tobixdev commented Apr 10, 2025

paleolimbot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

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 21, 2025

timsaucer commented Apr 21, 2025

paleolimbot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

berkaysynnada commented Apr 22, 2025

timsaucer commented Apr 22, 2025

crystalxyz left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

timsaucer Apr 24, 2025 • edited Loading

Choose a reason for hiding this comment

timsaucer commented Apr 28, 2025

berkaysynnada commented Apr 29, 2025

timsaucer commented Apr 29, 2025

alamb commented Apr 29, 2025

timsaucer commented Apr 8, 2025 •

edited

Loading

alamb commented Apr 14, 2025 •

edited

Loading

timsaucer Apr 24, 2025 •

edited

Loading