[Refactor] Refactor MV Rewrite partition prune

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/MaterializationContext.java

@@ -223,6 +223,12 @@ public boolean prune(OptimizerContext ctx, OptExpression queryExpression) {
         final List<Table> mvTables = getBaseTables();
         final OperatorType queryOp = queryExpression.getOp().getOpType();
 
+        // if a query has been applied this mv, return false directly.
+        List<LogicalScanOperator> scanOperators = MvUtils.getScanOperator(queryExpression);
+        if (scanOperators.stream().anyMatch(op -> op.isOpAppliedMV(mv.getId()))) {
+            return false;
+        }
+
         if (!checkOperatorCompatible(queryOp)) {
             return false;
         }

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/MvRewritePreprocessor.java

@@ -980,7 +980,7 @@ public static LogicalOlapScanOperator createScanMvOperator(MaterializedView mv,
             }
         }
         final PartitionNames partitionNames = new PartitionNames(false, selectedPartitionNames);
-
+        // MV's selected partition ids/tablet ids are necessary for MV rewrite.
         return LogicalOlapScanOperator.builder()
                 .setTable(mv)
                 .setColRefToColumnMetaMap(colRefToColumnMetaMapBuilder.build())

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Optimizer.java

@@ -128,6 +128,9 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static com.starrocks.sql.optimizer.operator.OpRuleBit.OP_MV_TRANSPARENT_REWRITE;
+import static com.starrocks.sql.optimizer.operator.OpRuleBit.OP_MV_UNION_REWRITE;
+import static com.starrocks.sql.optimizer.operator.OpRuleBit.OP_PARTITION_PRUNED;
 import static com.starrocks.sql.optimizer.rule.RuleType.TF_MATERIALIZED_VIEW;
 
 /**
@@ -406,7 +409,7 @@ private void pruneTables(OptExpression tree, TaskContext rootTaskContext, Column
      */
     private OptExpression transparentMVRewrite(OptExpression tree, TaskContext rootTaskContext) {
         ruleRewriteOnlyOnce(tree, rootTaskContext, new MaterializedViewTransparentRewriteRule());
-        if (Utils.isOptHasAppliedRule(tree, Operator.OP_TRANSPARENT_MV_BIT)) {
+        if (Utils.isOptHasAppliedRule(tree, OP_MV_TRANSPARENT_REWRITE)) {
             tree = new SeparateProjectRule().rewrite(tree, rootTaskContext);
         }
         return tree;
@@ -425,7 +428,12 @@ private void ruleBasedMaterializedViewRewrite(OptExpression tree,
         // NOTE: Since union rewrite will generate Filter -> Union -> OlapScan -> OlapScan, need to push filter below Union
         // and do partition predicate again.
         // TODO: Do it in CBO if needed later.
-        if (MvUtils.isAppliedMVUnionRewrite(tree)) {
+        boolean isNeedFurtherPartitionPrune = Utils.isOptHasAppliedRule(tree, op -> op.isOpRuleBitSet(OP_MV_UNION_REWRITE));
+        if (isNeedFurtherPartitionPrune && context.getQueryMaterializationContext().hasRewrittenSuccess()) {
+            // reset partition prune bit to do partition prune again.
+            MvUtils.getScanOperator(tree).forEach(scan -> {
+                scan.resetOpRuleBit(OP_PARTITION_PRUNED);
+            });
             // Do predicate push down if union rewrite successes.
             tree = new SeparateProjectRule().rewrite(tree, rootTaskContext);
             deriveLogicalProperty(tree);
@@ -520,7 +528,6 @@ private OptExpression logicalRuleRewrite(
         ruleRewriteIterative(tree, rootTaskContext, RuleSetType.PUSH_DOWN_PREDICATE);
         ruleRewriteOnlyOnce(tree, rootTaskContext, SchemaTableEvaluateRule.getInstance());
 
-
         ruleRewriteIterative(tree, rootTaskContext, new MergeTwoProjectRule());
         ruleRewriteOnlyOnce(tree, rootTaskContext, RuleSetType.ELIMINATE_OP_WITH_CONSTANT);
         ruleRewriteOnlyOnce(tree, rootTaskContext, EliminateAggRule.getInstance());

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/OptimizerTraceUtil.java

@@ -88,11 +88,11 @@ public static void logMVRewriteFailReason(String mvName, String format, Object..
     public static void logMVRewrite(MaterializationContext mvContext, String format, Object... object) {
         Tracers.log(Tracers.Module.MV, input -> {
             Object[] args = new Object[] {
-                    mvContext.getMv().getName(),
                     mvContext.getOptimizerContext().isInMemoPhase(),
+                    mvContext.getMv().getName(),
                     MessageFormatter.arrayFormat(format, object).getMessage()
             };
-            return MessageFormatter.arrayFormat("[MV TRACE] [REWRITE {}] [InMemo:{}] {}",
+            return MessageFormatter.arrayFormat("[MV TRACE] [REWRITE] [InMemo:{}] [{}] {}",
                     args).getMessage();
         });
     }
@@ -102,11 +102,11 @@ public static void logMVRewrite(MvRewriteContext mvRewriteContext, String format
         Tracers.log(Tracers.Module.MV, input -> {
             Object[] args = new Object[] {
                     mvRewriteContext.getRule().type().name(),
-                    mvContext.getMv().getName(),
                     mvRewriteContext.getMaterializationContext().getOptimizerContext().isInMemoPhase(),
+                    mvContext.getMv().getName(),
                     MessageFormatter.arrayFormat(format, object).getMessage()
             };
-            return MessageFormatter.arrayFormat("[MV TRACE] [REWRITE {} {}] [InMemo:{}] {}",
+            return MessageFormatter.arrayFormat("[MV TRACE] [REWRITE {}] [InMemo:{}] [{}] {}",
                     args).getMessage();
         });
     }

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Utils.java

@@ -55,7 +55,6 @@
 import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 import com.starrocks.sql.optimizer.rewrite.ReplaceColumnRefRewriter;
 import com.starrocks.sql.optimizer.rewrite.ScalarOperatorRewriter;
-import com.starrocks.sql.optimizer.rule.transformation.materialization.MvUtils;
 import com.starrocks.sql.optimizer.statistics.ColumnStatistic;
 import com.starrocks.sql.optimizer.statistics.StatisticsCalculator;
 import org.apache.commons.collections4.CollectionUtils;
@@ -847,78 +846,31 @@ public static OptExpression mergeProjection(OptExpression input,
         return input;
     }
 
-    /**
-     * Check if the operator has applied the rule
-     * @param op input operator to be checked
-     * @param ruleMask specific rule mask
-     * @return true if the operator has applied the rule, false otherwise
-     */
-    public static boolean isOpAppliedRule(Operator op, int ruleMask) {
-        if (op == null) {
-            return false;
-        }
-        // TODO: support cte inline
-        int opRuleMask = op.getOpRuleMask();
-        return (opRuleMask & ruleMask) != 0;
-    }
-
-    /**
-     * Set the rule mask to the operator
-     * @param op input operator
-     * @param ruleMask specific rule mask
-     */
-    public static void setOpAppliedRule(Operator op, int ruleMask) {
-        if (op == null) {
-            return;
-        }
-        op.setOpRuleMask(op.getOpRuleMask() | ruleMask);
-    }
-
-    /**
-     * Reset the rule mask to the operator
-     * @param op input operator
-     * @param ruleMask specific rule mask
-     */
-    public static void resetOpAppliedRule(Operator op, int ruleMask) {
-        if (op == null) {
-            return;
-        }
-        op.resetOpRuleMask(ruleMask);
-    }
-
     /**
      * Check if the optExpression has applied the rule in recursively
      * @param optExpression input optExpression to be checked
      * @param ruleMask specific rule mask
      * @return true if the optExpression or its children have applied the rule, false otherwise
      */
     public static boolean isOptHasAppliedRule(OptExpression optExpression, int ruleMask) {
+        return isOptHasAppliedRule(optExpression, op -> op.isOpRuleBitSet(ruleMask));
+    }
+
+    public static boolean isOptHasAppliedRule(OptExpression optExpression, Predicate<Operator> pred) {
         if (optExpression == null) {
             return false;
         }
-        if (isOpAppliedRule(optExpression.getOp(), ruleMask)) {
+        if (pred.test(optExpression.getOp())) {
             return true;
         }
         for (OptExpression child : optExpression.getInputs()) {
-            if (isOptHasAppliedRule(child, ruleMask)) {
+            if (isOptHasAppliedRule(child, pred)) {
                 return true;
             }
         }
         return false;
     }
 
-
-    public static void setOptScanOpsBit(OptExpression input,
-                                        int bit) {
-        List<LogicalScanOperator> scanOps = MvUtils.getScanOperator(input);
-        scanOps.stream().forEach(op -> op.setOpRuleMask(op.getOpRuleMask() | bit));
-    }
-
-    public static void setOpBit(OptExpression input,
-                                int bit) {
-        input.getOp().setOpRuleMask(input.getOp().getOpRuleMask() | bit);
-    }
-
     @SuppressWarnings("unchecked")
     public static <T, S extends T> Optional<S> downcast(T obj, Class<S> klass) {
         Preconditions.checkArgument(obj != null);

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/base/LogicalProperty.java

@@ -176,7 +176,8 @@ public OneTabletProperty visitLogicalViewScan(LogicalViewScanOperator node, Expr
         @Override
         public OneTabletProperty visitLogicalTableScan(LogicalScanOperator node, ExpressionContext context) {
             if (node instanceof LogicalOlapScanOperator) {
-                if (((LogicalOlapScanOperator) node).getSelectedTabletId().size() <= 1) {
+                LogicalOlapScanOperator olapScanOperator = (LogicalOlapScanOperator) node;
+                if (olapScanOperator.getSelectedTabletId() != null && olapScanOperator.getSelectedTabletId().size() <= 1) {
                     Set<String> distributionColumnNames = node.getTable().getDistributionColumnNames();
                     List<ColumnRefOperator> bucketColumns = Lists.newArrayList();
                     for (Map.Entry<ColumnRefOperator, Column> entry : node.getColRefToColumnMetaMap().entrySet()) {

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/operator/OpRuleBit.java

@@ -0,0 +1,33 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed 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
+//
+//     https://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 com.starrocks.sql.optimizer.operator;
+
+public class OpRuleBit {
+    // Like LogicalJoinOperator#transformMask, add a mask to avoid one operator's dead-loop in one transform rule.
+    // eg: MV's UNION-ALL RULE:
+    //                 UNION                         UNION
+    //               /        \                    /       \
+    //  OP -->   EXTRA-OP    MV-SCAN  -->     UNION    MV-SCAN     ---> ....
+    //                                       /      \
+    //                                  EXTRA-OP    MV-SCAN
+    // Operator has been union rewrite or not, if union all, no need to union again.
+    public static final int OP_MV_UNION_REWRITE = 0;
+    // Operator has been push down predicates or not, if push down predicates, no need to push down again.
+    public static final int OP_MV_AGG_PUSH_DOWN_REWRITE = 1;
+    // Operator has been transparent mv rewrite or not, if transparent mv rewrite, no need to rewrite again.
+    public static final int OP_MV_TRANSPARENT_REWRITE = 2;
+    // Operator has been partition pruned or not, if partition pruned, no need to prune again.
+    public static final int OP_PARTITION_PRUNED = 3;
+}

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/operator/Operator.java

@@ -14,6 +14,7 @@
 
 package com.starrocks.sql.optimizer.operator;
 
+import com.google.api.client.util.Sets;
 import com.google.common.collect.Lists;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.OptExpressionVisitor;
@@ -25,9 +26,11 @@
 import com.starrocks.sql.optimizer.operator.scalar.ScalarOperator;
 import com.starrocks.sql.optimizer.property.DomainProperty;
 
+import java.util.BitSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 
 public abstract class Operator {
     public static final long DEFAULT_LIMIT = -1;
@@ -55,20 +58,10 @@ public abstract class Operator {
     // or self reference of groups
     protected long salt = 0;
 
-    protected int opRuleMask = 0;
-    // Like LogicalJoinOperator#transformMask, add a mask to avoid one operator's dead-loop in one transform rule.
-    // eg: MV's UNION-ALL RULE:
-    //                 UNION                         UNION
-    //               /        \                    /       \
-    //  OP -->   EXTRA-OP    MV-SCAN  -->     UNION    MV-SCAN     ---> ....
-    //                                       /      \
-    //                                  EXTRA-OP    MV-SCAN
-    // Operator's rule mask: operator that has been union rewrite and no needs to rewrite again.
-    public static final int OP_UNION_ALL_BIT = 1 << 0;
-    // Operator's rule mask: operator that has been push down rewrite and no needs to rewrite again.
-    public static final int OP_PUSH_DOWN_BIT = 1 << 1;
-    public static final int OP_TRANSPARENT_MV_BIT = 1 << 2;
-    public static final int OP_PARTITION_PRUNE_BIT = 1 << 3;
+    // mark which rule(bit) has been applied to the operator.
+    protected BitSet opRuleBits = new BitSet();
+    // mark which mv has been applied to the operator
+    protected Set<Long> opAppliedMVs = Sets.newHashSet();
 
     // an operator logically equivalent to 'this' operator
     // used by view based mv rewrite
@@ -159,20 +152,24 @@ public long getSalt() {
         return salt;
     }
 
-    public int getOpRuleMask() {
-        return opRuleMask;
+    public void setOpRuleBit(int bit) {
+        this.opRuleBits.set(bit);
     }
 
-    public void setOpRuleMask(int bit) {
-        this.opRuleMask |= bit;
+    public void resetOpRuleBit(int bit) {
+        this.opRuleBits.clear(bit);
     }
 
-    public void resetOpRuleMask(int bit) {
-        this.opRuleMask &= (~ bit);
+    public boolean isOpRuleBitSet(int bit) {
+        return opRuleBits.get(bit);
     }
 
-    public boolean isOpRuleMaskSet(int bit) {
-        return (opRuleMask & bit) != 0;
+    public void setOpAppliedMV(long mvId) {
+        this.opAppliedMVs.add(mvId);
+    }
+
+    public boolean isOpAppliedMV(long mvId) {
+        return opAppliedMVs.contains(mvId);
     }
 
     public Operator getEquivalentOp() {
@@ -279,8 +276,9 @@ public B withOperator(O operator) {
             builder.predicate = operator.predicate;
             builder.projection = operator.projection;
             builder.salt = operator.salt;
-            builder.opRuleMask = operator.opRuleMask;
             builder.equivalentOp = operator.equivalentOp;
+            builder.opRuleBits.or(operator.opRuleBits);
+            builder.opAppliedMVs.addAll(operator.opAppliedMVs);
             return (B) this;
         }
 
@@ -326,8 +324,8 @@ public B addSalt() {
             return (B) this;
         }
 
-        public B setOpBitSet(int opRuleMask) {
-            builder.opRuleMask = opRuleMask;
+        public B setOpBitSet(BitSet opRuleMask) {
+            builder.opRuleBits = opRuleMask;
             return (B) this;
         }
     }

fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/transformation/ExternalScanPartitionPruneRule.java

@@ -18,8 +18,6 @@
 import com.google.common.collect.Lists;
 import com.starrocks.sql.optimizer.OptExpression;
 import com.starrocks.sql.optimizer.OptimizerContext;
-import com.starrocks.sql.optimizer.Utils;
-import com.starrocks.sql.optimizer.operator.Operator;
 import com.starrocks.sql.optimizer.operator.OperatorType;
 import com.starrocks.sql.optimizer.operator.logical.LogicalScanOperator;
 import com.starrocks.sql.optimizer.operator.pattern.Pattern;
@@ -28,8 +26,11 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
+import java.util.Collections;
 import java.util.List;
 
+import static com.starrocks.sql.optimizer.operator.OpRuleBit.OP_PARTITION_PRUNED;
+
 public class ExternalScanPartitionPruneRule extends TransformationRule {
     private static final Logger LOG = LogManager.getLogger(ExternalScanPartitionPruneRule.class);
 
@@ -56,21 +57,14 @@ public ExternalScanPartitionPruneRule(OperatorType logicalOperatorType) {
         super(RuleType.TF_PARTITION_PRUNE, Pattern.create(logicalOperatorType));
     }
 
-    @Override
-    public boolean check(final OptExpression input, OptimizerContext context) {
-        Operator op = input.getOp();
-        // if the partition id is already selected, no need to prune again
-        if (Utils.isOpAppliedRule(op, Operator.OP_PARTITION_PRUNE_BIT)) {
-            return false;
-        }
-        return true;
-    }
-
     @Override
     public List<OptExpression> transform(OptExpression input, OptimizerContext context) {
         LogicalScanOperator operator = (LogicalScanOperator) input.getOp();
+        if (operator.isOpRuleBitSet(OP_PARTITION_PRUNED)) {
+            return Collections.emptyList();
+        }
         OptExternalPartitionPruner.prunePartitions(context, operator);
-        Utils.setOpAppliedRule(operator, Operator.OP_PARTITION_PRUNE_BIT);
+        operator.setOpRuleBit(OP_PARTITION_PRUNED);
         return Lists.newArrayList();
     }
 }