From 4e50b2791ac349c4572a9ee9d0b7cd76503444de Mon Sep 17 00:00:00 2001
From: morrySnow <101034200+morrySnow@users.noreply.github.com>
Date: Tue, 12 Dec 2023 16:56:04 +0800
Subject: [PATCH] [fix](Nereids) pull up cte anchor should also pull up cte in
 apply (#28214)

---
 .../rules/rewrite/PullUpCteAnchor.java        | 32 ++++++++++++++++++-
 .../stats/ColumnStatsAdjustVisitor.java       |  2 +-
 .../nereids/trees/expressions/Exists.java     |  5 +++
 .../nereids/trees/expressions/InSubquery.java |  6 ++++
 .../nereids/trees/expressions/ListQuery.java  |  5 +++
 .../trees/expressions/ScalarSubquery.java     |  5 +++
 .../trees/expressions/SubqueryExpr.java       |  2 ++
 .../trees/plans/logical/LogicalApply.java     |  7 +++-
 8 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpCteAnchor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpCteAnchor.java
index c6461be271..4bfb36db81 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpCteAnchor.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpCteAnchor.java
@@ -18,12 +18,16 @@
 package org.apache.doris.nereids.rules.rewrite;
 
 import org.apache.doris.nereids.jobs.JobContext;
+import org.apache.doris.nereids.trees.expressions.SubqueryExpr;
 import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalApply;
 import org.apache.doris.nereids.trees.plans.logical.LogicalCTEAnchor;
 import org.apache.doris.nereids.trees.plans.logical.LogicalCTEProducer;
+import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter;
 import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 
 import java.util.List;
@@ -42,7 +46,7 @@ public class PullUpCteAnchor extends DefaultPlanRewriter<List<LogicalCTEProducer
         return rewriteRoot(plan, producers);
     }
 
-    private Plan rewriteRoot(Plan plan, List<LogicalCTEProducer<Plan>> producers) {
+    public Plan rewriteRoot(Plan plan, List<LogicalCTEProducer<Plan>> producers) {
         Plan root = plan.accept(this, producers);
         for (LogicalCTEProducer<Plan> producer : producers) {
             root = new LogicalCTEAnchor<>(producer.getCteId(), producer, root);
@@ -71,4 +75,30 @@ public class PullUpCteAnchor extends DefaultPlanRewriter<List<LogicalCTEProducer
         producers.addAll(childProducers);
         return newProducer;
     }
+
+    @Override
+    public Plan visitLogicalApply(LogicalApply<? extends Plan, ? extends Plan> apply,
+            List<LogicalCTEProducer<Plan>> producers) {
+        SubqueryExpr subqueryExpr = apply.getSubqueryExpr();
+        PullUpCteAnchor pullSubqueryExpr = new PullUpCteAnchor();
+        List<LogicalCTEProducer<Plan>> subqueryExprProducers = Lists.newArrayList();
+        Plan newPlanInExpr = pullSubqueryExpr.rewriteRoot(subqueryExpr.getQueryPlan(), subqueryExprProducers);
+        while (newPlanInExpr instanceof LogicalCTEAnchor) {
+            newPlanInExpr = ((LogicalCTEAnchor<?, ?>) newPlanInExpr).right();
+        }
+        SubqueryExpr newSubqueryExpr = subqueryExpr.withSubquery((LogicalPlan) newPlanInExpr);
+
+        Plan newApplyLeft = apply.left().accept(this, producers);
+
+        Plan applyRight = apply.right();
+        PullUpCteAnchor pullApplyRight = new PullUpCteAnchor();
+        List<LogicalCTEProducer<Plan>> childProducers = Lists.newArrayList();
+        Plan newApplyRight = pullApplyRight.rewriteRoot(applyRight, childProducers);
+        while (newApplyRight instanceof LogicalCTEAnchor) {
+            newApplyRight = ((LogicalCTEAnchor<?, ?>) newApplyRight).right();
+        }
+        producers.addAll(childProducers);
+        return apply.withSubqueryExprAndChildren(newSubqueryExpr,
+                ImmutableList.of(newApplyLeft, newApplyRight));
+    }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/ColumnStatsAdjustVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/ColumnStatsAdjustVisitor.java
index ff4290557d..eaf2ce3734 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/ColumnStatsAdjustVisitor.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/ColumnStatsAdjustVisitor.java
@@ -71,7 +71,7 @@ public class ColumnStatsAdjustVisitor extends ExpressionVisitor<ColumnStatistic,
     public ColumnStatistic visitCast(Cast cast, Statistics context) {
         ColumnStatistic colStats = context.findColumnStatistics(cast);
 
-        if (colStats != null) {
+        if (colStats != null && colStats.minExpr != null && colStats.maxExpr != null) {
             try {
                 DataType childNereidsType = cast.child().getDataType();
                 if (childNereidsType instanceof CharacterType) {
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Exists.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Exists.java
index e9f68d45a8..3d3bd17c70 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Exists.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Exists.java
@@ -102,4 +102,9 @@ public class Exists extends SubqueryExpr {
     public Expression withTypeCoercion(DataType dataType) {
         return this;
     }
+
+    @Override
+    public Exists withSubquery(LogicalPlan subquery) {
+        return new Exists(subquery, correlateSlots, typeCoercionExpr, isNot);
+    }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InSubquery.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InSubquery.java
index e04d1b2017..b04f9d3d78 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InSubquery.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/InSubquery.java
@@ -19,6 +19,7 @@ package org.apache.doris.nereids.trees.expressions;
 
 import org.apache.doris.nereids.exceptions.UnboundException;
 import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
+import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 import org.apache.doris.nereids.types.BooleanType;
 import org.apache.doris.nereids.types.DataType;
 
@@ -138,4 +139,9 @@ public class InSubquery extends SubqueryExpr {
                 : Optional.of(new Cast(listQuery.queryPlan.getOutput().get(0), dataType)),
             isNot);
     }
+
+    @Override
+    public InSubquery withSubquery(LogicalPlan subquery) {
+        return new InSubquery(compareExpr, listQuery.withSubquery(subquery), correlateSlots, typeCoercionExpr, isNot);
+    }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ListQuery.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ListQuery.java
index e01302ff39..214525d259 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ListQuery.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ListQuery.java
@@ -68,4 +68,9 @@ public class ListQuery extends SubqueryExpr {
                     ? Optional.of(queryPlan.getOutput().get(0))
                     : Optional.of(new Cast(queryPlan.getOutput().get(0), dataType)));
     }
+
+    @Override
+    public ListQuery withSubquery(LogicalPlan subquery) {
+        return new ListQuery(subquery, correlateSlots, typeCoercionExpr);
+    }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ScalarSubquery.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ScalarSubquery.java
index 48cf9491bf..88b354ae82 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ScalarSubquery.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ScalarSubquery.java
@@ -76,4 +76,9 @@ public class ScalarSubquery extends SubqueryExpr {
                     ? Optional.of(queryPlan.getOutput().get(0))
                     : Optional.of(new Cast(queryPlan.getOutput().get(0), dataType)));
     }
+
+    @Override
+    public ScalarSubquery withSubquery(LogicalPlan subquery) {
+        return new ScalarSubquery(subquery, correlateSlots, typeCoercionExpr);
+    }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SubqueryExpr.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SubqueryExpr.java
index 451fcad677..35d0e56647 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SubqueryExpr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SubqueryExpr.java
@@ -137,4 +137,6 @@ public abstract class SubqueryExpr extends Expression implements LeafExpression
     }
 
     public abstract Expression withTypeCoercion(DataType dataType);
+
+    public abstract SubqueryExpr withSubquery(LogicalPlan subquery);
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalApply.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalApply.java
index f6cf13dfc7..16e2e964ca 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalApply.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalApply.java
@@ -198,8 +198,13 @@ public class LogicalApply<LEFT_CHILD_TYPE extends Plan, RIGHT_CHILD_TYPE extends
                 .build();
     }
 
+    public LogicalApply<Plan, Plan> withSubqueryExprAndChildren(SubqueryExpr subqueryExpr, List<Plan> children) {
+        return new LogicalApply<>(correlationSlot, subqueryExpr, correlationFilter,
+                markJoinSlotReference, needAddSubOutputToProjects, inProject, children.get(0), children.get(1));
+    }
+
     @Override
-    public LogicalBinary<Plan, Plan> withChildren(List<Plan> children) {
+    public LogicalApply<Plan, Plan> withChildren(List<Plan> children) {
         Preconditions.checkArgument(children.size() == 2);
         return new LogicalApply<>(correlationSlot, subqueryExpr, correlationFilter,
                 markJoinSlotReference, needAddSubOutputToProjects, inProject,