[fix](mtmv) Fix query rewrite by mv fail when both use cte and group sets

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

@@ -63,24 +63,29 @@
 import javax.annotation.Nullable;
 
 /** NormalizeRepeat
- * eg: select sum(b + 1), grouping(a+1) from t1 group by grouping sets ((b+1));
+ * eg: SELECT
+ *      camp,
+ *      COUNT(occupation) AS occ_cnt,
+ *      GROUPING(camp) AS grouping
+ *     FROM
+ *      `roles`
+ *      GROUP BY ROLLUP(camp);
  * Original Plan:
- * LogicalRepeat ( groupingSets=[[(a#0 + 1)]],
- *                 outputExpressions=[sum((b#1 + 1)) AS `sum((b + 1))`#2,
- *                 Grouping((a#0 + 1)) AS `Grouping((a + 1))`#3] )
- *      +--LogicalOlapScan (t1)
- *
+ * LogicalRepeat ( groupingSets=[[camp#2], []], outputExpressions=[camp#2,
+ *                  count(occupation#1) AS `occ_cnt`#6, Grouping(camp#2) AS `grouping`#7], groupingId=Optional.empty )
+ *  +--LogicalFilter[10] ( predicates=(__DORIS_DELETE_SIGN__#4 = 0) )
+ *      +--LogicalOlapScan ( qualified=roles, indexName=index_not_selected, selectedIndexId=1765187322191,
+ *                              preAgg=UNSET, operativeCol=[], virtualColumns=[] )
  * After:
- * LogicalAggregate[62] ( groupByExpr=[(a + 1)#4, GROUPING_ID#7, GROUPING_PREFIX_(a + 1)#6],
- *                        outputExpr=[sum((b + 1)#5) AS `sum((b + 1))`#2,
- *                                    GROUPING_PREFIX_(a + 1)#6 AS `GROUPING_PREFIX_(a + 1)`#3] )
- *    +--LogicalRepeat ( groupingSets=[[(a + 1)#4]],
- *                       outputExpressions=[(a + 1)#4,
- *                                          (b + 1)#5,
- *                                          GROUPING_ID#7,
- *                                          GROUPING_PREFIX_(a + 1)#6] )
- *      +--LogicalProject[60] ( projects=[(a#0 + 1) AS `(a + 1)`#4, (b#1 + 1) AS `(b + 1)`#5], excepts=[]
- *          +--LogicalOlapScan ( t1 )
+ * LogicalAggregate[19] ( groupByExpr=[camp#2, GROUPING_PREFIX_camp#8, GROUPING_ID#9],
+ *                          outputExpr=[camp#2, count(occupation#1) AS `occ_cnt`#6,
+ *                              GROUPING_PREFIX_camp#8 AS `grouping`#7], hasRepeat=true )
+ *  +--LogicalRepeat ( groupingSets=[[camp#2], []], outputExpressions=[camp#2, occupation#1,
+ *                              Grouping(camp#2) AS `GROUPING_PREFIX_camp`#8], groupingId=Optional[GROUPING_ID#9] )
+ *      +--LogicalProject[17] ( distinct=false, projects=[camp#2, occupation#1] )
+ *          +--LogicalFilter[10] ( predicates=(__DORIS_DELETE_SIGN__#4 = 0) )
+ *              +--LogicalOlapScan ( qualified=roles, indexName=index_not_selected,
+ *                                  selectedIndexId=1765187322191, preAgg=UNSET, operativeCol=[], virtualColumns=[] )
  */
 public class NormalizeRepeat extends OneAnalysisRuleFactory {
     @Override
@@ -198,7 +203,16 @@ private static LogicalAggregate<Plan> normalizeRepeat(LogicalRepeat<Plan> repeat
 
         Plan normalizedChild = pushDownProject(pushedProject, repeat.child());
 
-        SlotReference groupingId = new SlotReference(Repeat.COL_GROUPING_ID, BigIntType.INSTANCE, false);
+        // If grouping id is not present, we need to add it, if repeat already has grouping id, use it directly
+        // which is for the case repeat is introduced by mv rewrite, should keep the rewritten grouping id
+        // is same to the original grouping id
+        SlotReference groupingId = repeat.getGroupingId().orElse(
+                new SlotReference(Repeat.COL_GROUPING_ID, BigIntType.INSTANCE, false));
+        // remove grouping id from repeat output expressions, grouping id should not in repeat output
+        // this keep consistent with original repeat behavior
+        normalizedRepeatOutput = normalizedRepeatOutput.stream()
+                .filter(expr -> !expr.equals(groupingId))
+                .collect(Collectors.toList());
         LogicalRepeat<Plan> normalizedRepeat = repeat.withNormalizedExpr(
                 (List) normalizedGroupingSets, normalizedRepeatOutput, groupingId, normalizedChild);
 

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewAggregateRule.java

@@ -240,7 +240,9 @@ protected LogicalAggregate<Plan> aggregateRewriteByView(
         if (queryAggregate.getSourceRepeat().isPresent()) {
             // construct group sets for repeat
             List<List<Expression>> rewrittenGroupSetsExpressions = new ArrayList<>();
-            List<List<Expression>> groupingSets = queryAggregate.getSourceRepeat().get().getGroupingSets();
+            List<List<Expression>> groupingSets = queryAggregate.collectFirst(LogicalRepeat.class::isInstance)
+                    .map(repeat -> ((LogicalRepeat<? extends Plan>) repeat).getGroupingSets())
+                    .orElse(queryAggregate.getSourceRepeat().get().getGroupingSets());
             for (List<Expression> groupingSet : groupingSets) {
                 if (groupingSet.isEmpty()) {
                     rewrittenGroupSetsExpressions.add(ImmutableList.of());
@@ -262,7 +264,7 @@ protected LogicalAggregate<Plan> aggregateRewriteByView(
                 }
             }
             LogicalRepeat<Plan> repeat = new LogicalRepeat<>(rewrittenGroupSetsExpressions,
-                    finalOutputExpressions, tempRewritedPlan);
+                    finalOutputExpressions, queryStructInfo.getGroupingId().get(), tempRewritedPlan);
             return NormalizeRepeat.doNormalize(repeat);
         }
         return new LogicalAggregate<>(finalGroupExpressions, finalOutputExpressions, tempRewritedPlan);

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewFilterAggregateRule.java

@@ -38,8 +38,11 @@ public class MaterializedViewFilterAggregateRule extends AbstractMaterializedVie
     @Override
     public List<Rule> buildRules() {
         return ImmutableList.of(
-                logicalFilter(logicalAggregate(any().when(LogicalPlan.class::isInstance))).thenApplyMultiNoThrow(
-                        ctx -> {
+                logicalFilter(logicalAggregate(any().when(LogicalPlan.class::isInstance))
+                        // Temporarily unsupported: do not rewrite when GROUP SETS is used together with HAVING.
+                        // Will be re-enabled after the fix is completed.
+                        .when(aggregate -> !aggregate.getSourceRepeat().isPresent()))
+                        .thenApplyMultiNoThrow(ctx -> {
                             LogicalFilter<LogicalAggregate<Plan>> root = ctx.root;
                             return rewrite(root, ctx.cascadesContext);
                         }).toRule(RuleType.MATERIALIZED_VIEW_FILTER_AGGREGATE));

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewFilterProjectAggregateRule.java

@@ -40,7 +40,10 @@ public class MaterializedViewFilterProjectAggregateRule extends AbstractMaterial
     @Override
     public List<Rule> buildRules() {
         return ImmutableList.of(
-                logicalFilter(logicalProject(logicalAggregate(any().when(LogicalPlan.class::isInstance))))
+                logicalFilter(logicalProject(logicalAggregate(any().when(LogicalPlan.class::isInstance)).when(
+                        // Temporarily unsupported: do not rewrite when GROUP SETS is used together with HAVING.
+                        // Will be re-enabled after the fix is completed.
+                        aggregate -> !aggregate.getSourceRepeat().isPresent())))
                         .thenApplyMultiNoThrow(ctx -> {
                             LogicalFilter<LogicalProject<LogicalAggregate<Plan>>> root = ctx.root;
                             return rewrite(root, ctx.cascadesContext);

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewProjectFilterAggregateRule.java

@@ -39,10 +39,14 @@ public class MaterializedViewProjectFilterAggregateRule extends AbstractMaterial
 
     @Override
     public List<Rule> buildRules() {
-        return ImmutableList.of(logicalProject(logicalFilter(logicalAggregate(
-                any().when(LogicalPlan.class::isInstance)))).thenApplyMultiNoThrow(ctx -> {
-                    LogicalProject<LogicalFilter<LogicalAggregate<Plan>>> root = ctx.root;
-                    return rewrite(root, ctx.cascadesContext);
-                }).toRule(RuleType.MATERIALIZED_VIEW_PROJECT_FILTER_AGGREGATE));
+        return ImmutableList.of(
+                logicalFilter(logicalProject(logicalAggregate(any().when(LogicalPlan.class::isInstance)).when(
+                        // Temporarily unsupported: do not rewrite when GROUP SETS is used together with HAVING.
+                        // Will be re-enabled after the fix is completed.
+                        aggregate -> !aggregate.getSourceRepeat().isPresent())))
+                        .thenApplyMultiNoThrow(ctx -> {
+                            LogicalFilter<LogicalProject<LogicalAggregate<Plan>>> root = ctx.root;
+                            return rewrite(root, ctx.cascadesContext);
+                        }).toRule(RuleType.MATERIALIZED_VIEW_FILTER_PROJECT_AGGREGATE));
     }
 }

regression-test/data/nereids_rules_p0/mv/grouping_sets/grouping_sets.out

@@ -233,3 +233,63 @@ o	2023-12-10	a	46.00	33.50	12.50	2	\N
 o	2023-12-11	c	43.20	43.20	43.20	1	\N
 o	2023-12-12	c	57.40	56.20	1.20	2	\N
 
+-- !query15_0_before --
+\N	\N
+\N	2023-12-11
+3	\N
+3	2023-12-11
+
+-- !query15_0_after --
+\N	\N
+\N	2023-12-11
+3	\N
+3	2023-12-11
+
+-- !query16_0_before --
+\N	\N	\N	1	1	3	7	43.20	43.20	43.20	1	0
+\N	\N	2023-12-11	1	0	3	6	43.20	43.20	43.20	1	0
+\N	3	\N	0	1	2	5	43.20	43.20	43.20	1	0
+\N	3	2023-12-11	0	0	2	4	43.20	43.20	43.20	1	0
+3	\N	\N	1	1	1	3	43.20	43.20	43.20	1	0
+3	\N	2023-12-11	1	0	1	2	43.20	43.20	43.20	1	0
+3	3	\N	0	1	0	1	43.20	43.20	43.20	1	0
+3	3	2023-12-11	0	0	0	0	43.20	43.20	43.20	1	0
+
+-- !query16_0_after --
+\N	\N	\N	1	1	3	7	43.20	43.20	43.20	1	0
+\N	\N	2023-12-11	1	0	3	6	43.20	43.20	43.20	1	0
+\N	3	\N	0	1	2	5	43.20	43.20	43.20	1	0
+\N	3	2023-12-11	0	0	2	4	43.20	43.20	43.20	1	0
+3	\N	\N	1	1	1	3	43.20	43.20	43.20	1	0
+3	\N	2023-12-11	1	0	1	2	43.20	43.20	43.20	1	0
+3	3	\N	0	1	0	1	43.20	43.20	43.20	1	0
+3	3	2023-12-11	0	0	0	0	43.20	43.20	43.20	1	0
+
+-- !query17_before --
+\N	\N	\N	1	1	3	7	43.20	43.20	43.20	1	0
+\N	\N	2023-12-11	1	0	3	6	43.20	43.20	43.20	1	0
+3	\N	\N	1	1	1	3	43.20	43.20	43.20	1	0
+3	\N	2023-12-11	1	0	1	2	43.20	43.20	43.20	1	0
+
+-- !query17_after --
+\N	\N	\N	1	1	3	7	43.20	43.20	43.20	1	0
+\N	\N	2023-12-11	1	0	3	6	43.20	43.20	43.20	1	0
+3	\N	\N	1	1	1	3	43.20	43.20	43.20	1	0
+3	\N	2023-12-11	1	0	1	2	43.20	43.20	43.20	1	0
+
+-- !query18_before --
+\N	\N	\N	1	1	3	7	43.20	43.20	43.20	1	0
+\N	\N	2023-12-11	1	0	3	6	43.20	43.20	43.20	1	0
+\N	3	\N	0	1	2	5	43.20	43.20	43.20	1	0
+\N	3	2023-12-11	0	0	2	4	43.20	43.20	43.20	1	0
+3	\N	\N	1	1	1	3	43.20	43.20	43.20	1	0
+3	\N	2023-12-11	1	0	1	2	43.20	43.20	43.20	1	0
+
+-- !query18_after --
+\N	\N	\N	1	1	3	7	43.20	43.20	43.20	1	0
+\N	\N	2023-12-11	1	0	3	6	43.20	43.20	43.20	1	0
+\N	3	\N	0	1	2	5	43.20	43.20	43.20	1	0
+\N	3	2023-12-11	0	0	2	4	43.20	43.20	43.20	1	0
+3	\N	\N	1	1	1	3	43.20	43.20	43.20	1	0
+3	\N	2023-12-11	1	0	1	2	43.20	43.20	43.20	1	0
+

regression-test/suites/nereids_rules_p0/mv/grouping_sets/grouping_sets.groovy

@@ -474,7 +474,6 @@ suite("materialized_view_grouping_sets") {
     order_qt_query10_0_after "${query10_0}"
     sql """ DROP MATERIALIZED VIEW IF EXISTS mv10_0"""
 
-
     // single table rollup with grouping scalar function and filter
     def mv10_1 =
             """
@@ -510,7 +509,6 @@ suite("materialized_view_grouping_sets") {
     order_qt_query10_1_after "${query10_1}"
     sql """ DROP MATERIALIZED VIEW IF EXISTS mv10_1"""
 
-
     // multi table rollup without grouping scalar function
     def mv11_0 =
             """
@@ -648,5 +646,171 @@ suite("materialized_view_grouping_sets") {
     async_mv_rewrite_fail(db, mv14_0, query14_0, "mv14_0")
     order_qt_query14_0_after "${query14_0}"
     sql """ DROP MATERIALIZED VIEW IF EXISTS mv14_0"""
+
+    // group sets with cte
+    def mv15_0 =
+            """
+            select l_shipdate, o_orderdate, l_partkey, l_suppkey,
+            sum(o_totalprice) as sum_total,
+            max(o_totalprice) as max_total,
+            min(o_totalprice) as min_total,
+            count(*) as count_all,
+            bitmap_union(to_bitmap(case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)) as bitmap_union_basic
+            from lineitem
+            left join orders on lineitem.l_orderkey = orders.o_orderkey and l_shipdate = o_orderdate
+            group by
+            l_shipdate,
+            o_orderdate,
+            l_partkey,
+            l_suppkey;
+            """
+    def query15_0 =
+            """
+            with t as (
+                select t1.l_partkey, t1.l_suppkey, o_orderdate,
+                grouping(t1.l_suppkey),
+                grouping(o_orderdate),
+                grouping_id(t1.l_partkey, t1.l_suppkey),
+                grouping_id(t1.l_partkey, t1.l_suppkey, o_orderdate),
+                sum(o_totalprice),
+                max(o_totalprice),
+                min(o_totalprice),
+                count(*),
+                count(distinct case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)
+                from (select * from lineitem where l_shipdate = '2023-12-11') t1
+                left join orders on t1.l_orderkey = orders.o_orderkey and t1.l_shipdate = o_orderdate
+                group by
+                GROUPING SETS ((l_shipdate, o_orderdate, l_partkey), (l_partkey, l_suppkey), (l_suppkey), ())
+            )
+            select t1.l_suppkey, t2.o_orderdate
+            from
+            t t1
+            inner join
+            t t2 on t1.l_partkey = t2.l_partkey;
+            """
+    order_qt_query15_0_before "${query15_0}"
+    async_mv_rewrite_success(db, mv15_0, query15_0, "mv15_0")
+    order_qt_query15_0_after "${query15_0}"
+    sql """ DROP MATERIALIZED VIEW IF EXISTS mv15_0"""
+
+    // group sets with alias and the alias is grouping id
+    def mv16_0 =
+            """
+            select l_shipdate, o_orderdate, l_partkey, l_suppkey,
+            sum(o_totalprice) as sum_total,
+            max(o_totalprice) as max_total,
+            min(o_totalprice) as min_total,
+            count(*) as count_all,
+            bitmap_union(to_bitmap(case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)) as bitmap_union_basic
+            from lineitem
+            left join orders on lineitem.l_orderkey = orders.o_orderkey and l_shipdate = o_orderdate
+            group by
+            l_shipdate,
+            o_orderdate,
+            l_partkey,
+            l_suppkey;
+            """
+    def query16_0 =
+            """
+            select t1.l_partkey as p_alias, t1.l_suppkey as s_alias, o_orderdate,
+            grouping(t1.l_suppkey),
+            grouping(o_orderdate),
+            grouping_id(t1.l_partkey, t1.l_suppkey),
+            grouping_id(t1.l_partkey, t1.l_suppkey, o_orderdate),
+            sum(o_totalprice),
+            max(o_totalprice),
+            min(o_totalprice),
+            count(*),
+            count(distinct case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)
+            from (select * from lineitem where l_shipdate = '2023-12-11') t1
+            left join orders on t1.l_orderkey = orders.o_orderkey and t1.l_shipdate = o_orderdate
+            group by
+            CUBE (t1.l_partkey, t1.l_suppkey, o_orderdate);
+            """
+    order_qt_query16_0_before "${query16_0}"
+    async_mv_rewrite_success(db, mv16_0, query16_0, "mv16_0")
+    order_qt_query16_0_after "${query16_0}"
+    sql """ DROP MATERIALIZED VIEW IF EXISTS mv16_0"""
+
+    // group sets with alias and the alias is grouping id and query has filter on column
+    def mv17 =
+            """
+            select l_shipdate, o_orderdate, l_partkey, l_suppkey,
+            sum(o_totalprice) as sum_total,
+            max(o_totalprice) as max_total,
+            min(o_totalprice) as min_total,
+            count(*) as count_all,
+            bitmap_union(to_bitmap(case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)) as bitmap_union_basic
+            from lineitem
+            left join orders on lineitem.l_orderkey = orders.o_orderkey and l_shipdate = o_orderdate
+            group by
+            l_shipdate,
+            o_orderdate,
+            l_partkey,
+            l_suppkey;
+            """
+    def query17 =
+            """
+            select t1.l_partkey as p_alias, t1.l_suppkey as s_alias, o_orderdate,
+            grouping(t1.l_suppkey),
+            grouping(o_orderdate),
+            grouping_id(t1.l_partkey, t1.l_suppkey),
+            grouping_id(t1.l_partkey, t1.l_suppkey, o_orderdate),
+            sum(o_totalprice),
+            max(o_totalprice),
+            min(o_totalprice),
+            count(*),
+            count(distinct case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)
+            from (select * from lineitem where l_shipdate = '2023-12-11') t1
+            left join orders on t1.l_orderkey = orders.o_orderkey and t1.l_shipdate = o_orderdate
+            group by
+            CUBE (t1.l_partkey, t1.l_suppkey, o_orderdate)
+            having grouping(t1.l_suppkey) >= 1;
+            """
+    order_qt_query17_before "${query17}"
+    async_mv_rewrite_success(db, mv17, query17, "mv17")
+    order_qt_query17_after "${query17}"
+    sql """ DROP MATERIALIZED VIEW IF EXISTS mv17"""
+
+
+    // group sets with alias and the alias is grouping id and query has filter on column
+    def mv18 =
+            """
+            select l_shipdate, o_orderdate, l_partkey, l_suppkey,
+            sum(o_totalprice) as sum_total,
+            max(o_totalprice) as max_total,
+            min(o_totalprice) as min_total,
+            count(*) as count_all,
+            bitmap_union(to_bitmap(case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)) as bitmap_union_basic
+            from lineitem
+            left join orders on lineitem.l_orderkey = orders.o_orderkey and l_shipdate = o_orderdate
+            group by
+            l_shipdate,
+            o_orderdate,
+            l_partkey,
+            l_suppkey;
+            """
+    def query18 =
+            """
+            select t1.l_partkey as p_alias, t1.l_suppkey as s_alias, o_orderdate,
+            grouping(t1.l_suppkey),
+            grouping(o_orderdate),
+            grouping_id(t1.l_partkey, t1.l_suppkey),
+            grouping_id(t1.l_partkey, t1.l_suppkey, o_orderdate),
+            sum(o_totalprice),
+            max(o_totalprice),
+            min(o_totalprice),
+            count(*),
+            count(distinct case when o_shippriority > 1 and o_orderkey IN (1, 3) then o_custkey else null end)
+            from (select * from lineitem where l_shipdate = '2023-12-11') t1
+            left join orders on t1.l_orderkey = orders.o_orderkey and t1.l_shipdate = o_orderdate
+            group by
+            CUBE (t1.l_partkey, t1.l_suppkey, o_orderdate)
+            having grouping_id(t1.l_partkey, t1.l_suppkey) >= 1;
+            """
+    order_qt_query18_before "${query18}"
+    async_mv_rewrite_success(db, mv18, query18, "mv18")
+    order_qt_query18_after "${query18}"
+    sql """ DROP MATERIALIZED VIEW IF EXISTS mv18"""
 }