Skip to content

Commit

Permalink
Refactor MV Rewrite partition prune
Browse files Browse the repository at this point in the history
Signed-off-by: shuming.li <[email protected]>
  • Loading branch information
LiShuMing committed Oct 25, 2024
1 parent 6080065 commit 9b8da69
Show file tree
Hide file tree
Showing 23 changed files with 222 additions and 190 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand Down Expand Up @@ -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;
Expand All @@ -425,7 +428,9 @@ 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) || !op.isOpRuleBitSet(OP_PARTITION_PRUNED));
if (isNeedFurtherPartitionPrune && context.getQueryMaterializationContext().hasRewrittenSuccess()) {
// Do predicate push down if union rewrite successes.
tree = new SeparateProjectRule().rewrite(tree, rootTaskContext);
deriveLogicalProperty(tree);
Expand Down Expand Up @@ -520,7 +525,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());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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();
});
}
Expand All @@ -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();
});
}
Expand Down
60 changes: 6 additions & 54 deletions fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Utils.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);

Expand All @@ -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();
}
}
Loading

0 comments on commit 9b8da69

Please sign in to comment.