optimize code usage

fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVCache.java

@@ -51,14 +51,18 @@ public class MTMVCache {
     // The materialized view plan which should be optimized by the same rules to query
     // and will remove top sink and unused sort
     private final Plan logicalPlan;
-    // The original plan of mv def sql
+    // The original rewritten plan of mv def sql
     private final Plan originalPlan;
+    // The analyzed plan of mv def sql, which is used by tableCollector,should not be optimized by rbo
+    private final Plan analyzedPlan;
     private final Statistics statistics;
     private final StructInfo structInfo;
 
-    public MTMVCache(Plan logicalPlan, Plan originalPlan, Statistics statistics, StructInfo structInfo) {
+    public MTMVCache(Plan logicalPlan, Plan originalPlan, Plan analyzedPlan,
+            Statistics statistics, StructInfo structInfo) {
         this.logicalPlan = logicalPlan;
         this.originalPlan = originalPlan;
+        this.analyzedPlan = analyzedPlan;
         this.statistics = statistics;
         this.structInfo = structInfo;
     }
@@ -71,6 +75,10 @@ public Plan getOriginalPlan() {
         return originalPlan;
     }
 
+    public Plan getAnalyzedPlan() {
+        return analyzedPlan;
+    }
+
     public Statistics getStatistics() {
         return statistics;
     }
@@ -117,7 +125,7 @@ public Plan visitLogicalResultSink(LogicalResultSink<? extends Plan> logicalResu
         Optional<StructInfo> structInfoOptional = MaterializationContext.constructStructInfo(mvPlan, originPlan,
                 planner.getCascadesContext(),
                 new BitSet());
-        return new MTMVCache(mvPlan, originPlan, needCost
+        return new MTMVCache(mvPlan, originPlan, planner.getAnalyzedPlan(), needCost
                 ? planner.getCascadesContext().getMemo().getRoot().getStatistics() : null,
                 structInfoOptional.orElseGet(() -> null));
     }

fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVPlanUtil.java

@@ -41,7 +41,6 @@
 import org.apache.doris.qe.ConnectContext;
 
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Sets;
 
 import java.util.List;
 import java.util.Optional;
@@ -97,32 +96,32 @@ public static MTMVRelation generateMTMVRelation(MTMV mtmv, ConnectContext ctx) {
         // Should not make table without data to empty relation when analyze the related table,
         // so add disable rules
         Plan plan = getAnalyzePlanBySql(mtmv.getQuerySql(), ctx);
-        return generateMTMVRelation(plan);
+        return generateMTMVRelation(plan, ctx);
     }
 
-    public static MTMVRelation generateMTMVRelation(Plan plan) {
-        return new MTMVRelation(getBaseTables(plan, true, true),
-                getBaseTables(plan, false, true),
-                getBaseViews(plan, true, true));
+    public static MTMVRelation generateMTMVRelation(Plan plan, ConnectContext connectContext) {
+        return new MTMVRelation(getBaseTables(plan, true, true, connectContext),
+                getBaseTables(plan, false, true, connectContext),
+                getBaseViews(plan, true, true, connectContext));
     }
 
     private static Set<BaseTableInfo> getBaseTables(Plan plan, boolean expandMaterializedView,
-            boolean expandView) {
+            boolean expandView, ConnectContext connectContext) {
         TableCollectorContext collectorContext =
                 new TableCollector.TableCollectorContext(
                         com.google.common.collect.Sets
-                                .newHashSet(TableType.values()), expandMaterializedView, expandView);
+                                .newHashSet(TableType.values()), connectContext, expandMaterializedView, expandView);
         plan.accept(TableCollector.INSTANCE, collectorContext);
         Set<TableIf> collectedTables = collectorContext.getCollectedTables();
         return transferTableIfToInfo(collectedTables);
     }
 
     private static Set<BaseTableInfo> getBaseViews(Plan plan, boolean expandMaterializedView,
-            boolean expandView) {
+            boolean expandView, ConnectContext connectContext) {
         TableCollectorContext collectorContext =
                 new TableCollector.TableCollectorContext(
                         com.google.common.collect.Sets
-                                .newHashSet(TableType.VIEW), expandMaterializedView, expandView);
+                                .newHashSet(TableType.VIEW), connectContext, expandMaterializedView, expandView);
         plan.accept(TableCollector.INSTANCE, collectorContext);
         Set<TableIf> collectedTables = collectorContext.getCollectedTables();
         return transferTableIfToInfo(collectedTables);

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/InitMaterializationContextHook.java

@@ -80,13 +80,13 @@ public void initMaterializationContext(CascadesContext cascadesContext) {
      */
     protected void doInitMaterializationContext(CascadesContext cascadesContext) {
         // Only collect the table or mv which query use directly, to avoid useless mv partition in rewrite
+        // Keep use one connection context when in query, if new connect context,
+        // the ConnectionContext.get() will change
         TableCollectorContext collectorContext = new TableCollectorContext(Sets.newHashSet(),
+                cascadesContext.getConnectContext(),
                 false, true);
         try {
             Plan rewritePlan = cascadesContext.getRewritePlan();
-            // Keep use one connection context when in query, if new connect context,
-            // the ConnectionContext.get() will change
-            collectorContext.setConnectContext(cascadesContext.getConnectContext());
             rewritePlan.accept(TableCollector.INSTANCE, collectorContext);
         } catch (Exception e) {
             LOG.warn(String.format("MaterializationContext init table collect fail, current queryId is %s",

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewUtils.java

@@ -343,7 +343,7 @@ public Plan visitLogicalResultSink(LogicalResultSink<? extends Plan> logicalResu
                     ImmutableList.of(Rewriter.custom(RuleType.ELIMINATE_SORT, EliminateSort::new))).execute();
             return childContext.getRewritePlan();
         }, mvPlan, originPlan);
-        return new MTMVCache(mvPlan, originPlan,
+        return new MTMVCache(mvPlan, originPlan, planner.getAnalyzedPlan(),
                 planner.getCascadesContext().getMemo().getRoot().getStatistics(), null);
     }
 

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateMTMVInfo.java

@@ -252,8 +252,6 @@ public void analyzeQuery(ConnectContext ctx, Map<String, String> mvProperties) t
         NereidsPlanner planner = new NereidsPlanner(statementContext);
         // this is for expression column name infer when not use alias
         LogicalSink<Plan> logicalSink = new UnboundResultSink<>(logicalQuery);
-        // must disable constant folding by be, because be constant folding may return wrong type
-        ctx.getSessionVariable().disableConstantFoldingByBEOnce();
         // Should not make table without data to empty relation when analyze the related table,
         // so add disable rules
         Set<String> tempDisableRules = ctx.getSessionVariable().getDisableNereidsRuleNames();
@@ -326,7 +324,7 @@ private void analyzeKeys() {
 
     // Should use analyzed plan for collect views and tables
     private void getRelation(NereidsPlanner planner) {
-        this.relation = MTMVPlanUtil.generateMTMVRelation(planner.getAnalyzedPlan());
+        this.relation = MTMVPlanUtil.generateMTMVRelation(planner.getAnalyzedPlan(), planner.getConnectContext());
     }
 
     private PartitionDesc generatePartitionDesc(ConnectContext ctx) {

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/TableCollector.java

@@ -21,8 +21,8 @@
 import org.apache.doris.catalog.TableIf;
 import org.apache.doris.catalog.TableIf.TableType;
 import org.apache.doris.catalog.View;
+import org.apache.doris.common.AnalysisException;
 import org.apache.doris.mtmv.MTMVCache;
-import org.apache.doris.mtmv.MTMVPlanUtil;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalView;
@@ -88,9 +88,15 @@ private void expandMvAndCollect(MTMV mtmv, TableCollectorContext context) {
             return;
         }
         // Make sure use only one connection context when in query to avoid ConnectionContext.get() wrong
-        MTMVCache expandedMv = MTMVCache.from(mtmv, context.getConnectContext() == null
-                ? MTMVPlanUtil.createMTMVContext(mtmv) : context.getConnectContext(), false);
-        expandedMv.getLogicalPlan().accept(this, context);
+        MTMVCache expandedMvCache;
+        try {
+            expandedMvCache = mtmv.getOrGenerateCache(context.getConnectContext());
+        } catch (AnalysisException exception) {
+            LOG.warn(String.format("expandMvAndCollect getOrGenerateCache fail, mtmv name is %s", mtmv.getName()),
+                    exception);
+            expandedMvCache = MTMVCache.from(mtmv, context.getConnectContext(), false);
+        }
+        expandedMvCache.getAnalyzedPlan().accept(this, context);
     }
 
     /**
@@ -103,11 +109,13 @@ public static final class TableCollectorContext {
         // if expand the mv or not
         private final boolean expandMaterializedView;
         private final boolean expandView;
-        private ConnectContext connectContext;
+        private final ConnectContext connectContext;
 
-        public TableCollectorContext(Set<TableType> targetTableTypes, boolean expandMaterializedView,
+        public TableCollectorContext(Set<TableType> targetTableTypes,
+                ConnectContext connectContext, boolean expandMaterializedView,
                 boolean expandView) {
             this.targetTableTypes = targetTableTypes;
+            this.connectContext = connectContext;
             this.expandMaterializedView = expandMaterializedView;
             this.expandView = expandView;
         }
@@ -120,16 +128,12 @@ public Set<TableType> getTargetTableTypes() {
             return targetTableTypes;
         }
 
-        public boolean isExpandMaterializedView() {
-            return expandMaterializedView;
-        }
-
         public ConnectContext getConnectContext() {
             return connectContext;
         }
 
-        public void setConnectContext(ConnectContext connectContext) {
-            this.connectContext = connectContext;
+        public boolean isExpandMaterializedView() {
+            return expandMaterializedView;
         }
 
         public boolean isExpandView() {