Skip to content

Commit

Permalink
[draft](mtmv) support rewrite when mv has date_trunc but query doesn'…
Browse files Browse the repository at this point in the history
…t have
  • Loading branch information
seawinde committed Dec 3, 2024
1 parent 3e3948c commit 06d35f1
Show file tree
Hide file tree
Showing 7 changed files with 151 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,8 @@ protected Plan rewriteQueryByView(MatchMode matchMode,
queryTopPlan,
materializationContext.getShuttledExprToScanExprMapping(),
viewToQuerySlotMapping,
queryStructInfo.getTableBitSet());
queryStructInfo.getTableBitSet(),
cascadesContext);
boolean isRewrittenQueryExpressionValid = true;
if (!rewrittenQueryExpressions.isEmpty()) {
List<NamedExpression> projects = new ArrayList<>();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,8 @@ protected Plan rewriteQueryByView(MatchMode matchMode,
queryStructInfo.getTopPlan(),
materializationContext.getShuttledExprToScanExprMapping(),
targetToSourceMapping,
queryStructInfo.getTableBitSet()
queryStructInfo.getTableBitSet(),
cascadesContext
);
// Can not rewrite, bail out
if (expressionsRewritten.isEmpty()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@
import org.apache.doris.nereids.rules.exploration.mv.mapping.ExpressionMapping;
import org.apache.doris.nereids.rules.exploration.mv.mapping.RelationMapping;
import org.apache.doris.nereids.rules.exploration.mv.mapping.SlotMapping;
import org.apache.doris.nereids.rules.expression.ExpressionOptimization;
import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext;
import org.apache.doris.nereids.rules.rewrite.MergeProjects;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.Expression;
Expand All @@ -45,6 +47,7 @@
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.scalar.DateTrunc;
import org.apache.doris.nereids.trees.expressions.functions.scalar.ElementAt;
import org.apache.doris.nereids.trees.expressions.functions.scalar.NonNullable;
import org.apache.doris.nereids.trees.expressions.functions.scalar.Nullable;
Expand Down Expand Up @@ -242,7 +245,7 @@ protected List<Plan> doRewrite(StructInfo queryStructInfo, CascadesContext casca
// Try to rewrite compensate predicates by using mv scan
List<Expression> rewriteCompensatePredicates = rewriteExpression(compensatePredicates.toList(),
queryPlan, materializationContext.getShuttledExprToScanExprMapping(),
viewToQuerySlotMapping, queryStructInfo.getTableBitSet());
viewToQuerySlotMapping, queryStructInfo.getTableBitSet(), cascadesContext);
if (rewriteCompensatePredicates.isEmpty()) {
materializationContext.recordFailReason(queryStructInfo,
"Rewrite compensate predicate by view fail",
Expand Down Expand Up @@ -563,7 +566,8 @@ protected Plan rewriteQueryByView(MatchMode matchMode, StructInfo queryStructInf
* then use the corresponding value of mapping to replace it
*/
protected List<Expression> rewriteExpression(List<? extends Expression> sourceExpressionsToWrite, Plan sourcePlan,
ExpressionMapping targetExpressionMapping, SlotMapping targetToSourceMapping, BitSet sourcePlanBitSet) {
ExpressionMapping targetExpressionMapping, SlotMapping targetToSourceMapping, BitSet sourcePlanBitSet,
CascadesContext cascadesContext) {
// Firstly, rewrite the target expression using source with inverse mapping
// then try to use the target expression to represent the query. if any of source expressions
// can not be represented by target expressions, return null.
Expand All @@ -582,18 +586,47 @@ protected List<Expression> rewriteExpression(List<? extends Expression> sourceEx
rewrittenExpressions.add(expressionShuttledToRewrite);
continue;
}
final Set<Object> slotsToRewrite =
final Set<Slot> slotsToRewrite =
expressionShuttledToRewrite.collectToSet(expression -> expression instanceof Slot);

final Set<SlotReference> variants =
expressionShuttledToRewrite.collectToSet(expression -> expression instanceof SlotReference
&& ((SlotReference) expression).getDataType() instanceof VariantType);
&& ((SlotReference) expression).getDataType() instanceof VariantType);
extendMappingByVariant(variants, targetToTargetReplacementMappingQueryBased);
Expression replacedExpression = ExpressionUtils.replace(expressionShuttledToRewrite,
targetToTargetReplacementMappingQueryBased);
if (replacedExpression.anyMatch(slotsToRewrite::contains)) {
// if contains any slot to rewrite, which means can not be rewritten by target, bail out
return ImmutableList.of();
Set<Slot> replacedExpressionSlotQueryUsed = replacedExpression.collect(slotsToRewrite::contains);
if (!replacedExpressionSlotQueryUsed.isEmpty()) {
// if contains any slot to rewrite, which means can not be rewritten by target,
// shuttled query expr is slot#0 > '2024-01-01' but mv plan output is date_trunc(slot#0, 'day')
// which would try to rewrite
// slotDateTruncMap is {date_trunc(slot#0, 'day') : mv_scan_date_trunc_slot#10}
Map<Slot, DateTrunc> slotDateTruncMap = new HashMap<>();
targetToTargetReplacementMappingQueryBased.keySet().forEach(expr -> {
if (expr instanceof DateTrunc && expr.child(0) instanceof Slot) {
slotDateTruncMap.put((Slot) expr.child(0), (DateTrunc) expr);
}
});

if (slotDateTruncMap.isEmpty()) {
//mv date_trunc slot can not offer slot for query,
// can not try to rewrite by date_trunc, bail out
return ImmutableList.of();
}
Expression replacedWithDateTrunc =
ExpressionUtils.replace(expressionShuttledToRewrite, slotDateTruncMap);
// check date_trunc(slot#0, 'day') > '2024-01-01' can simplify
replacedWithDateTrunc = new ExpressionOptimization().rewrite(replacedWithDateTrunc,
new ExpressionRewriteContext(cascadesContext));
if (replacedWithDateTrunc.equals(expressionShuttledToRewrite)) {
// after date_trunc simplify if equals to original expression, could rewritten by mv
replacedExpression = ExpressionUtils.replace(expressionShuttledToRewrite,
targetToTargetReplacementMappingQueryBased,
slotDateTruncMap);
}
if (replacedExpression.anyMatch(slotsToRewrite::contains)) {
return ImmutableList.of();
}
}
rewrittenExpressions.add(replacedExpression);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ protected Plan rewriteQueryByView(MatchMode matchMode,
queryStructInfo.getTopPlan(),
materializationContext.getShuttledExprToScanExprMapping(),
targetToSourceMapping,
queryStructInfo.getTableBitSet()
queryStructInfo.getTableBitSet(),
cascadesContext
);
// Can not rewrite, bail out
if (expressionsRewritten.isEmpty()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import org.apache.doris.nereids.rules.expression.rules.NullSafeEqualToEqual;
import org.apache.doris.nereids.rules.expression.rules.OrToIn;
import org.apache.doris.nereids.rules.expression.rules.SimplifyComparisonPredicate;
import org.apache.doris.nereids.rules.expression.rules.SimplifyDateTrunc;
import org.apache.doris.nereids.rules.expression.rules.SimplifyDecimalV3Comparison;
import org.apache.doris.nereids.rules.expression.rules.SimplifyInPredicate;
import org.apache.doris.nereids.rules.expression.rules.SimplifyRange;
Expand All @@ -42,6 +43,7 @@
public class ExpressionOptimization extends ExpressionRewrite {
public static final List<ExpressionRewriteRule> OPTIMIZE_REWRITE_RULES = ImmutableList.of(
bottomUp(
SimplifyDateTrunc.INSTANCE,
ExtractCommonFactorRule.INSTANCE,
DistinctPredicatesRule.INSTANCE,
SimplifyComparisonPredicate.INSTANCE,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
// 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.rules.expression.rules;

import org.apache.doris.nereids.rules.expression.AbstractExpressionRewriteRule;
import org.apache.doris.nereids.rules.expression.ExpressionMatchingContext;
import org.apache.doris.nereids.rules.expression.ExpressionPatternMatcher;
import org.apache.doris.nereids.rules.expression.ExpressionPatternRuleFactory;
import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext;
import org.apache.doris.nereids.trees.expressions.ComparisonPredicate;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.functions.scalar.DateTrunc;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.util.ExpressionUtils;

import com.google.common.collect.ImmutableList;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
* date_trunc(slot#0, 'day') > '2024-01-01' can be simplify to slot#0 > '2024-01-01'
* */
public class SimplifyDateTrunc extends AbstractExpressionRewriteRule implements ExpressionPatternRuleFactory {

public static SimplifyDateTrunc INSTANCE = new SimplifyDateTrunc();

@Override
public List<ExpressionPatternMatcher<? extends Expression>> buildRules() {
return ImmutableList.of(
matchesType(ComparisonPredicate.class)
.when(predicate ->
(predicate.getArgument(0) instanceof DateTrunc
&& predicate.getArgument(1) instanceof Literal)
|| (predicate.getArgument(1) instanceof DateTrunc
&& predicate.getArgument(0) instanceof Literal)
)
.thenApply(SimplifyDateTrunc::rewriteWhenComparison)
);
}

private static Expression rewriteWhenComparison(ExpressionMatchingContext<? extends Expression> context) {

Expression predicate = context.expr;
DateTrunc dateTrunc = predicate.getArgument(0) instanceof DateTrunc
? (DateTrunc) predicate.getArgument(0) : (DateTrunc) predicate.getArgument(1);
Literal literal = predicate.getArgument(0) instanceof Literal
? (Literal) predicate.getArgument(0) : (Literal) predicate.getArgument(1);
Expression argument = dateTrunc.getArgument(0);
if (!(argument instanceof Slot)) {
return predicate;
}
Map<Expression, Expression> exprMapping = new HashMap<>();
exprMapping.put(argument, literal);
Expression replacedExpr = ExpressionUtils.replace(argument, exprMapping);
Expression evaluatedExpr = FoldConstantRuleOnFE.evaluate(replacedExpr,
new ExpressionRewriteContext(context.cascadesContext));
if (evaluatedExpr.equals(literal)) {
return predicate.withChildren(ImmutableList.of(argument, literal));
}
return predicate;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -403,6 +403,29 @@ public static Expression replace(Expression expr, Map<? extends Expression, ? ex
});
}

/**
* Replace expression node in the expression tree by `replaceMap` in top-down manner.
* For example.
* <pre>
* input expression: a > 1
* replaceMap: a -> b + c
*
* output:
* b + c > 1
* </pre>
*/
public static Expression replace(Expression expr, Map<? extends Expression, ? extends Expression> replaceMap,
Map<? extends Expression, ? extends Expression> transferMap) {
return expr.rewriteDownShortCircuit(e -> {
Expression replacedExpr = replaceMap.get(e);
if (replacedExpr != null) {
return replacedExpr;
}
replacedExpr = replaceMap.get(transferMap.get(e));
return replacedExpr == null ? e : replacedExpr;
});
}

public static List<Expression> replace(List<Expression> exprs,
Map<? extends Expression, ? extends Expression> replaceMap) {
ImmutableList.Builder<Expression> result = ImmutableList.builderWithExpectedSize(exprs.size());
Expand Down

0 comments on commit 06d35f1

Please sign in to comment.