branch-4.0: [test](mtmv) Add ut test for session variable guard when generate mtmv cache #58863

fe/fe-core/src/test/java/org/apache/doris/nereids/mv/IdStatisticsMapTest.java

@@ -67,6 +67,15 @@ public boolean canBeCandidate() {
                 return true;
             }
         };
+
+        new MockUp<MTMV>() {
+            @Mock
+            public boolean canBeCandidate() {
+                return true;
+            }
+        };
+        connectContext.getState().setIsQuery(true);
+
         connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
         connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
         connectContext.getSessionVariable().setPreMaterializedViewRewriteStrategy(PreRewriteStrategy.NOT_IN_RBO.name());
@@ -124,6 +133,8 @@ public boolean canBeCandidate() {
                 return true;
             }
         };
+        connectContext.getState().setIsQuery(true);
+
         connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
         connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
         connectContext.getSessionVariable().setPreMaterializedViewRewriteStrategy(

fe/fe-core/src/test/java/org/apache/doris/nereids/mv/MTMVCacheTest.java

@@ -0,0 +1,132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.nereids.mv;
+
+import org.apache.doris.catalog.MTMV;
+import org.apache.doris.mtmv.MTMVCache;
+import org.apache.doris.mtmv.MTMVRelationManager;
+import org.apache.doris.nereids.CascadesContext;
+import org.apache.doris.nereids.rules.exploration.mv.AsyncMaterializationContext;
+import org.apache.doris.nereids.rules.exploration.mv.MaterializationContext;
+import org.apache.doris.nereids.sqltest.SqlTestBase;
+import org.apache.doris.nereids.trees.expressions.SessionVarGuardExpr;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
+import org.apache.doris.nereids.util.PlanChecker;
+import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.SessionVariable;
+import org.apache.doris.qe.SqlModeHelper;
+
+import mockit.Mock;
+import mockit.MockUp;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.BitSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+/**
+ * Relevant test case about mtmv cache.
+ */
+public class MTMVCacheTest extends SqlTestBase {
+
+    @Test
+    void testMTMVCacheIsCorrect() throws Exception {
+        connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION");
+        BitSet disableNereidsRules = connectContext.getSessionVariable().getDisableNereidsRules();
+        new MockUp<SessionVariable>() {
+            @Mock
+            public BitSet getDisableNereidsRules() {
+                return disableNereidsRules;
+            }
+        };
+        new MockUp<MTMVRelationManager>() {
+            @Mock
+            public boolean isMVPartitionValid(MTMV mtmv, ConnectContext ctx, boolean forceConsistent,
+                                              Map<List<String>, Set<String>> queryUsedPartitions) {
+                return true;
+            }
+        };
+
+        new MockUp<MTMV>() {
+            @Mock
+            public boolean canBeCandidate() {
+                return true;
+            }
+        };
+        connectContext.getState().setIsQuery(true);
+
+        connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
+        connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
+        createMvByNereids("create materialized view mv1 BUILD IMMEDIATE REFRESH COMPLETE ON MANUAL\n"
+                + "        DISTRIBUTED BY RANDOM BUCKETS 1\n"
+                + "        PROPERTIES ('replication_num' = '1') \n"
+                + "        as select T1.id, sum(score) from T1 group by T1.id;");
+        CascadesContext c1 = createCascadesContext(
+                "select T1.id, sum(score) from T1 group by T1.id;",
+                connectContext
+        );
+        PlanChecker.from(c1)
+                .analyze()
+                .rewrite()
+                .optimize()
+                .printlnBestPlanTree();
+        List<MaterializationContext> normalMaterializationContexts = c1.getMaterializationContexts();
+        Assertions.assertEquals(1, normalMaterializationContexts.size());
+
+        MTMV mtmv = ((AsyncMaterializationContext) normalMaterializationContexts.get(0)).getMtmv();
+        MTMVCache cacheWithoutGuard = mtmv.getOrGenerateCache(connectContext);
+
+        Optional<LogicalAggregate<? extends Plan>> aggregate = cacheWithoutGuard.getAllRulesRewrittenPlanAndStructInfo().key()
+                .collectFirst(LogicalAggregate.class::isInstance);
+        Assertions.assertTrue(aggregate.isPresent());
+        // should not contain SessionVarGuardExpr
+        Assertions.assertTrue(aggregate.get().getOutputExpressions().stream()
+                .noneMatch(expr -> expr.containsType(SessionVarGuardExpr.class)));
+
+        // set guard check session var
+        connectContext.getSessionVariable().setSqlMode(SqlModeHelper.MODE_NO_UNSIGNED_SUBTRACTION);
+        CascadesContext c2 = createCascadesContext(
+                "select T1.id, sum(score) from T1 group by T1.id;",
+                connectContext
+        );
+        connectContext.getState().setIsQuery(true);
+        PlanChecker.from(c2)
+                .analyze()
+                .rewrite()
+                .optimize()
+                .printlnBestPlanTree();
+
+        List<MaterializationContext> sessionChangedMaterializationContexts = c2.getMaterializationContexts();
+        Assertions.assertEquals(1, sessionChangedMaterializationContexts.size());
+
+        MTMV mvWithGuard = ((AsyncMaterializationContext) sessionChangedMaterializationContexts.get(0)).getMtmv();
+        MTMVCache cacheWithGuard = mvWithGuard.getOrGenerateCache(connectContext);
+
+        aggregate = cacheWithGuard.getAllRulesRewrittenPlanAndStructInfo().key()
+                .collectFirst(LogicalAggregate.class::isInstance);
+        Assertions.assertTrue(aggregate.isPresent());
+        // should contain SessionVarGuardExpr
+        Assertions.assertTrue(aggregate.get().getOutputExpressions().stream()
+                .anyMatch(expr -> expr.containsType(SessionVarGuardExpr.class)));
+        dropMvByNereids("drop materialized view mv1");
+    }
+}

fe/fe-core/src/test/java/org/apache/doris/nereids/mv/MtmvCacheNewConnectContextTest.java

@@ -18,7 +18,6 @@
 package org.apache.doris.nereids.mv;
 
 import org.apache.doris.catalog.MTMV;
-import org.apache.doris.mtmv.BaseTableInfo;
 import org.apache.doris.mtmv.MTMVRelationManager;
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.sqltest.SqlTestBase;
@@ -32,6 +31,7 @@
 import org.junit.jupiter.api.Test;
 
 import java.util.BitSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -55,11 +55,18 @@ public BitSet getDisableNereidsRules() {
         };
         new MockUp<MTMVRelationManager>() {
             @Mock
-            public boolean isMVPartitionValid(MTMV mtmv, ConnectContext ctx, boolean isMVPartitionValid,
-                    Map<BaseTableInfo, Set<String>> queryUsedRelatedTablePartitionsMap) {
+            public boolean isMVPartitionValid(MTMV mtmv, ConnectContext ctx, boolean forceConsistent,
+                                              Map<List<String>, Set<String>> queryUsedPartitions) {
                 return true;
             }
         };
+        new MockUp<MTMV>() {
+            @Mock
+            public boolean canBeCandidate() {
+                return true;
+            }
+        };
+        connectContext.getState().setIsQuery(true);
         connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
         connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
 

fe/fe-core/src/test/java/org/apache/doris/nereids/mv/MvTableIdIsLongTest.java

@@ -63,6 +63,7 @@ public boolean canBeCandidate() {
                 return true;
             }
         };
+        connectContext.getState().setIsQuery(true);
         connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
         connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
         createMvByNereids("create materialized view mv1 BUILD IMMEDIATE REFRESH COMPLETE ON MANUAL\n"

fe/fe-core/src/test/java/org/apache/doris/nereids/mv/OptimizeGetAvailableMvsTest.java

@@ -18,6 +18,7 @@
 package org.apache.doris.nereids.mv;
 
 import org.apache.doris.catalog.DistributionInfo;
+import org.apache.doris.catalog.MTMV;
 import org.apache.doris.catalog.MaterializedIndex;
 import org.apache.doris.catalog.MaterializedIndex.IndexState;
 import org.apache.doris.catalog.OlapTable;
@@ -112,6 +113,14 @@ public List<Long> getSelectedPartitionIds() {
             }
         };
 
+        new MockUp<MTMV>() {
+            @Mock
+            public boolean canBeCandidate() {
+                return true;
+            }
+        };
+        connectContext.getState().setIsQuery(true);
+
         connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
         connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
         createMvByNereids("create materialized view mv1 "

fe/fe-core/src/test/java/org/apache/doris/nereids/mv/PointQueryShouldNotMvRewriteTest.java

@@ -33,6 +33,8 @@
 
 import java.util.BitSet;
 import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -53,7 +55,7 @@ public BitSet getDisableNereidsRules() {
         new MockUp<MTMVRelationManager>() {
             @Mock
             public boolean isMVPartitionValid(MTMV mtmv, ConnectContext ctx, boolean forceConsistent,
-                    Set<String> relatedPartitions) {
+                                              Map<List<String>, Set<String>> queryUsedPartitions) {
                 return true;
             }
         };
@@ -63,6 +65,7 @@ public boolean canBeCandidate() {
                 return true;
             }
         };
+        connectContext.getState().setIsQuery(true);
         connectContext.getSessionVariable().enableMaterializedViewRewrite = true;
         connectContext.getSessionVariable().enableMaterializedViewNestRewrite = true;
 

regression-test/suites/nereids_rules_p0/mv/union_all_compensate/union_all_compensate.groovy

@@ -125,6 +125,9 @@ suite("union_all_compensate") {
         sql """analyze table test_table1 with sync"""
         sql """analyze table test_table2 with sync"""
 
+        sql """alter table test_table1 modify column num set stats ('row_count'='20');"""
+        sql """alter table test_table2 modify column num set stats ('row_count'='16');"""
+
         // Aggregate, scalar aggregate, should not compensate union all
         sql """ DROP MATERIALIZED VIEW IF EXISTS test_agg_mv"""
         sql """
@@ -194,9 +197,6 @@ suite("union_all_compensate") {
         sql "set enable_sql_cache=true"
         order_qt_query1_1_after_use_sql_cache "${query1_0}"
 
-        sql """alter table test_table1 modify column num set stats ('row_count'='20');"""
-        sql """alter table test_table2 modify column num set stats ('row_count'='16');"""
-
 
         // Aggregate, if query group by expression doesn't use the partition column, but the invalid partition is in the
         // grace_period, should not compensate union all, but should rewritten successfully

regression-test/suites/nereids_rules_p0/mv/union_rewrite/partition_curd_union_rewrite.groovy

@@ -78,9 +78,6 @@ suite ("partition_curd_union_rewrite") {
     );
     """
 
-    sql """alter table orders modify column o_comment set stats ('row_count'='3');"""
-    sql """alter table lineitem modify column l_comment set stats ('row_count'='3');"""
-
     sql"""
     insert into orders values 
     (1, 1, 'ok', 99.5, '2023-10-17', 'a', 'b', 1, 'yy'),
@@ -107,6 +104,9 @@ suite ("partition_curd_union_rewrite") {
     (3, 2, 3, 6, 7.5, 8.5, 9.5, 10.5, 'k', 'o', '2023-10-19', '2023-10-19', '2023-10-19', 'c', 'd', 'xxxxxxxxx');
     """
 
+    sql """alter table orders modify column o_comment set stats ('row_count'='9');"""
+    sql """alter table lineitem modify column l_comment set stats ('row_count'='9');"""
+
     sql """analyze table orders with sync;"""
     sql """analyze table lineitem with sync;"""