modified based on comment

apache · Oct 31, 2023 · 143ec8a · 143ec8a
1 parent 2ba7dfe
commit 143ec8a
Showing 1 changed file with 20 additions and 4 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubExprAnalyzer.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/SubExprAnalyzer.java
@@ -31,13 +31,15 @@
 import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
 import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral;
 import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
+import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
 import org.apache.doris.nereids.trees.plans.logical.LogicalLimit;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 
+import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
@@ -221,10 +223,24 @@ public boolean hasGroupBy() {
         }
 
         public boolean hasCorrelatedSlotsUnderAgg() {
-            if (!correlatedSlots.isEmpty()) {
-                return logicalPlan.anyMatch(planTreeNode -> planTreeNode instanceof LogicalAggregate
-                        && ((LogicalAggregate<?>) planTreeNode)
-                                .containsSlots(ImmutableSet.copyOf(correlatedSlots)));
+            return correlatedSlots.isEmpty() ? false
+                    : findAggContainsCorrelatedSlots(logicalPlan, ImmutableSet.copyOf(correlatedSlots));
+        }
+
+        private boolean findAggContainsCorrelatedSlots(Plan rootPlan, ImmutableSet<Slot> slots) {
+            ArrayDeque<Plan> planQueue = new ArrayDeque<>();
+            planQueue.add(rootPlan);
+            while (!planQueue.isEmpty()) {
+                Plan plan = planQueue.poll();
+                if (plan instanceof LogicalAggregate) {
+                    if (plan.containsSlots(slots)) {
+                        return true;
+                    }
+                } else {
+                    for (Plan child : plan.children()) {
+                        planQueue.add(child);
+                    }
+                }
             }
             return false;
         }