From 9f5c97e7bf5f37ee15e23ae8d5f111893f135ac9 Mon Sep 17 00:00:00 2001
From: seawinde <daydayup005@yeah.net>
Date: Fri, 22 Dec 2023 16:59:22 +0800
Subject: [PATCH] [fix](nereids) Fix query rewrite fail when query cache build
 quickly

---
 .../mv/AbstractMaterializedViewRule.java      | 21 ++-----------------
 .../mv/MaterializationContext.java            |  8 +++++++
 2 files changed, 10 insertions(+), 19 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java
index 6dcf1885563691..54cc61b708dd9d 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java
@@ -26,7 +26,6 @@
 import org.apache.doris.catalog.TableIf;
 import org.apache.doris.common.AnalysisException;
 import org.apache.doris.mtmv.BaseTableInfo;
-import org.apache.doris.mtmv.MTMVCache;
 import org.apache.doris.mtmv.MTMVPartitionInfo;
 import org.apache.doris.mtmv.MTMVUtil;
 import org.apache.doris.nereids.CascadesContext;
@@ -104,13 +103,8 @@ protected List<Plan> rewrite(Plan queryPlan, CascadesContext cascadesContext) {
                 logger.debug(currentClassName + " this group is already rewritten so skip");
                 continue;
             }
-            MTMV mtmv = materializationContext.getMTMV();
-            MTMVCache mtmvCache = getCacheFromMTMV(mtmv);
-            if (mtmvCache == null) {
-                logger.warn(currentClassName + " mv cache is null so return");
-                return rewriteResults;
-            }
-            List<StructInfo> viewStructInfos = extractStructInfo(mtmvCache.getLogicalPlan(), cascadesContext);
+            List<StructInfo> viewStructInfos = extractStructInfo(materializationContext.getMvPlan(),
+                    cascadesContext);
             if (viewStructInfos.size() > 1) {
                 // view struct info should only have one
                 logger.warn(currentClassName + " the num of view struct info is more then one so return");
@@ -289,17 +283,6 @@ protected boolean checkPartitionIsValid(
                 && relatedTalbeValidSet.containsAll(relatedTableSelectedPartitionToCheck);
     }
 
-    private MTMVCache getCacheFromMTMV(MTMV mtmv) {
-        MTMVCache cache;
-        try {
-            cache = mtmv.getOrGenerateCache();
-        } catch (AnalysisException analysisException) {
-            logger.warn(this.getClass().getSimpleName() + " get mtmv cache analysisException", analysisException);
-            return null;
-        }
-        return cache;
-    }
-
     /**
      * Rewrite query by view, for aggregate or join rewriting should be different inherit class implementation
      */
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java
index 19b1f5be8bb10a..153688ecc2d3dd 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializationContext.java
@@ -51,6 +51,8 @@ public class MaterializationContext {
     // generate form mv scan plan
     private ExpressionMapping mvExprToMvScanExprMapping;
     private boolean available = true;
+    // the mv plan from cache at present, record it to make sure query rewrite by mv is right when cache change.
+    private Plan mvPlan;
 
     /**
      * MaterializationContext, this contains necessary info for query rewriting by mv
@@ -81,6 +83,8 @@ public MaterializationContext(MTMV mtmv,
                         mtmvCache.getMvOutputExpressions(),
                         mtmvCache.getLogicalPlan()),
                 mvScanPlan.getExpressions());
+        // copy the plan from cache, which the plan in cache may change
+        this.mvPlan = mtmvCache.getLogicalPlan();
     }
 
     public Set<GroupId> getMatchedGroups() {
@@ -119,6 +123,10 @@ public boolean isAvailable() {
         return available;
     }
 
+    public Plan getMvPlan() {
+        return mvPlan;
+    }
+
     /**
      * MaterializationContext fromMaterializedView
      */