[fix](mtmv) Disable sql_limit variable when query rewrite by material…

…ize view
seawinde · Aug 29, 2024 · 1194638 · 1194638
1 parent 6b81529
commit 1194638
Show file tree

Hide file tree

Showing 7 changed files with 372 additions and 7 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVCache.java b/fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVCache.java
@@ -39,6 +39,7 @@
 import org.apache.doris.statistics.Statistics;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 
 import java.util.BitSet;
 import java.util.Optional;
@@ -112,7 +113,7 @@ public Plan visitLogicalResultSink(LogicalResultSink<? extends Plan> logicalResu
             Rewriter.getCteChildrenRewriter(childContext,
                     ImmutableList.of(Rewriter.custom(RuleType.ELIMINATE_SORT, EliminateSort::new))).execute();
             return childContext.getRewritePlan();
-        }, mvPlan, originPlan);
+        }, mvPlan, originPlan, ImmutableSet.of(RuleType.ADD_DEFAULT_LIMIT));
         // Construct structInfo once for use later
         Optional<StructInfo> structInfoOptional = MaterializationContext.constructStructInfo(mvPlan, originPlan,
                 planner.getCascadesContext(),

diff --git a/.../org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewAggregateRule.java b/.../org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewAggregateRule.java
@@ -21,6 +21,7 @@
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.jobs.executor.Rewriter;
 import org.apache.doris.nereids.properties.DataTrait;
+import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.rules.analysis.NormalizeRepeat;
 import org.apache.doris.nereids.rules.exploration.mv.AbstractMaterializedViewAggregateRule.AggregateExpressionRewriteContext.ExpressionRewriteMode;
 import org.apache.doris.nereids.rules.exploration.mv.StructInfo.PlanCheckContext;
@@ -513,7 +514,7 @@ private static boolean isGroupByEqualsAfterGroupByEliminate(Set<Expression> quer
                     Rewriter.getCteChildrenRewriter(childContext,
                             ImmutableList.of(Rewriter.topDown(new EliminateGroupByKey()))).execute();
                     return childContext.getRewritePlan();
-                }, project, project);
+                }, project, project, ImmutableSet.of(RuleType.ADD_DEFAULT_LIMIT));
 
         Optional<LogicalAggregate<Plan>> aggreagateOptional =
                 rewrittenPlan.collectFirst(LogicalAggregate.class::isInstance);

diff --git a/...main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java b/...main/java/org/apache/doris/nereids/rules/exploration/mv/AbstractMaterializedViewRule.java
@@ -31,6 +31,7 @@
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.jobs.executor.Rewriter;
 import org.apache.doris.nereids.properties.LogicalProperties;
+import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.rules.exploration.ExplorationRuleFactory;
 import org.apache.doris.nereids.rules.exploration.mv.Predicates.SplitPredicate;
 import org.apache.doris.nereids.rules.exploration.mv.StructInfo.PartitionRemover;
@@ -265,7 +266,7 @@ protected List<Plan> doRewrite(StructInfo queryStructInfo, CascadesContext casca
                     childContext -> {
                         Rewriter.getWholeTreeRewriter(childContext).execute();
                         return childContext.getRewritePlan();
-                    }, rewrittenPlan, queryPlan);
+                    }, rewrittenPlan, queryPlan, ImmutableSet.of(RuleType.ADD_DEFAULT_LIMIT));
             if (rewrittenPlan == null) {
                 continue;
             }

diff --git a/...re/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewUtils.java b/...re/src/main/java/org/apache/doris/nereids/rules/exploration/mv/MaterializedViewUtils.java
@@ -73,6 +73,7 @@
 import org.apache.doris.nereids.util.ExpressionUtils;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.qe.OriginStatement;
+import org.apache.doris.qe.SessionVariable;
 
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
@@ -252,7 +253,8 @@ public static Plan generateMvScanPlan(OlapTable table, long indexId,
     public static Plan rewriteByRules(
             CascadesContext cascadesContext,
             Function<CascadesContext, Plan> planRewriter,
-            Plan rewrittenPlan, Plan originPlan) {
+            Plan rewrittenPlan, Plan originPlan,
+            Set<RuleType> disableRuleSet) {
         if (originPlan == null || rewrittenPlan == null) {
             return null;
         }
@@ -267,7 +269,26 @@ public static Plan rewriteByRules(
         CascadesContext rewrittenPlanContext = CascadesContext.initContext(
                 cascadesContext.getStatementContext(), rewrittenPlan,
                 cascadesContext.getCurrentJobContext().getRequiredProperties());
-        rewrittenPlan = planRewriter.apply(rewrittenPlanContext);
+        // Tmp old disable rule variable
+        Set<String> oldDisableRuleNames = rewrittenPlanContext.getStatementContext().getConnectContext()
+                .getSessionVariable()
+                .getDisableNereidsRuleNames();
+        if (!disableRuleSet.isEmpty()) {
+            Set<String> disableRules = disableRuleSet.stream().map(RuleType::name).collect(Collectors.toSet());
+            rewrittenPlanContext.getStatementContext().getConnectContext().getSessionVariable()
+                    .setDisableNereidsRules(String.join(",", disableRules));
+            rewrittenPlanContext.getStatementContext().invalidCache(SessionVariable.DISABLE_NEREIDS_RULES);
+        }
+        try {
+            rewrittenPlan = planRewriter.apply(rewrittenPlanContext);
+        } finally {
+            // Recover old disable rules variable
+            if (!disableRuleSet.isEmpty()) {
+                rewrittenPlanContext.getStatementContext().getConnectContext().getSessionVariable()
+                        .setDisableNereidsRules(String.join(",", oldDisableRuleNames));
+                rewrittenPlanContext.getStatementContext().invalidCache(SessionVariable.DISABLE_NEREIDS_RULES);
+            }
+        }
         Map<ExprId, Slot> exprIdToNewRewrittenSlot = Maps.newLinkedHashMap();
         for (Slot slot : rewrittenPlan.getOutput()) {
             exprIdToNewRewrittenSlot.put(slot.getExprId(), slot);
@@ -325,7 +346,7 @@ public Plan visitLogicalResultSink(LogicalResultSink<? extends Plan> logicalResu
             Rewriter.getCteChildrenRewriter(childContext,
                     ImmutableList.of(Rewriter.custom(RuleType.ELIMINATE_SORT, EliminateSort::new))).execute();
             return childContext.getRewritePlan();
-        }, mvPlan, originPlan);
+        }, mvPlan, originPlan, ImmutableSet.of(RuleType.ADD_DEFAULT_LIMIT));
         return new MTMVCache(mvPlan, originPlan,
                 planner.getCascadesContext().getMemo().getRoot().getStatistics(), null);
     }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/StructInfo.java
@@ -28,6 +28,7 @@
 import org.apache.doris.nereids.jobs.joinorder.hypergraph.node.StructInfoNode;
 import org.apache.doris.nereids.memo.Group;
 import org.apache.doris.nereids.memo.GroupExpression;
+import org.apache.doris.nereids.rules.RuleType;
 import org.apache.doris.nereids.rules.exploration.mv.Predicates.SplitPredicate;
 import org.apache.doris.nereids.trees.copier.DeepCopierContext;
 import org.apache.doris.nereids.trees.copier.LogicalPlanDeepCopier;
@@ -782,7 +783,7 @@ public static Pair<Plan, Boolean> addFilterOnTableScan(Plan queryPlan, Map<BaseT
         return Pair.of(MaterializedViewUtils.rewriteByRules(parentCascadesContext, context -> {
             Rewriter.getWholeTreeRewriter(context).execute();
             return context.getRewritePlan();
-        }, queryPlanWithUnionFilter, queryPlan), true);
+        }, queryPlanWithUnionFilter, queryPlan, ImmutableSet.of(RuleType.ADD_DEFAULT_LIMIT)), true);
     }
 
     /**

diff --git a/regression-test/data/nereids_rules_p0/mv/with_sql_limit/query_with_sql_limit.out b/regression-test/data/nereids_rules_p0/mv/with_sql_limit/query_with_sql_limit.out
@@ -0,0 +1,19 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !query1_0_before --
+4
+
+-- !query1_0_after --
+4
+
+-- !query2_0_before --
+4
+
+-- !query2_0_after --
+4
+
+-- !query3_0_before --
+4
+
+-- !query3_0_after --
+4
+