[BugFix] Fix CTE distinct grouping sets rewrite generate invalid plan (…

…#48765)

Signed-off-by: stdpain <[email protected]>
(cherry picked from commit 6076570)

# Conflicts:
#	fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java
#	fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/RuleType.java

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java

@@ -35,6 +35,7 @@
 import com.starrocks.sql.optimizer.rule.mv.MaterializedViewRule;
 import com.starrocks.sql.optimizer.rule.transformation.ApplyExceptionRule;
 import com.starrocks.sql.optimizer.rule.transformation.ArrayDistinctAfterAggRule;
+import com.starrocks.sql.optimizer.rule.transformation.CTEProduceAddProjectionRule;
 import com.starrocks.sql.optimizer.rule.transformation.ConvertToEqualForNullRule;
 import com.starrocks.sql.optimizer.rule.transformation.DeriveRangeJoinPredicateRule;
 import com.starrocks.sql.optimizer.rule.transformation.GroupByCountDistinctRewriteRule;
@@ -410,7 +411,12 @@ private OptExpression logicalRuleRewrite(ConnectContext connectContext,
         ruleRewriteOnlyOnce(tree, rootTaskContext, RuleSetType.PARTITION_PRUNE);
         ruleRewriteIterative(tree, rootTaskContext, new RewriteMultiDistinctRule());
         ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PUSH_DOWN_PREDICATE);
+<<<<<<< HEAD
         ruleRewriteOnlyOnce(tree, rootTaskContext, LimitPruneTabletsRule.getInstance());
+=======
+        ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PRUNE_EMPTY_OPERATOR);
+        ruleRewriteIterative(tree, rootTaskContext, new CTEProduceAddProjectionRule());
+>>>>>>> 6076570358 ([BugFix] Fix CTE distinct grouping sets rewrite generate invalid plan (#48765))
         ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PRUNE_PROJECT);
 
         // ArrayDistinctAfterAggRule must run before pushDownAggregation,

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/RuleType.java

@@ -174,6 +174,14 @@ public enum RuleType {
     TF_CONVERT_TO_EQUAL_FOR_NULL_RULE,
     TF_ARRAY_DISTINCT_AFTER_AGG,
     TF_ELIMINATE_GROUP_BY_CONSTANT,
+<<<<<<< HEAD
+=======
+    TF_ELIMINATE_AGG,
+
+    TF_CTE_ADD_PROJECTION,
+
+    TF_PREDICATE_PROPAGATE,
+>>>>>>> 6076570358 ([BugFix] Fix CTE distinct grouping sets rewrite generate invalid plan (#48765))
 
     // The following are implementation rules:
     IMP_OLAP_LSCAN_TO_PSCAN,

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/CTEProduceAddProjectionRule.java

@@ -0,0 +1,64 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed 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
+//
+//     https://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 com.starrocks.sql.optimizer.rule.transformation;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Lists;
+import com.starrocks.sql.optimizer.OptExpression;
+import com.starrocks.sql.optimizer.OptimizerContext;
+import com.starrocks.sql.optimizer.RowOutputInfo;
+import com.starrocks.sql.optimizer.operator.ColumnOutputInfo;
+import com.starrocks.sql.optimizer.operator.OperatorType;
+import com.starrocks.sql.optimizer.operator.Projection;
+import com.starrocks.sql.optimizer.operator.logical.LogicalRepeatOperator;
+import com.starrocks.sql.optimizer.operator.pattern.Pattern;
+import com.starrocks.sql.optimizer.operator.scalar.ColumnRefOperator;
+import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
+import com.starrocks.sql.optimizer.rule.RuleType;
+
+import java.util.Collections;
+import java.util.List;
+
+public class CTEProduceAddProjectionRule extends TransformationRule {
+
+    public CTEProduceAddProjectionRule() {
+        super(RuleType.TF_CTE_ADD_PROJECTION,
+                Pattern.create(OperatorType.LOGICAL_CTE_PRODUCE, OperatorType.LOGICAL_REPEAT));
+    }
+
+    @Override
+    public List<OptExpression> transform(OptExpression input, OptimizerContext context) {
+        final OptExpression repeatOpt = input.getInputs().get(0);
+        final LogicalRepeatOperator repeat = (LogicalRepeatOperator) repeatOpt.getOp();
+        if (repeat.getProjection() == null) {
+            final RowOutputInfo rowOutputInfo = repeatOpt.getInputs().get(0).getRowOutputInfo();
+            final ImmutableMap.Builder<ColumnRefOperator, ScalarOperator> builder = ImmutableMap.builder();
+
+            for (ColumnOutputInfo columnOutputInfo : rowOutputInfo.getColumnOutputInfo()) {
+                final ColumnRefOperator columnRef = columnOutputInfo.getColumnRef();
+                builder.put(columnRef, columnRef);
+            }
+
+            // add grouping id to projection
+            for (ColumnRefOperator columnRefOperator : repeat.getOutputGrouping()) {
+                builder.put(columnRefOperator, columnRefOperator);
+            }
+            repeat.setProjection(new Projection(builder.build()));
+            return Lists.newArrayList(input);
+        } else {
+            return Collections.emptyList();
+        }
+    }
+}

fe/fe-core/src/test/java/com/starrocks/sql/plan/AggregateTest.java

@@ -1528,6 +1528,32 @@ public void testMultiDistinctAggregate() throws Exception {
                 "  |  colocate: false, reason: \n" +
                 "  |  equal join conjunct: 16: t1c <=> 27: t1c\n" +
                 "  |  equal join conjunct: 17: expr <=> 28: expr");
+        // count distinct with grouping sets
+
+        sql = "select avg(distinct t1b) as cn_t1b, sum(distinct t1b), " +
+                "count(distinct t1b, t1c) cn_t1b_t1c from test_all_type group by rollup(t1c, t1b)";
+        plan = getFragmentPlan(sql);
+        // make sure repeat + project + multi cast sink
+        assertContains(plan, "  MultiCastDataSinks\n" +
+                "  STREAM DATA SINK\n" +
+                "    EXCHANGE ID: 04\n" +
+                "    RANDOM\n" +
+                "  STREAM DATA SINK\n" +
+                "    EXCHANGE ID: 10\n" +
+                "    RANDOM\n" +
+                "  STREAM DATA SINK\n" +
+                "    EXCHANGE ID: 17\n" +
+                "    RANDOM\n" +
+                "\n" +
+                "  3:Project\n" +
+                "  |  <slot 2> : 2: t1b\n" +
+                "  |  <slot 3> : 3: t1c\n" +
+                "  |  <slot 13> : 13: expr\n" +
+                "  |  <slot 14> : 14: expr\n" +
+                "  |  <slot 18> : 18: GROUPING_ID\n" +
+                "  |  \n" +
+                "  2:REPEAT_NODE\n" +
+                "  |  repeat: repeat 2 lines [[], [3], [2, 3]]");
     }
 
     @Test