add rule decompose_repeat

Author: feiniaofeiafei

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -60,6 +60,7 @@
 import org.apache.doris.nereids.rules.rewrite.CountDistinctRewrite;
 import org.apache.doris.nereids.rules.rewrite.CountLiteralRewrite;
 import org.apache.doris.nereids.rules.rewrite.CreatePartitionTopNFromWindow;
+import org.apache.doris.nereids.rules.rewrite.DecomposeRepeatWithPreAggregation;
 import org.apache.doris.nereids.rules.rewrite.DecoupleEncodeDecode;
 import org.apache.doris.nereids.rules.rewrite.DeferMaterializeTopNResult;
 import org.apache.doris.nereids.rules.rewrite.DistinctAggStrategySelector;
@@ -902,7 +903,8 @@ private static List<RewriteJob> getWholeTreeRewriteJobs(
                     rewriteJobs.addAll(jobs(topic("or expansion",
                             custom(RuleType.OR_EXPANSION, () -> OrExpansion.INSTANCE))));
                 }
-
+                rewriteJobs.add(topic("repeat rewrite",
+                        custom(RuleType.DECOMPOSE_REPEAT, () -> DecomposeRepeatWithPreAggregation.INSTANCE)));
                 rewriteJobs.addAll(jobs(topic("split multi distinct",
                         custom(RuleType.DISTINCT_AGG_STRATEGY_SELECTOR, () -> DistinctAggStrategySelector.INSTANCE))));
 

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java

@@ -415,7 +415,7 @@ public enum RuleType {
     SHOR_CIRCUIT_POINT_QUERY(RuleTypeClass.REWRITE),
     // skew rewrtie
     SALT_JOIN(RuleTypeClass.REWRITE),
-
+    DECOMPOSE_REPEAT(RuleTypeClass.REWRITE),
     DISTINCT_AGGREGATE_SPLIT(RuleTypeClass.REWRITE),
     PROCESS_SCALAR_AGG_MUST_USE_MULTI_DISTINCT(RuleTypeClass.REWRITE),
 

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/DecomposeRepeatWithPreAggregation.java

@@ -171,6 +171,10 @@ public Alias withChildren(List<Expression> children) {
         return new Alias(exprId, children, name, qualifier, nameFromChild);
     }
 
+    public Alias withExprId(ExprId exprId) {
+        return new Alias(exprId, children, name, qualifier, nameFromChild);
+    }
+
     public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
         return visitor.visitAlias(this, context);
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Alias.java

@@ -0,0 +1,99 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.util.MemoPatternMatchSupported;
+import org.apache.doris.nereids.util.PlanChecker;
+import org.apache.doris.utframe.TestWithFeService;
+
+import org.junit.jupiter.api.Test;
+
+/**
+ * UT for {@link DecomposeRepeatWithPreAggregation}.
+ */
+public class DecomposeRepeatWithPreAggregationTest extends TestWithFeService implements MemoPatternMatchSupported {
+
+    @Override
+    protected void runBeforeAll() throws Exception {
+        createDatabase("decompose_repeat_with_preagg");
+        createTable(
+                "create table decompose_repeat_with_preagg.t1 (\n"
+                        + "a int, b int, c int, d int\n"
+                        + ")\n"
+                        + "distributed by hash(a) buckets 1\n"
+                        + "properties('replication_num' = '1');"
+        );
+        connectContext.setDatabase("default_cluster:decompose_repeat_with_preagg");
+        connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION");
+    }
+
+    @Test
+    void rewriteRollupSumShouldGenerateCteAndUnion() {
+        String sql = "select a,b,c,sum(d) from t1 group by rollup(a,b,c);";
+        PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .matches(logicalCTEAnchor());
+    }
+
+    @Test
+    void noRewriteWhenGroupingSetsSizeLe3() {
+        String sql = "select a,b,sum(d) from t1 group by rollup(a,b);";
+        PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .nonMatch(logicalCTEAnchor());
+    }
+
+    @Test
+    void noRewriteWhenDistinctAgg() {
+        String sql = "select a,b,c,sum(distinct d) from t1 group by rollup(a,b,c);";
+        PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .nonMatch(logicalCTEAnchor());
+    }
+
+    @Test
+    void noRewriteWhenUnsupportedAgg() {
+        String sql = "select a,b,c,count(d) from t1 group by rollup(a,b,c);";
+        PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .nonMatch(logicalCTEAnchor());
+
+    }
+
+    @Test
+    void noRewriteWhenHasGroupingScalarFunction() {
+        String sql = "select a,b,c,sum(d),grouping_id(a) from t1 group by rollup(a,b,c);";
+        PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .nonMatch(logicalCTEAnchor());
+    }
+
+    @Test
+    void rewriteWhenMaxGroupingSetNotFirst() {
+        String sql = "select a,b,c,sum(d) from t1 group by grouping sets((a),(a,b,c),(a,b),());";
+        PlanChecker.from(connectContext)
+                .analyze(sql)
+                .rewrite()
+                .matches(logicalCTEAnchor());
+    }
+}

fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/DecomposeRepeatWithPreAggregationTest.java

@@ -0,0 +1,111 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !sum --
+\N	\N	\N	10
+1	\N	\N	10
+1	2	\N	8
+1	2	1	1
+1	2	3	7
+1	3	\N	2
+1	3	2	2
+
+-- !agg_func_gby_key_same_col --
+\N	\N	\N	\N	10
+1	\N	\N	\N	10
+1	2	\N	\N	8
+1	2	1	\N	1
+1	2	1	1	1
+1	2	3	\N	7
+1	2	3	3	3
+1	2	3	4	4
+1	3	\N	\N	2
+1	3	2	\N	2
+1	3	2	2	2
+
+-- !multi_agg_func --
+\N	\N	\N	10	9	1
+1	\N	\N	10	9	1
+1	2	\N	8	7	1
+1	2	1	1	1	1
+1	2	1	1	1	1
+1	2	3	3	3	1
+1	2	3	4	3	1
+1	2	3	7	6	1
+1	3	\N	2	2	1
+1	3	2	2	2	1
+1	3	2	2	2	1
+
+-- !nest_rewrite --
+\N	\N	\N	\N
+1	\N	\N	\N
+1	\N	\N	\N
+1	\N	\N	\N
+1	\N	\N	10
+1	2	\N	\N
+1	2	1	\N
+1	2	1	1
+1	2	3	\N
+1	2	3	3
+1	2	3	4
+1	2	3	7
+1	3	\N	\N
+1	3	2	\N
+1	3	2	2
+
+-- !upper_ref --
+11	1	2	1
+12	1	3	\N
+12	1	3	2
+17	1	2	3
+18	1	2	\N
+20	\N	\N	\N
+20	1	\N	\N
+
+-- !grouping_func --
+1	\N	\N	\N	10	0
+1	2	1	\N	1	0
+1	2	1	\N	1	0
+1	2	1	1	1	0
+1	2	3	\N	7	0
+1	2	3	\N	7	0
+1	2	3	3	3	0
+1	2	3	4	4	0
+1	3	2	\N	2	0
+1	3	2	\N	2	0
+1	3	2	2	2	0
+
+-- !avg --
+1	\N	\N	\N	2.5
+1	2	1	\N	1
+1	2	1	\N	1
+1	2	1	1	1
+1	2	3	\N	3.5
+1	2	3	\N	3.5
+1	2	3	3	3
+1	2	3	4	4
+1	3	2	\N	2
+1	3	2	\N	2
+1	3	2	2	2
+
+-- !distinct --
+1	\N	\N	\N	10
+1	2	1	\N	1
+1	2	1	\N	1
+1	2	1	1	1
+1	2	3	\N	7
+1	2	3	\N	7
+1	2	3	3	3
+1	2	3	4	4
+1	3	2	\N	2
+1	3	2	\N	2
+1	3	2	2	2
+
+-- !less_equal_than_3 --
+\N	\N	\N	\N	10
+1	2	1	\N	1
+1	2	1	1	1
+1	2	3	\N	7
+1	2	3	3	3
+1	2	3	4	4
+1	3	2	\N	2
+1	3	2	2	2
+

regression-test/data/nereids_rules_p0/decompose_repeat/decompose_repeat.out

@@ -0,0 +1,41 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+suite("decompose_repeat") {
+//    sql "set disable_nereids_rules='DECOMPOSE_REPEAT';"
+    sql "drop table if exists t1;"
+    sql "create table t1(a int, b int, c int, d int) distributed by hash(a) properties('replication_num'='1');"
+    sql "insert into t1 values(1,2,3,4),(1,2,3,3),(1,2,1,1),(1,3,2,2);"
+    order_qt_sum "select a,b,c,sum(d) from t1 group by rollup(a,b,c);"
+    order_qt_agg_func_gby_key_same_col "select a,b,c,d,sum(d) from t1 group by rollup(a,b,c,d);"
+    order_qt_multi_agg_func "select a,b,c,sum(d),sum(c),max(a) from t1 group by rollup(a,b,c,d);"
+    order_qt_nest_rewrite """
+    select a,b,c,c1 from (
+    select a,b,c,d,sum(d) c1 from t1 group by grouping sets((a,b,c),(a,b,c,d),(a),(a,b,c,c))
+    ) t group by rollup(a,b,c,c1);
+    """
+    order_qt_upper_ref """
+    select c1+10,a,b,c from (select a,b,c,sum(d) c1 from t1 group by rollup(a,b,c)) t group by c1+10,a,b,c;
+    """
+
+    // negative case
+    order_qt_grouping_func "select a,b,c,d,sum(d),grouping_id(a) from t1 group by grouping sets((a,b,c),(a,b,c,d),(a),(a,b,c,c))"
+    order_qt_avg "select a,b,c,d,avg(d) from t1 group by grouping sets((a,b,c),(a,b,c,d),(a),(a,b,c,c));"
+    order_qt_distinct "select a,b,c,d,sum(distinct d) from t1 group by grouping sets((a,b,c),(a,b,c,d),(a),(a,b,c,c));"
+    order_qt_less_equal_than_3 "select a,b,c,d,sum(distinct d) from t1 group by grouping sets((a,b,c),(a,b,c,d),());"
+
+}