[fix](Nereids) could not run query with repeat node in cte (apache#26330) (apache#26531)

morrySnow · gnehil · commit 3e11826410c5 · 2023-12-04T10:58:06.000+08:00
pick from master PR: apache#26330 commit id: a89477e ExpressionDeepCopier not process VirtualReference, so we generate inline plan with mistake.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/ExpressionDeepCopier.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/ExpressionDeepCopier.java
@@ -26,9 +26,15 @@
 import org.apache.doris.nereids.trees.expressions.ScalarSubquery;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
+import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator;
+import org.apache.doris.nereids.trees.expressions.VirtualSlotReference;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.GroupingScalarFunction;
 import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
+import org.apache.doris.nereids.trees.plans.algebra.Repeat.GroupingSetShapes;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 
+import com.google.common.base.Function;
+
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -75,6 +81,30 @@ public Expression visitSlotReference(SlotReference slotReference, DeepCopierCont
         }
     }
 
+    @Override
+    public Expression visitVirtualReference(VirtualSlotReference virtualSlotReference, DeepCopierContext context) {
+        Map<ExprId, ExprId> exprIdReplaceMap = context.exprIdReplaceMap;
+        ExprId newExprId;
+        if (exprIdReplaceMap.containsKey(virtualSlotReference.getExprId())) {
+            newExprId = exprIdReplaceMap.get(virtualSlotReference.getExprId());
+        } else {
+            newExprId = StatementScopeIdGenerator.newExprId();
+        }
+        // according to VirtualReference generating logic in Repeat.java
+        // generateVirtualGroupingIdSlot and generateVirtualSlotByFunction
+        Optional<GroupingScalarFunction> newOriginExpression = virtualSlotReference.getOriginExpression()
+                .map(func -> (GroupingScalarFunction) func.accept(this, context));
+        Function<GroupingSetShapes, List<Long>> newFunction = newOriginExpression
+                .<Function<GroupingSetShapes, List<Long>>>map(f -> f::computeVirtualSlotValue)
+                .orElseGet(() -> GroupingSetShapes::computeVirtualGroupingIdValue);
+        VirtualSlotReference newOne = new VirtualSlotReference(newExprId,
+                virtualSlotReference.getName(), virtualSlotReference.getDataType(),
+                virtualSlotReference.nullable(), virtualSlotReference.getQualifier(),
+                newOriginExpression, newFunction);
+        exprIdReplaceMap.put(virtualSlotReference.getExprId(), newOne.getExprId());
+        return newOne;
+    }
+
     @Override
     public Expression visitExistsSubquery(Exists exists, DeepCopierContext context) {
         LogicalPlan logicalPlan = LogicalPlanDeepCopier.INSTANCE.deepCopy(exists.getQueryPlan(), context);
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/Repeat.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/Repeat.java
@@ -60,13 +60,13 @@ default List<Expression> getGroupByExpressions() {
 
     static VirtualSlotReference generateVirtualGroupingIdSlot() {
         return new VirtualSlotReference(COL_GROUPING_ID, BigIntType.INSTANCE, Optional.empty(),
-                shapes -> shapes.computeVirtualGroupingIdValue());
+                GroupingSetShapes::computeVirtualGroupingIdValue);
     }
 
     static VirtualSlotReference generateVirtualSlotByFunction(GroupingScalarFunction function) {
         return new VirtualSlotReference(
                 generateVirtualSlotName(function), function.getDataType(), Optional.of(function),
-                shapes -> function.computeVirtualSlotValue(shapes));
+                function::computeVirtualSlotValue);
     }
 
     /**
@@ -175,7 +175,7 @@ default List<Set<Integer>> getGroupingSetsIndexesInOutput() {
                 if (index == null) {
                     throw new AnalysisException("Can not find grouping set expression in output: " + expression);
                 }
-                if (groupingSetsIndex.contains(index)) {
+                if (groupingSetIndex.contains(index)) {
                     throw new AnalysisException("expression duplicate in grouping set: " + expression);
                 }
                 groupingSetIndex.add(index);
@@ -228,14 +228,6 @@ public GroupingSetShapes(Set<Expression> flattenGroupingSetExpression, List<Grou
             this.shapes = ImmutableList.copyOf(shapes);
         }
 
-        public GroupingSetShape getGroupingSetShape(int index) {
-            return shapes.get(index);
-        }
-
-        public Expression getExpression(int index) {
-            return flattenGroupingSetExpression.get(index);
-        }
-
         // compute a long value that backend need to fill to the GROUPING_ID slot
         public List<Long> computeVirtualGroupingIdValue() {
             return shapes.stream()
diff --git a/regression-test/data/nereids_syntax_p0/cte.out b/regression-test/data/nereids_syntax_p0/cte.out
@@ -82,6 +82,17 @@ ASIA	1
 15
 29
 
+-- !cte_with_repeat --
+\N	\N	1
+\N	1	1
+\N	2	1
+\N	6	1
+1309892	\N	0
+1309892	1	0
+1309892	2	0
+1310179	\N	0
+1310179	6	0
+
 -- !test --
 1	2023-08-25 00:00:00	10	10
 1	2023-08-25 01:00:00	20	30
diff --git a/regression-test/suites/nereids_syntax_p0/cte.groovy b/regression-test/suites/nereids_syntax_p0/cte.groovy
@@ -253,7 +253,7 @@ suite("cte") {
                 ORDER BY dd.s_suppkey;
     """
 
-    sql "set experimental_enable_pipeline_engine=true"
+    sql "set enable_pipeline_engine=true"
 
     qt_cte14 """
             SELECT abs(dd.s_suppkey)
@@ -306,6 +306,13 @@ suite("cte") {
         notContains "MultiCastDataSinks"
     }
 
+    sql "WITH cte_0 AS ( SELECT 1 AS a ) SELECT * from cte_0 t1 LIMIT 10 UNION SELECT * from cte_0 t1 LIMIT 10"
+
+    qt_cte_with_repeat """
+        with cte_0 as (select lo_orderkey, lo_linenumber, grouping_id(lo_orderkey) as id from lineorder group by cube(lo_orderkey, lo_linenumber))
+        select * from cte_0 order by lo_orderkey, lo_linenumber, id
+    """
+
     qt_test """
         SELECT * FROM (
         WITH temptable as (
