[fix](Nereids) storage later agg rule process agg children by mistake (…

…#26101) update Project#findProject agg function's children could be any expression rather than only slot. we use Project#findProject to process them. But this util could only process slot. This PR update this util to let it could process all type expression.
apache · Nov 6, 2023 · c0ed5f7 · c0ed5f7
1 parent 1e2a614
commit c0ed5f7
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 20 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
@@ -397,7 +397,7 @@ private LogicalAggregate<? extends Plan> storageLayerAggregate(
 
         if (project != null) {
             argumentsOfAggregateFunction = Project.findProject(
-                        (List<SlotReference>) (List) argumentsOfAggregateFunction, project.getProjects())
+                        argumentsOfAggregateFunction, project.getProjects())
                     .stream()
                     .map(p -> p instanceof Alias ? p.child(0) : p)
                     .collect(ImmutableList.toImmutableList());
@@ -431,8 +431,8 @@ private LogicalAggregate<? extends Plan> storageLayerAggregate(
         Set<SlotReference> aggUsedSlots =
                 ExpressionUtils.collect(argumentsOfAggregateFunction, SlotReference.class::isInstance);
 
-        List<SlotReference> usedSlotInTable = (List<SlotReference>) (List) Project.findProject(aggUsedSlots,
-                (List<NamedExpression>) (List) logicalScan.getOutput());
+        List<SlotReference> usedSlotInTable = (List<SlotReference>) Project.findProject(aggUsedSlots,
+                logicalScan.getOutput());
 
         for (SlotReference slot : usedSlotInTable) {
             Column column = slot.getColumn().get();

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateAggregate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateAggregate.java
@@ -41,8 +41,8 @@ public Rule build() {
             if (!onlyHasSlots(outerAgg.getOutputExpressions())) {
                 return outerAgg;
             }
-            List<NamedExpression> prunedInnerAggOutput = Project.findProject(outerAgg.getOutputSet(),
-                    innerAgg.getOutputExpressions());
+            List<NamedExpression> prunedInnerAggOutput = (List<NamedExpression>) Project.findProject(
+                    outerAgg.getOutputSet(), innerAgg.getOutputExpressions());
             return innerAgg.withAggOutput(prunedInnerAggOutput);
         }).toRule(RuleType.ELIMINATE_AGGREGATE);
     }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/Project.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/algebra/Project.java
@@ -23,9 +23,9 @@
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.NamedExpression;
 import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.util.ExpressionUtils;
 import org.apache.doris.nereids.util.PlanUtils;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import java.util.Collection;
@@ -78,22 +78,25 @@ default List<NamedExpression> mergeProjections(Project childProject) {
     /**
      * find projects, if not found the slot, then throw AnalysisException
      */
-    static List<NamedExpression> findProject(
-            Collection<? extends Slot> slotReferences,
-            List<NamedExpression> projects) throws AnalysisException {
+    static List<? extends Expression> findProject(
+            Collection<? extends Expression> expressions,
+            List<? extends NamedExpression> projects) throws AnalysisException {
         Map<ExprId, NamedExpression> exprIdToProject = projects.stream()
-                .collect(ImmutableMap.toImmutableMap(p -> p.getExprId(), p -> p));
+                .collect(ImmutableMap.toImmutableMap(NamedExpression::getExprId, p -> p));
 
-        return slotReferences.stream()
-                .map(slot -> {
-                    ExprId exprId = slot.getExprId();
-                    NamedExpression project = exprIdToProject.get(exprId);
-                    if (project == null) {
-                        throw new AnalysisException("ExprId " + slot.getExprId() + " no exists in " + projects);
+        return ExpressionUtils.rewriteDownShortCircuit(expressions,
+                expr -> {
+                    if (expr instanceof Slot) {
+                        Slot slot = (Slot) expr;
+                        ExprId exprId = slot.getExprId();
+                        NamedExpression project = exprIdToProject.get(exprId);
+                        if (project == null) {
+                            throw new AnalysisException("ExprId " + slot.getExprId() + " no exists in " + projects);
+                        }
+                        return project;
                     }
-                    return project;
-                })
-                .collect(ImmutableList.toImmutableList());
+                    return expr;
+                });
     }
 }
 

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/ExpressionUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/ExpressionUtils.java
@@ -282,7 +282,7 @@ public static Set<Expression> replace(Set<Expression> exprs,
     }
 
     public static <E extends Expression> List<E> rewriteDownShortCircuit(
-            List<E> exprs, Function<Expression, Expression> rewriteFunction) {
+            Collection<E> exprs, Function<Expression, Expression> rewriteFunction) {
         return exprs.stream()
                 .map(expr -> (E) expr.rewriteDownShortCircuit(rewriteFunction))
                 .collect(ImmutableList.toImmutableList());

diff --git a/regression-test/suites/query_p0/aggregate/aggregate_decimal256.groovy b/regression-test/suites/query_p0/aggregate/aggregate_decimal256.groovy
@@ -17,6 +17,7 @@
 
 suite("aggregate_decimal256") {
     sql "set enable_nereids_planner = true;"
+    sql "set enable_fallback_to_original_planner = false;"
     sql "set enable_decimal256 = true;"
     sql "drop table if exists test_aggregate_decimal256_sum;"
     sql """ create table test_aggregate_decimal256_sum(k1 int, v1 decimal(38, 6), v2 decimal(38, 6))