4 phase

924060929 · Jun 4, 2024 · 3143262 · 3143262
1 parent df006bd
commit 3143262
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 5 deletions.
diff --git a/...core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java b/...core/src/main/java/org/apache/doris/nereids/rules/implementation/AggregateStrategies.java
@@ -363,7 +363,7 @@ && couldConvertToMulti(agg))
              */
             RuleType.FOUR_PHASE_AGGREGATE_WITH_DISTINCT_WITH_FULL_DISTRIBUTE.build(
                 basePattern
-                    .when(agg -> agg.getDistinctArguments().size() == 1 && !agg.getGroupByExpressions().isEmpty())
+                    .when(agg -> agg.everyDistinctArgumentsIsOne() && !agg.getGroupByExpressions().isEmpty())
                     .when(agg ->
                         ImmutableSet.builder()
                             .addAll(agg.getGroupByExpressions())

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/Aggregate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/Aggregate.java
@@ -29,6 +29,7 @@
 
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Common interface for logical/physical Aggregate.
@@ -68,4 +69,25 @@ default Set<Expression> getDistinctArguments() {
         }
         return distinctArguments.build();
     }
+
+    /** everyDistinctAggregateFunctionIsOne */
+    default boolean everyDistinctArgumentsIsOne() {
+        AtomicBoolean hasDistinctArguments = new AtomicBoolean(false);
+        for (NamedExpression outputExpression : getOutputExpressions()) {
+            boolean distinctArgumentSizeNotOne = outputExpression.anyMatch(expr -> {
+                if (expr instanceof AggregateFunction) {
+                    AggregateFunction aggFun = (AggregateFunction) expr;
+                    if (aggFun.isDistinct()) {
+                        hasDistinctArguments.set(true);
+                        return aggFun.getDistinctArguments().size() != 1;
+                    }
+                }
+                return false;
+            });
+            if (distinctArgumentSizeNotOne) {
+                return false;
+            }
+        }
+        return hasDistinctArguments.get();
+    }
 }
diff --git a/regression-test/data/nereids_p0/aggregate/aggregate.out b/regression-test/data/nereids_p0/aggregate/aggregate.out
@@ -687,5 +687,5 @@ TESTING	AGAIN
 7	-32767.0
 
 -- !four_phase_full_distribute --
-2	43
-
+hello	1	1
+world	1	1
diff --git a/regression-test/suites/nereids_p0/aggregate/aggregate.groovy b/regression-test/suites/nereids_p0/aggregate/aggregate.groovy
@@ -364,8 +364,9 @@ suite("aggregate") {
     sql "insert into test_four_phase_full_distribute values(1, 21, 'hello'), (2, 22, 'world')"
     sql " sync "
     order_qt_four_phase_full_distribute """select
-        /*+SET_VAR(disable_nereids_rules='THREE_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI,THREE_PHASE_AGGREGATE_WITH_DISTINCT,TWO_PHASE_AGGREGATE_SINGLE_DISTINCT_TO_MULTI,FOUR_PHASE_AGGREGATE_WITH_DISTINCT')*/
-        count(distinct name), sum(age)
+        /*+SET_VAR(disable_nereids_rules='TWO_PHASE_AGGREGATE_SINGLE_DISTINCT_TO_MULTI,TWO_PHASE_AGGREGATE_WITH_MULTI_DISTINCT,THREE_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI,THREE_PHASE_AGGREGATE_WITH_DISTINCT,FOUR_PHASE_AGGREGATE_WITH_DISTINCT')*/
+        name, count(distinct name), count(distinct age)
         from test_four_phase_full_distribute
+        group by name
         """
 }