Move conversion of FIRST/LAST Aggregate function to independent physical optimizer rule #10061

jayzhan211 · 2024-04-12T13:08:39Z

Which issue does this PR close?

Closes #9972.

Rationale for this change

We plan to make FIRST / LAST UDAF. This rule does the conversion between FIRST/LAST, it will eventually be moved to aggregate-functions crate. The first step is to move it out to an independent rule.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

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

jayzhan211 · 2024-04-12T13:31:49Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+            &input_order_mode,
+        );
+
+        let aggr_exec = aggr_exec.new_with_aggr_expr_and_ordering_info(


keep the logic similar to AggregateExec::try_new_with_schema

Ok(AggregateExec { mode, group_by, aggr_expr, filter_expr, input, schema, input_schema, metrics: ExecutionPlanMetricsSet::new(), required_input_ordering, limit: None, input_order_mode, cache, })

datafusion/core/src/physical_optimizer/convert_first_last.rs

ozankabak · 2024-04-12T14:19:40Z

We will review and comment on this next week. @mustafasrepo

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

jayzhan211 · 2024-04-13T07:26:20Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+///
+/// Similar to the one in datafusion/physical-plan/src/aggregates/mod.rs, but this
+/// function care only the possible conversion between FIRST_VALUE and LAST_VALUE
+fn get_aggregate_exprs_requirement(


only first/last rule is moved to here.

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

alamb

Thank you @jayzhan211 -- this is a very nice contribution towards extracting aggregate functions out of the core. I think we should fix the double recursion but otherwise this code looks (really) nice to me.

cc @mustafasrepo and @ozankabak FYI

alamb · 2024-04-13T13:25:27Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+/// so we can convert the aggregate expression to FirstValue(c1 order by asc),
+/// since the current ordering is already satisfied, it saves our time!
+#[derive(Default)]
+pub struct ConvertFirstLast {}


I wonder if we could call this something more general, like OptimizeAggregateOrder so it could potentially be used for aggregates other than FIRST_VALUE and LAST_VALUE 🤔

alamb · 2024-04-13T13:25:31Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+    }
+
+    fn name(&self) -> &str {
+        "SimpleOrdering"


I think this name should match the name of the structure -- that is "ConvertFirstLast" in this case

alamb · 2024-04-13T13:26:36Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+fn get_common_requirement_of_aggregate_input(
+    plan: Arc<dyn ExecutionPlan>,
+) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+    // Optimize children


since this rule already calls transform_up which handles the recursion up the tree of ExecutionPlan and managine the transformedflag, I don't think you also need to recursively walk down the children here again. I think you can probably just call optimize_internal directly

Recursing back down the tree is also like N^2 (or worse) in the number of plan nodes so I think we should avoid it for performance reasons (in addition to making the code simpler)

Amazing, it works! I think I have no idea what is going on in transform_up 😞

/// Convenience utility for writing optimizer rules: Recursively apply the /// given function `f` to all children of a node, and then to the node itself /// (post-order traversal). When `f` does not apply to a given node, it is /// left unchanged.

I didn't notice that the children is updated to parent too, so I do it manually.

$SELF.map_children($F_CHILD)?.transform_parent(|n| $F_UP(n))

I think transform_parent here is actually doing transform_self 🤔

alamb · 2024-04-13T13:31:34Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+    }
+}
+
+/// In `create_initial_plan` for LogicalPlan::Aggregate, we have a nested AggregateExec where the first layer


thank you for this comment. It makes things much clearer

alamb · 2024-04-13T13:33:21Z

datafusion/core/src/physical_optimizer/convert_first_last.rs

+        let reverse_aggr_req =
+            PhysicalSortRequirement::from_sort_exprs(&reverse_aggr_req);
+
+        if let Some(first_value) = aggr_expr.as_any().downcast_ref::<FirstValue>() {


Eventually (some other PR) it would be amazing if we can move this code into FirstValue somehow. As it is now, there is a coupling between the optimizer rule and the actual PhysicalExpr -- which means among other things this same optimization can't be used by user defined aggregates

Yes, this should be similar to FunctionRewrite, registerable optimize rule.

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

ozankabak · 2024-04-13T14:18:33Z

Thank you @jayzhan211 -- this is a very nice contribution towards extracting aggregate functions out of the core. I think we should fix the double recursion but otherwise this code looks (really) nice to me.

cc @mustafasrepo and @ozankabak FYI

We will review this Monday

alamb · 2024-04-13T20:11:11Z

Sounds good -- thank you @ozankabak -- let's wait for that review prior to merging this PR

mustafasrepo · 2024-04-15T11:58:00Z

datafusion/core/src/physical_optimizer/optimizer.rs

@@ -89,6 +90,8 @@ impl PhysicalOptimizer {
            // as that rule may inject other operations in between the different AggregateExecs.
            // Applying the rule early means only directly-connected AggregateExecs must be examined.
            Arc::new(LimitedDistinctAggregation::new()),
+            // Run once before PartialFinalAggregation is rewritten to ensure the rule is applied correctly
+            Arc::new(OptimizeAggregateOrder::new()),


I think we can remove this rule from here. Using it only in below place should be enough.

mustafasrepo

Thanks @jayzhan211 for this PR. I have left some minor comment. Please note that, applying my suggestion changes couple of tests. However, I have verified that those changes are both valid, not harmful for the execution. We can merge this PR as is also.

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

jayzhan211 · 2024-04-15T12:23:31Z

Arc::new(OptimizeAggregateOrder::new()),

Let me apply your suggestion! Thanks for your review @mustafasrepo and @alamb

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

jayzhan211 added 23 commits April 10, 2024 14:55

move out the ordering ruel

e506662

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

introduce rule

2812886

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

revert test result

d20dde3

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

pass mulit order test

c671033

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

cleanup

3f43dc1

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

with new childes

c8b46a4

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

revert slt

fcb297b

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

revert back

ab10168

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

rm rewrite in new child

fffb54a

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

backup

fec3eaf

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

only move conversion to optimizer

5b68b48

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

find test that do reverse

2ab74c9

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

add test for first and last

03abf19

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

pass all test

83c721a

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

upd test

34b8f01

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

upd test

b17e8f7

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

cleanup

a313af0

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

cleanup

a37fb3e

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

cleanup

8de501c

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

add aggregate test

107f4c8

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

cleanup

345d4c6

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

final draft

030d7b9

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

cleanup again

cf3f33e

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

github-actions bot added core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) labels Apr 12, 2024

jayzhan211 commented Apr 12, 2024

View reviewed changes

datafusion/core/src/physical_optimizer/convert_first_last.rs Outdated Show resolved Hide resolved

jayzhan211 commented Apr 12, 2024

View reviewed changes

datafusion/core/src/physical_optimizer/convert_first_last.rs Show resolved Hide resolved

pull out finer ordering code and reuse

6d9ee9f

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

jayzhan211 added 2 commits April 13, 2024 09:22

clippy

abd40c6

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

remove finer in optimize rule

3e2c186

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

jayzhan211 commented Apr 13, 2024

View reviewed changes

jayzhan211 added 4 commits April 13, 2024 15:28

add comments and clenaup

d7c2590

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

rename fun

fed356a

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

rename fun

0d3d461

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

fmt

d3daa4e

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

alamb reviewed Apr 13, 2024

View reviewed changes

avoid unnecessary recursion and rename

dcb3e90

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

mustafasrepo reviewed Apr 15, 2024

View reviewed changes

mustafasrepo approved these changes Apr 15, 2024

View reviewed changes

remove unnecessary rule

2951fbc

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

jayzhan211 added 3 commits April 15, 2024 20:23

Merge remote-tracking branch 'upstream/main' into aggr-rule

af45a22

Merge remote-tracking branch 'upstream/main' into aggr-rule

35b84ae

fix merge

fffcd3b

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

jayzhan211 merged commit 4e9f2d5 into apache:main Apr 15, 2024
24 checks passed

mustafasrepo mentioned this pull request Apr 15, 2024

Simplify Rule synnada-ai/datafusion-upstream#17

Closed

Move conversion of FIRST/LAST Aggregate function to independent physical optimizer rule #10061

Move conversion of FIRST/LAST Aggregate function to independent physical optimizer rule #10061

Uh oh!

Conversation

jayzhan211 commented Apr 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ozankabak commented Apr 12, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ozankabak commented Apr 13, 2024

Uh oh!

alamb commented Apr 13, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mustafasrepo left a comment

Choose a reason for hiding this comment

Uh oh!

jayzhan211 commented Apr 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jayzhan211 commented Apr 12, 2024 •

edited

Loading

jayzhan211 commented Apr 15, 2024 •

edited

Loading