UDAF: Add more fields to state fields #10391

jayzhan211 · 2024-05-06T12:20:26Z

Which issue does this PR close?

Closes #.

I check the field and state_fields that are used for most of the aggregate functions. I think it is nice to split them out as a PR

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 · 2024-05-06T12:21:28Z

datafusion/physical-expr-common/src/aggregate/mod.rs

@@ -152,16 +157,21 @@ pub trait AggregateExpr: Send + Sync + Debug + PartialEq<dyn Any> {
 pub struct AggregateFunctionExpr {
    fun: AggregateUDF,
    args: Vec<Arc<dyn PhysicalExpr>>,
+    /// input type
+    input_type: DataType,


I found most of the builtin function name is input_data_type so I rename it.

jayzhan211 · 2024-05-06T12:22:17Z

datafusion/expr/src/udaf.rs

+    ///
+    /// [`format_state_name`]: crate::utils::format_state_name
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+        not_impl_err!("state_fields hasn't been implemented for {self:?} yet")


I found most of the state_fields have their own version, so return err if not defined now

This would be an API change -- I think if we want to change it we should also update the documentation to reflect the change

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb

Thanks @jayzhan211 -- I almost feel like this PR is a solution in search of a problem. Specifically, is there a practical problem that the current API can't handle that this could ?

If this doesn't add any value to users but costs them (in terms of API churn) I am not sure it makes sense to make this change

alamb · 2024-05-06T19:13:39Z

datafusion/expr/src/udaf.rs

@@ -287,18 +286,23 @@ pub trait AggregateUDFImpl: Debug + Send + Sync {
    /// aggregate function was called.
    fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>>;

+    /// Return the fields for the function
+    fn field(&self, _args: FieldArgs) -> Result<Field> {


Is this the same information as return_type? If so, perhaps we should deprecate the return_type function?

return_type is part of the information in field and state_field not replaceable to field

alamb · 2024-05-06T19:14:26Z

datafusion/expr/src/udaf.rs

+    ///
+    /// [`format_state_name`]: crate::utils::format_state_name
+    fn state_fields(&self, _args: StateFieldsArgs) -> Result<Vec<Field>> {
+        not_impl_err!("state_fields hasn't been implemented for {self:?} yet")


This would be an API change -- I think if we want to change it we should also update the documentation to reflect the change

alamb · 2024-05-06T19:15:10Z

datafusion/functions-aggregate/src/covariance.rs

+            Field::new(
+                format_state_name(args.name, "mean1"),
+                DataType::Float64,
+                args.nullable,


aren't some of these fields nullable even if their input is not nullable (e.g. mean with no inputs is null)?

yes, this should be true

alamb · 2024-05-06T19:16:12Z

datafusion/expr/src/function.rs

+    pub nullable: bool,
+}
+
+impl<'a> FieldArgs<'a> {


Given all the fields are pub, I wonder how much benefit this API adds over simply creating FieldArgs directly 🤔

But like name, it includes arguments, and not every time is the same, how do they get the correct field without args.name?

jayzhan211 · 2024-05-07T00:24:48Z

@alamb This change is based on what the built-in function does, and what we need for UDAF to have the same behavior. If the change is not clear if it is necessary, we can also change it when moving the function to UDAF gradually.

This PR exists because 1. breaking API changes at once, so we don't need to keep breaking API after that. 2. Split the change out can also reduce the review size of other functions moving (but less straightforward if the change is worth it)

if we found out there is a better way to rewrite what the current builtin function does, the change here might be useless, so I think both way has pros and cons

jayzhan211 · 2024-05-07T00:51:58Z

StateFieldsArgs
DistinctArrayAgg needs nullable
Sum needs return_type
Most of the functions needs input_type.

FieldArgs
~~Avg needs return_type~~
~~Most of the functions needs input_type.~~
OrderSensitiveArrayAgg needs nullable and order_by_data_types

Upd:
I was mislead by the naming, it seems we just need the value_type what we have now, which most of the functions are input_type, while functions like Sum or Count has a different type (return_type).

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

alamb · 2024-05-08T13:06:48Z

@alamb This change is based on what the built-in function does, and what we need for UDAF to have the same behavior. If the change is not clear if it is necessary, we can also change it when moving the function to UDAF gradually.

Ah, interesting, I didn't realize this. Sorry

I think it would be useful to port UDFs over one by one, and as we find gaps in the APIs we can adjust them

Or if we want to discover all the potential issues we could do a WIP PR (as you have done already) to try to move the aggregates and when we hit a block fix the underlying issues in a different PR

add more fields to state fields

4c243c2

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

github-actions bot added the logical-expr Logical plan and expressions label May 6, 2024

jayzhan211 commented May 6, 2024

View reviewed changes

jayzhan211 added the api change Changes the API exposed to users of the crate label May 6, 2024

add fields for udfimpl

d252ff1

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

github-actions bot added the core Core datafusion crate label May 6, 2024

jayzhan211 added 3 commits May 6, 2024 21:55

sue return type

49e9a2a

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fix partialeq

994ecb8

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

fix doctest

b840a50

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 marked this pull request as ready for review May 6, 2024 14:49

alamb mentioned this pull request May 6, 2024

DataFusion weekly project plan (Andrew Lamb) - May 6, 2024 #10395

Closed

7 tasks

alamb reviewed May 6, 2024

View reviewed changes

jayzhan211 added 2 commits May 7, 2024 08:53

fix nullable in covar

5eb43ea

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

upd doc

f8310ca

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

jayzhan211 marked this pull request as draft May 8, 2024 05:22

jayzhan211 mentioned this pull request May 15, 2024

UDAF: Extend more args to state_fields and groups_accumulator_supported and introduce ReversedUDAF #10525

Merged

jayzhan211 closed this in #10525 May 16, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

UDAF: Add more fields to state fields #10391

UDAF: Add more fields to state fields #10391

jayzhan211 commented May 6, 2024 •

edited

jayzhan211 May 6, 2024

jayzhan211 May 6, 2024 •

edited

alamb May 6, 2024

alamb left a comment

alamb May 6, 2024

jayzhan211 May 7, 2024

alamb May 6, 2024

alamb May 6, 2024

jayzhan211 May 7, 2024

alamb May 6, 2024

jayzhan211 May 7, 2024 •

edited

jayzhan211 commented May 7, 2024

jayzhan211 commented May 7, 2024 •

edited

alamb commented May 8, 2024

UDAF: Add more fields to state fields #10391

UDAF: Add more fields to state fields #10391

Conversation

jayzhan211 commented May 6, 2024 • edited

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

jayzhan211 May 6, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alamb 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jayzhan211 May 7, 2024 • edited

Choose a reason for hiding this comment

jayzhan211 commented May 7, 2024

jayzhan211 commented May 7, 2024 • edited

alamb commented May 8, 2024

jayzhan211 commented May 6, 2024 •

edited

jayzhan211 May 6, 2024 •

edited

jayzhan211 May 7, 2024 •

edited

jayzhan211 commented May 7, 2024 •

edited