Move array `ArrayAgg` to a `UserDefinedAggregate` #11448

jayzhan211 · 2024-07-13T07:37:00Z

Which issue does this PR close?

Closes #.
Part of #10999, #8708

TODO

OrderSensitiveArrayAgg

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

This follows how it done for input_type and only provide a single value. But might need to be changed into a Vec in the future. This is need when we are moving `arrag_agg` to udaf where one of the states nullability will depend on the nullability of the input.

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 · 2024-07-17T06:12:08Z

datafusion/physical-expr/src/aggregate/array_agg_distinct.rs

-}
-
-#[cfg(test)]
-mod tests {


The test for distinct array agg is already coverdd in

datafusion/datafusion/sqllogictest/test_files/aggregate.slt

Lines 173 to 192 in d67b0fb

# test array_agg_distinct with list data type

statement ok

CREATE TABLE array_agg_distinct_list_table AS VALUES

('w', [0,1]),

('w', [0,1]),

('w', [1,0]),

('b', [1,0]),

('b', [1,0]),

('b', [1,0]),

('b', [0,1])

;

# Apply array_sort to have deterministic result, higher dimension nested array also works but not for array sort,

# so they are covered in `datafusion/physical-expr/src/aggregate/array_agg_distinct.rs`

query ??

select array_sort(c1), array_sort(c2) from (

select array_agg(distinct column1) as c1, array_agg(distinct column2) as c2 from array_agg_distinct_list_table

);

----

[b, w] [[0, 1], [1, 0]]

Remove the test here doesn't reduce test coverage

(though it looks like the comments should be adjusted):

datafusion/datafusion/sqllogictest/test_files/aggregate.slt

Line 186 in d67b0fb

# so they are covered in `datafusion/physical-expr/src/aggregate/array_agg_distinct.rs`

jayzhan211 · 2024-07-17T06:15:08Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -702,6 +702,7 @@ async fn roundtrip_expr_api() -> Result<()> {
        string_agg(col("a").cast_to(&DataType::Utf8, &schema)?, lit("|")),
        bool_and(lit(true)),
        bool_or(lit(true)),
+        array_agg(lit(1)),


I didn't add array_agg_distinct since it is not exist, and we can get distinct one with builder mode.

What is builder mode ? You mean a builder for array_agg

I think @jayzhan211 means that there is a single array_agg udf function now, and internally it has special accumulators for distinct and non distinct

Thus there is no difference in how to serialize it I think

Maybe we could use the builder to test

array_agg(lit(1)).distinct().build().unwrap(),

🤔

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 · 2024-07-17T07:21:13Z

datafusion/physical-expr/src/expressions/mod.rs

@@ -63,160 +61,3 @@ pub use negative::{negative, NegativeExpr};
 pub use no_op::NoOp;
 pub use not::{not, NotExpr};
 pub use try_cast::{try_cast, TryCastExpr};
-
-#[cfg(test)]
-pub(crate) mod tests {


dharanad · 2024-07-18T17:18:30Z

datafusion/functions-aggregate/src/array_agg.rs

+
+    // TODO: change name to lowercase
+    fn name(&self) -> &str {
+        "ARRAY_AGG"


Any specific reason for having this in uppercase ?

Interestingly it is inconsistent with count which uses lower case:

datafusion/datafusion/functions-aggregate/src/count.rs

Lines 112 to 113 in b685e2d

fn name(&self) -> &str {

"count"

I'm not sure why they are upper case initially, however we plan to convert them all to lowercase

dharanad · 2024-07-18T17:20:29Z

datafusion/functions-aggregate/src/array_agg.rs

+        if values.is_empty() {
+            return Ok(());
+        }
+        assert!(values.len() == 1, "array_agg can only take 1 param!");


This will invoke the panic!, should we return an Err here ?

dharanad · 2024-07-18T17:20:54Z

datafusion/functions-aggregate/src/array_agg.rs

+        if states.is_empty() {
+            return Ok(());
+        }
+        assert!(states.len() == 1, "array_agg states must be singleton!");


Same as above. Can we remove all assert! and return Err

dharanad · 2024-07-18T17:23:28Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -702,6 +702,7 @@ async fn roundtrip_expr_api() -> Result<()> {
        string_agg(col("a").cast_to(&DataType::Utf8, &schema)?, lit("|")),
        bool_and(lit(true)),
        bool_or(lit(true)),
+        array_agg(lit(1)),


What is builder mode ? You mean a builder for array_agg

dharanad · 2024-07-18T17:28:17Z

datafusion/proto/src/physical_plan/to_proto.rs

-        distinct = true;
-        protobuf::AggregateFunction::ArrayAgg
-    } else if aggr_expr.downcast_ref::<OrderSensitiveArrayAgg>().is_some() {
+    // TODO: remove


Nit: I think it can be removed.

OrderSensitiveArrayAgg is not yet removed

dharanad · 2024-07-18T17:29:38Z

datafusion/functions-aggregate/src/array_agg.rs

+    fn default() -> Self {
+        Self {
+            signature: Signature::any(1, Volatility::Immutable),
+            alias: vec!["array_agg".to_string()],


Why are we using the smallcase as alias ? Shouldn't this be the function name itself ?

dharanad · 2024-07-18T17:31:23Z

datafusion/functions-aggregate/src/array_agg.rs

+    ArrayAgg,
+    array_agg,
+    expression,
+    "input values, including nulls, concatenated into an array",


I think this is not clear & can we improve for clarity just the other functions ?

dharanad · 2024-07-18T17:31:46Z

datafusion/expr/src/expr_fn.rs

@@ -171,6 +171,7 @@ pub fn max(expr: Expr) -> Expr {
    ))
 }

+// TODO: remove


Nit: I think we can remove this

dharanad · 2024-07-18T17:36:42Z

datafusion/functions-aggregate/src/array_agg.rs

+        Ok(())
+    }
+
+    // Append value like ListArray(Int64Array(1,2,3), Int64Array(4,5,6))


Nit: If this is code document, then we can move this inside the function.

alamb

Thanks @jayzhan211 and @dharanad for the review

I think this PR is good enough to merge, though I think addressing @dharanad 's comments beforehand would make it even better.

We are so close to completing user defined aggregate migration. so close!

alamb · 2024-07-18T19:50:49Z

datafusion/functions-aggregate/src/array_agg.rs

+// specific language governing permissions and limitations
+// under the License.
+
+//! Defines physical expressions that can evaluated at runtime during query execution


Suggested change

//! Defines physical expressions that can evaluated at runtime during query execution

//! `ARRAY_AGG` aggregate implementation: [`ArrayAgg`]

alamb · 2024-07-18T19:52:15Z

datafusion/functions-aggregate/src/array_agg.rs

+
+    // TODO: change name to lowercase
+    fn name(&self) -> &str {
+        "ARRAY_AGG"


Interestingly it is inconsistent with count which uses lower case:

datafusion/datafusion/functions-aggregate/src/count.rs

Lines 112 to 113 in b685e2d

fn name(&self) -> &str {

"count"

alamb · 2024-07-18T19:56:13Z

datafusion/physical-expr/src/aggregate/array_agg_distinct.rs

-}
-
-#[cfg(test)]
-mod tests {


(though it looks like the comments should be adjusted):

datafusion/datafusion/sqllogictest/test_files/aggregate.slt

Line 186 in d67b0fb

# so they are covered in `datafusion/physical-expr/src/aggregate/array_agg_distinct.rs`

alamb · 2024-07-18T19:58:07Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -702,6 +702,7 @@ async fn roundtrip_expr_api() -> Result<()> {
        string_agg(col("a").cast_to(&DataType::Utf8, &schema)?, lit("|")),
        bool_and(lit(true)),
        bool_or(lit(true)),
+        array_agg(lit(1)),


I think @jayzhan211 means that there is a single array_agg udf function now, and internally it has special accumulators for distinct and non distinct

Thus there is no difference in how to serialize it I think

Maybe we could use the builder to test

array_agg(lit(1)).distinct().build().unwrap(),

🤔

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 · 2024-07-19T03:58:55Z

Thanks @dharanad and @alamb

alamb · 2024-07-19T21:36:56Z

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -702,6 +702,8 @@ async fn roundtrip_expr_api() -> Result<()> {
        string_agg(col("a").cast_to(&DataType::Utf8, &schema)?, lit("|")),
        bool_and(lit(true)),
        bool_or(lit(true)),
+        array_agg(lit(1)),
+        array_agg(lit(1)).distinct().build().unwrap(),


* Add input_nullable to UDAF args StateField/AccumulatorArgs This follows how it done for input_type and only provide a single value. But might need to be changed into a Vec in the future. This is need when we are moving `arrag_agg` to udaf where one of the states nullability will depend on the nullability of the input. * Make ArragAgg (not ordered or distinct) into a UDAF * Add roundtrip_expr_api test case * Address PR comments * Propegate input nullability for aggregates * Remove from accumulator args * first draft Signed-off-by: jayzhan211 <[email protected]> * cleanup Signed-off-by: jayzhan211 <[email protected]> * fix test Signed-off-by: jayzhan211 <[email protected]> * distinct Signed-off-by: jayzhan211 <[email protected]> * fix Signed-off-by: jayzhan211 <[email protected]> * address comment Signed-off-by: jayzhan211 <[email protected]> --------- Signed-off-by: jayzhan211 <[email protected]> Co-authored-by: Emil Ejbyfeldt <[email protected]>

github-actions bot added logical-expr Logical plan and expressions physical-expr Physical Expressions core Core DataFusion crate labels Jul 13, 2024

jayzhan211 force-pushed the array-agg-udf branch from c7f10e1 to a38089a Compare July 13, 2024 07:40

jayzhan211 changed the title ~~Convert ArrayAgg to UDAF~~ ArrayAgg UDAF Jul 13, 2024

jayzhan211 mentioned this pull request Jul 15, 2024

Remove element's nullability of array_agg function #11447

Merged

eejbyfeldt and others added 9 commits July 17, 2024 13:35

Make ArragAgg (not ordered or distinct) into a UDAF

d1634bc

Add roundtrip_expr_api test case

dc712f4

Address PR comments

41b9c8a

Propegate input nullability for aggregates

a7712e7

Remove from accumulator args

d5d21b0

first draft

7a1056c

Signed-off-by: jayzhan211 <[email protected]>

cleanup

5ddaddb

Signed-off-by: jayzhan211 <[email protected]>

fix test

9d17c1c

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 force-pushed the array-agg-udf branch from 7088d00 to 9d17c1c Compare July 17, 2024 05:38

jayzhan211 commented Jul 17, 2024

View reviewed changes

distinct

ac640fb

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 force-pushed the array-agg-udf branch from d7fd244 to ac640fb Compare July 17, 2024 06:15

fix

f67d8f7

Signed-off-by: jayzhan211 <[email protected]>

jayzhan211 commented Jul 17, 2024

View reviewed changes

jayzhan211 marked this pull request as ready for review July 17, 2024 08:12

jayzhan211 requested a review from alamb July 18, 2024 07:36

dharanad reviewed Jul 18, 2024

View reviewed changes

alamb changed the title ~~ArrayAgg UDAF~~ Move array ArrayAgg to a UserDefinedAggregate Jul 18, 2024

alamb approved these changes Jul 18, 2024

View reviewed changes

address comment

a9b925b

Signed-off-by: jayzhan211 <[email protected]>

github-actions bot added the sqllogictest SQL Logic Tests (.slt) label Jul 19, 2024

jayzhan211 merged commit 12d82c4 into apache:main Jul 19, 2024
23 checks passed

jayzhan211 deleted the array-agg-udf branch July 19, 2024 03:59

alamb reviewed Jul 19, 2024

View reviewed changes

eejbyfeldt mentioned this pull request Aug 6, 2024

Convert ArrayAgg to UDAF #10999

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Move array `ArrayAgg` to a `UserDefinedAggregate` #11448

Move array `ArrayAgg` to a `UserDefinedAggregate` #11448

jayzhan211 commented Jul 13, 2024 •

edited

Loading

jayzhan211 Jul 17, 2024

alamb Jul 18, 2024

jayzhan211 Jul 17, 2024

dharanad Jul 18, 2024

alamb Jul 18, 2024

jayzhan211 Jul 17, 2024

dharanad Jul 18, 2024

alamb Jul 18, 2024

jayzhan211 Jul 18, 2024

dharanad Jul 18, 2024

dharanad Jul 18, 2024

dharanad Jul 18, 2024

dharanad Jul 18, 2024

jayzhan211 Jul 18, 2024

dharanad Jul 18, 2024 •

edited

Loading

dharanad Jul 18, 2024

dharanad Jul 18, 2024

dharanad Jul 18, 2024

alamb left a comment

alamb Jul 18, 2024

alamb Jul 18, 2024

alamb Jul 18, 2024

alamb Jul 18, 2024

jayzhan211 commented Jul 19, 2024

alamb Jul 19, 2024

	# test array_agg_distinct with list data type
	statement ok
	CREATE TABLE array_agg_distinct_list_table AS VALUES
	('w', [0,1]),
	('w', [0,1]),
	('w', [1,0]),
	('b', [1,0]),
	('b', [1,0]),
	('b', [1,0]),
	('b', [0,1])
	;

	# Apply array_sort to have deterministic result, higher dimension nested array also works but not for array sort,
	# so they are covered in `datafusion/physical-expr/src/aggregate/array_agg_distinct.rs`
	query ??
	select array_sort(c1), array_sort(c2) from (
	select array_agg(distinct column1) as c1, array_agg(distinct column2) as c2 from array_agg_distinct_list_table
	);
	----
	[b, w] [[0, 1], [1, 0]]

	//! Defines physical expressions that can evaluated at runtime during query execution
	//! `ARRAY_AGG` aggregate implementation: [`ArrayAgg`]

Move array ArrayAgg to a UserDefinedAggregate #11448

Move array ArrayAgg to a UserDefinedAggregate #11448

Conversation

jayzhan211 commented Jul 13, 2024 • edited Loading

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

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

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

dharanad Jul 18, 2024 • edited Loading

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

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

jayzhan211 commented Jul 19, 2024

Choose a reason for hiding this comment

Move array `ArrayAgg` to a `UserDefinedAggregate` #11448

Move array `ArrayAgg` to a `UserDefinedAggregate` #11448

jayzhan211 commented Jul 13, 2024 •

edited

Loading

dharanad Jul 18, 2024 •

edited

Loading