AbstractMaterializedViewAggregateRule.java
// 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.exploration.mv;
import org.apache.doris.catalog.Column;
import org.apache.doris.common.Pair;
import org.apache.doris.mtmv.BaseColInfo;
import org.apache.doris.mtmv.BaseTableInfo;
import org.apache.doris.mtmv.MTMVPartitionInfo;
import org.apache.doris.nereids.CascadesContext;
import org.apache.doris.nereids.properties.DataTrait;
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;
import org.apache.doris.nereids.rules.exploration.mv.StructInfo.PlanSplitContext;
import org.apache.doris.nereids.rules.exploration.mv.mapping.SlotMapping;
import org.apache.doris.nereids.rules.exploration.mv.rollup.AggFunctionRollUpHandler;
import org.apache.doris.nereids.rules.exploration.mv.rollup.BothCombinatorRollupHandler;
import org.apache.doris.nereids.rules.exploration.mv.rollup.ContainDistinctFunctionRollupHandler;
import org.apache.doris.nereids.rules.exploration.mv.rollup.DirectRollupHandler;
import org.apache.doris.nereids.rules.exploration.mv.rollup.MappingRollupHandler;
import org.apache.doris.nereids.rules.exploration.mv.rollup.SingleCombinatorRollupHandler;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.ExprId;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.functions.Function;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate;
import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.trees.plans.logical.LogicalRepeat;
import org.apache.doris.nereids.trees.plans.visitor.ExpressionLineageReplacer;
import org.apache.doris.nereids.util.ExpressionUtils;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Sets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Supplier;
import java.util.stream.Collectors;
/**
* AbstractMaterializedViewAggregateRule
* This is responsible for common aggregate rewriting
*/
public abstract class AbstractMaterializedViewAggregateRule extends AbstractMaterializedViewRule {
public static final List<AggFunctionRollUpHandler> ROLL_UP_HANDLERS =
ImmutableList.of(DirectRollupHandler.INSTANCE,
MappingRollupHandler.INSTANCE,
SingleCombinatorRollupHandler.INSTANCE,
BothCombinatorRollupHandler.INSTANCE,
ContainDistinctFunctionRollupHandler.INSTANCE);
protected static final AggregateExpressionRewriter AGGREGATE_EXPRESSION_REWRITER =
new AggregateExpressionRewriter();
@Override
protected Plan rewriteQueryByView(MatchMode matchMode,
StructInfo queryStructInfo,
StructInfo viewStructInfo,
SlotMapping viewToQuerySlotMapping,
Plan tempRewritedPlan,
MaterializationContext materializationContext,
CascadesContext cascadesContext) {
// get view and query aggregate and top plan correspondingly
Pair<Plan, LogicalAggregate<Plan>> viewTopPlanAndAggPair = splitToTopPlanAndAggregate(viewStructInfo);
if (viewTopPlanAndAggPair == null) {
materializationContext.recordFailReason(queryStructInfo,
"Split view to top plan and agg fail, view doesn't not contain aggregate",
() -> String.format("view plan = %s\n", viewStructInfo.getOriginalPlan().treeString()));
return null;
}
Pair<Plan, LogicalAggregate<Plan>> queryTopPlanAndAggPair = splitToTopPlanAndAggregate(queryStructInfo);
if (queryTopPlanAndAggPair == null) {
materializationContext.recordFailReason(queryStructInfo,
"Split query to top plan and agg fail",
() -> String.format("query plan = %s\n", queryStructInfo.getOriginalPlan().treeString()));
return null;
}
Plan queryTopPlan = queryTopPlanAndAggPair.key();
LogicalAggregate<Plan> queryAggregate = queryTopPlanAndAggPair.value();
if (!checkCompatibility(queryStructInfo, queryAggregate, viewTopPlanAndAggPair.value(),
materializationContext)) {
return null;
}
boolean queryContainsGroupSets = queryAggregate.getSourceRepeat().isPresent();
// If group by expression between query and view is equals, try to rewrite expression directly
if (!queryContainsGroupSets && isGroupByEquals(queryTopPlanAndAggPair, viewTopPlanAndAggPair,
viewToQuerySlotMapping, queryStructInfo, viewStructInfo, tempRewritedPlan, materializationContext,
cascadesContext)) {
List<Expression> rewrittenQueryExpressions = rewriteExpression(queryTopPlan.getOutput(),
queryTopPlan,
materializationContext.getShuttledExprToScanExprMapping(),
viewToQuerySlotMapping,
ImmutableMap.of(), cascadesContext);
boolean isRewrittenQueryExpressionValid = true;
if (!rewrittenQueryExpressions.isEmpty()) {
List<NamedExpression> projects = new ArrayList<>();
for (Expression expression : rewrittenQueryExpressions) {
if (expression.containsType(AggregateFunction.class)) {
// record the reason and then try to roll up aggregate function
materializationContext.recordFailReason(queryStructInfo,
"rewritten expression contains aggregate functions when group equals aggregate rewrite",
() -> String.format("aggregate functions = %s\n", rewrittenQueryExpressions));
isRewrittenQueryExpressionValid = false;
}
projects.add(expression instanceof NamedExpression
? (NamedExpression) expression : new Alias(expression));
}
if (isRewrittenQueryExpressionValid) {
return new LogicalProject<>(projects, tempRewritedPlan);
}
}
// if fails, record the reason and then try to roll up aggregate function
materializationContext.recordFailReason(queryStructInfo,
"Can not rewrite expression when no roll up",
() -> String.format("expressionToWrite = %s,\n mvExprToMvScanExprMapping = %s,\n"
+ "viewToQuerySlotMapping = %s",
queryTopPlan.getOutput(),
materializationContext.getShuttledExprToScanExprMapping(),
viewToQuerySlotMapping));
}
return aggregateRewriteByView(queryStructInfo,
viewToQuerySlotMapping,
queryTopPlanAndAggPair,
tempRewritedPlan,
materializationContext,
ExpressionRewriteMode.EXPRESSION_DIRECT,
ExpressionRewriteMode.EXPRESSION_ROLL_UP);
}
/**
* Aggregate function and group by expression rewrite impl
*/
protected Plan aggregateRewriteByView(
StructInfo queryStructInfo,
SlotMapping viewToQuerySlotMapping,
Pair<Plan, LogicalAggregate<Plan>> queryTopPlanAndAggPair,
Plan tempRewritedPlan,
MaterializationContext materializationContext,
ExpressionRewriteMode groupByMode,
ExpressionRewriteMode aggregateFunctionMode) {
// try to roll up.
// try to rewrite, contains both roll up aggregate functions and aggregate group expression
List<NamedExpression> finalOutputExpressions = new ArrayList<>();
List<Expression> finalGroupExpressions = new ArrayList<>();
// permute the mv expr mapping to query based
Map<Expression, Expression> mvExprToMvScanExprQueryBased =
materializationContext.getShuttledExprToScanExprMapping().keyPermute(viewToQuerySlotMapping)
.flattenMap().get(0);
Plan queryTopPlan = queryStructInfo.getTopPlan();
LogicalAggregate<Plan> queryAggregate = queryTopPlanAndAggPair.value();
List<Expression> queryGroupByExpressions = queryAggregate.getGroupByExpressions();
if (queryAggregate.getSourceRepeat().isPresent()) {
// The group by/function classification of the query top plan output expressions is only
// used by the repeat rewrite, so it is computed lazily inside this branch to avoid paying
// the full plan lineage walk for every ordinary aggregate rewrite.
// split the query top plan expressions to group expressions and functions, if can not, bail out.
Pair<Set<? extends Expression>, Set<? extends Expression>> queryGroupAndFunctionPair
= topPlanSplitToGroupAndFunction(queryTopPlanAndAggPair, queryStructInfo);
Set<? extends Expression> queryTopPlanGroupBySet = queryGroupAndFunctionPair.key();
Set<? extends Expression> queryTopPlanFunctionSet = queryGroupAndFunctionPair.value();
// try to rewrite the query top plan expressions, the query top plan output expressions
// are used as the repeat output expressions directly
for (Expression topExpression : queryTopPlan.getOutput()) {
if (queryTopPlanFunctionSet.contains(topExpression)) {
// if agg function, try to roll up and rewrite
Expression rollupedExpression = tryRewriteExpression(queryStructInfo, topExpression,
mvExprToMvScanExprQueryBased, aggregateFunctionMode, materializationContext,
"Query function roll up fail",
() -> String.format("queryExpression = %s,\n mvExprToMvScanExprQueryBased = %s",
topExpression, mvExprToMvScanExprQueryBased));
if (rollupedExpression == null) {
return null;
}
finalOutputExpressions.add(new Alias(rollupedExpression));
} else {
// if group by dimension, try to rewrite
Expression rewrittenGroupByExpression = tryRewriteExpression(queryStructInfo, topExpression,
mvExprToMvScanExprQueryBased, groupByMode, materializationContext,
"View dimensions doesn't not cover the query dimensions",
() -> String.format("mvExprToMvScanExprQueryBased is %s,\n queryExpression is %s",
mvExprToMvScanExprQueryBased, topExpression));
if (rewrittenGroupByExpression == null) {
// group expr can not rewrite by view
return null;
}
NamedExpression groupByExpression = rewrittenGroupByExpression instanceof NamedExpression
? (NamedExpression) rewrittenGroupByExpression : new Alias(rewrittenGroupByExpression);
finalOutputExpressions.add(groupByExpression);
finalGroupExpressions.add(groupByExpression);
}
}
// handle the scene that query top plan not use the group by in query bottom aggregate
if (needCompensateGroupBy(queryTopPlanGroupBySet, queryGroupByExpressions)) {
for (Expression expression : queryGroupByExpressions) {
if (queryTopPlanGroupBySet.contains(expression)) {
continue;
}
Expression rewrittenGroupByExpression = tryRewriteExpression(queryStructInfo, expression,
mvExprToMvScanExprQueryBased, groupByMode, materializationContext,
"View dimensions doesn't not cover the query dimensions in bottom agg ",
() -> String.format("mvExprToMvScanExprQueryBased is %s,\n expression is %s",
mvExprToMvScanExprQueryBased, expression));
if (rewrittenGroupByExpression == null) {
return null;
}
NamedExpression groupByExpression = rewrittenGroupByExpression instanceof NamedExpression
? (NamedExpression) rewrittenGroupByExpression : new Alias(rewrittenGroupByExpression);
finalGroupExpressions.add(groupByExpression);
}
}
// construct group sets for repeat
List<List<Expression>> rewrittenGroupSetsExpressions = new ArrayList<>();
List<List<Expression>> groupingSets = queryAggregate.collectFirst(LogicalRepeat.class::isInstance)
.map(repeat -> ((LogicalRepeat<? extends Plan>) repeat).getGroupingSets())
.orElse(queryAggregate.getSourceRepeat().get().getGroupingSets());
for (List<Expression> groupingSet : groupingSets) {
if (groupingSet.isEmpty()) {
rewrittenGroupSetsExpressions.add(ImmutableList.of());
} else {
List<Expression> rewrittenGroupSetExpressions = new ArrayList<>();
for (Expression expression : groupingSet) {
Expression rewrittenGroupByExpression = tryRewriteExpression(queryStructInfo, expression,
mvExprToMvScanExprQueryBased, ExpressionRewriteMode.EXPRESSION_DIRECT,
materializationContext,
"View dimensions doesn't not cover the query group set dimensions",
() -> String.format("mvExprToMvScanExprQueryBased is %s,\n queryExpression is %s",
mvExprToMvScanExprQueryBased, expression));
if (rewrittenGroupByExpression == null) {
return null;
}
rewrittenGroupSetExpressions.add(rewrittenGroupByExpression);
}
rewrittenGroupSetsExpressions.add(rewrittenGroupSetExpressions);
}
}
LogicalRepeat<Plan> repeat = new LogicalRepeat<>(rewrittenGroupSetsExpressions,
finalOutputExpressions, queryStructInfo.getGroupingId().get(),
queryAggregate.getSourceRepeat().get().getGroupingIdValues().orElse(null),
queryAggregate.getSourceRepeat().get().getRepeatType(), tempRewritedPlan);
return NormalizeRepeat.doNormalize(repeat);
}
// The rewritten aggregate should group by the query bottom aggregate's group by expressions,
// and its output expressions should be the rewritten group by expressions and the rolled up
// aggregate functions. The query top plan output expressions are recomputed by a project above
// the rewritten aggregate, so the projection of a group by key in the query top plan (such as
// `select cast(date_trunc(ts, 'day') as string) from t group by date_trunc(ts, 'day')`) will not
// be wrongly treated as a group by key of the rewritten aggregate.
// The mapping from the query bottom aggregate output slot to the new aggregate output expression
// is used to rewrite the query top plan output expressions to reference the new aggregate output.
// The shuttled query bottom aggregate outputs are used as the map keys and are passed to the
// rewriter directly, and the shuttled query top plan outputs restore the projection expressions,
// so the whole top plan is traversed only twice instead of once per output expression.
List<? extends Expression> shuttledBottomAggOutputs = ExpressionUtils.shuttleExpressionWithLineage(
queryAggregate.getOutputExpressions(), queryTopPlan);
List<? extends Expression> shuttledTopPlanOutputs = ExpressionUtils.shuttleExpressionWithLineage(
queryTopPlan.getOutput(), queryTopPlan);
Map<Expression, Expression> bottomAggOutputToNewExprMap = new HashMap<>();
Set<Expression> queryGroupByExpressionSet = new HashSet<>(queryGroupByExpressions);
List<NamedExpression> queryAggregateOutputs = queryAggregate.getOutputExpressions();
for (int i = 0; i < queryAggregateOutputs.size(); i++) {
NamedExpression queryAggregateOutput = queryAggregateOutputs.get(i);
Expression shuttledQueryAggregateOutput = shuttledBottomAggOutputs.get(i);
if (queryGroupByExpressionSet.contains(queryAggregateOutput)) {
// if it is a group by expression, rewrite it to the new aggregate group by key
Expression rewrittenGroupByExpression = rewriteShuttledExpression(queryStructInfo,
shuttledQueryAggregateOutput, mvExprToMvScanExprQueryBased, groupByMode,
materializationContext,
"View dimensions doesn't not cover the query dimensions",
() -> String.format("mvExprToMvScanExprQueryBased is %s,\n queryExpression is %s",
mvExprToMvScanExprQueryBased, queryAggregateOutput));
if (rewrittenGroupByExpression == null) {
return null;
}
NamedExpression groupByOutput = rewrittenGroupByExpression instanceof NamedExpression
? (NamedExpression) rewrittenGroupByExpression : new Alias(rewrittenGroupByExpression);
finalGroupExpressions.add(groupByOutput);
finalOutputExpressions.add(groupByOutput);
bottomAggOutputToNewExprMap.put(shuttledQueryAggregateOutput, groupByOutput.toSlot());
} else {
// if it is an aggregate function, try to roll up and rewrite
Expression rewrittenFunction = rewriteShuttledExpression(queryStructInfo,
shuttledQueryAggregateOutput, mvExprToMvScanExprQueryBased, aggregateFunctionMode,
materializationContext,
"Query function roll up fail",
() -> String.format("queryExpression = %s,\n mvExprToMvScanExprQueryBased = %s",
queryAggregateOutput, mvExprToMvScanExprQueryBased));
if (rewrittenFunction == null) {
return null;
}
NamedExpression functionOutput = new Alias(rewrittenFunction);
finalOutputExpressions.add(functionOutput);
bottomAggOutputToNewExprMap.put(shuttledQueryAggregateOutput, functionOutput.toSlot());
}
}
LogicalAggregate<Plan> rewrittenAggregate =
new LogicalAggregate<>(finalGroupExpressions, finalOutputExpressions, tempRewritedPlan);
// rewrite the query top plan output expressions to reference the rewritten aggregate output,
// the query top plan output slot is shuttled by lineage firstly to restore the projection
// expression, so a projection of the group by key in the query top plan can be recomputed
// by a project above the rewritten aggregate.
List<NamedExpression> topProjectExpressions = new ArrayList<>();
// Preserve the original top output expr id multiplicity: several query top plan output positions
// can share one original expr id (e.g. the unaliased `select k, k from t group by k` references the
// same output slot twice). Such repeated positions must also share one rewritten output expr id,
// otherwise the rewritten output set is inflated and MaterializedViewUtils.rewriteByRules returns at
// its output-set-size guard before the whole-tree normalization and partition pruning, which can
// miscompute invalid-partition compensation or leave the new aggregate unnormalized. Conversely,
// positions with distinct original expr ids that are rewritten to the same aggregate output (e.g.
// `select sum(v) as s1, sum(v) as s2 from t group by a`) still need distinct output expr ids, so the
// rewritten output set is not collapsed.
List<Slot> queryTopPlanOutputs = queryTopPlan.getOutput();
Map<ExprId, NamedExpression> originalExprIdToRewritten = new HashMap<>();
Set<ExprId> usedRewrittenExprIds = new HashSet<>();
for (int i = 0; i < shuttledTopPlanOutputs.size(); i++) {
ExprId originalExprId = queryTopPlanOutputs.get(i).getExprId();
NamedExpression groupRewrittenExpression = originalExprIdToRewritten.get(originalExprId);
if (groupRewrittenExpression == null) {
Expression replacedExpression = ExpressionUtils.replace(shuttledTopPlanOutputs.get(i),
bottomAggOutputToNewExprMap);
groupRewrittenExpression = replacedExpression instanceof NamedExpression
? (NamedExpression) replacedExpression : new Alias(replacedExpression);
if (!usedRewrittenExprIds.add(groupRewrittenExpression.getExprId())) {
// The rewritten expr id is already used by another original expr id group, keep a
// distinct output expr id for this group.
groupRewrittenExpression = new Alias(replacedExpression);
}
originalExprIdToRewritten.put(originalExprId, groupRewrittenExpression);
}
topProjectExpressions.add(groupRewrittenExpression);
}
// If the query top plan output expressions can be produced by the rewritten aggregate directly,
// return the aggregate, otherwise compute them by a project above the rewritten aggregate.
// Note the query top plan output may be a strict matching prefix of the rewritten aggregate output
// (e.g. `select k1 from t group by k1, k2`), in which case the redundant aggregate outputs must be
// projected away. Otherwise the rewritten plan output count differs from the query and the candidate
// is rejected by MaterializedViewUtils.normalizeExpressions, so a valid sync MV is silently not used.
boolean needTopProject = topProjectExpressions.size() != finalOutputExpressions.size();
for (int i = 0; i < topProjectExpressions.size(); i++) {
if (i >= finalOutputExpressions.size()
|| !topProjectExpressions.get(i).toSlot().equals(finalOutputExpressions.get(i).toSlot())) {
needTopProject = true;
break;
}
}
if (!needTopProject) {
return rewrittenAggregate;
}
return new LogicalProject<>(topProjectExpressions, rewrittenAggregate);
}
/**
* handle the scene that query top plan not use the group by in query bottom aggregate
* If mv is select o_orderdate from orders group by o_orderdate;
* query is select 1 from orders group by o_orderdate.
* Or mv is select o_orderdate from orders group by o_orderdate
* query is select o_orderdate from orders group by o_orderdate, o_orderkey;
* if the slot which query top project use can not cover the slot which query bottom aggregate group by slot
* should compensate group by to make sure the data is right.
* For example:
* mv is select o_orderdate from orders group by o_orderdate;
* query is select o_orderdate from orders group by o_orderdate, o_orderkey;
*
* @param queryGroupByExpressions query bottom aggregate group by is o_orderdate, o_orderkey
* @param queryTopProject query top project is o_orderdate
* @return need to compensate group by if true or not need
*
*/
private static boolean needCompensateGroupBy(Set<? extends Expression> queryTopProject,
List<Expression> queryGroupByExpressions) {
Set<Expression> queryGroupByExpressionSet = new HashSet<>(queryGroupByExpressions);
if (queryGroupByExpressionSet.size() != queryTopProject.size()) {
return true;
}
Set<NamedExpression> queryTopPlanGroupByUseNamedExpressions = new HashSet<>();
Set<NamedExpression> queryGroupByUseNamedExpressions = new HashSet<>();
for (Expression expr : queryTopProject) {
queryTopPlanGroupByUseNamedExpressions.addAll(expr.collect(NamedExpression.class::isInstance));
}
for (Expression expr : queryGroupByExpressionSet) {
queryGroupByUseNamedExpressions.addAll(expr.collect(NamedExpression.class::isInstance));
}
// if the slots query top project use can not cover the slots which query bottom aggregate use
// Should compensate.
return !queryTopPlanGroupByUseNamedExpressions.containsAll(queryGroupByUseNamedExpressions);
}
/**
* Try to rewrite query expression by view, contains both group by dimension and aggregate function
*/
protected Expression tryRewriteExpression(StructInfo queryStructInfo, Expression queryExpression,
Map<Expression, Expression> mvShuttledExprToMvScanExprQueryBased, ExpressionRewriteMode rewriteMode,
MaterializationContext materializationContext, String summaryIfFail, Supplier<String> detailIfFail) {
Expression queryFunctionShuttled = ExpressionUtils.shuttleExpressionWithLineage(
queryExpression,
queryStructInfo.getTopPlan());
return rewriteShuttledExpression(queryStructInfo, queryFunctionShuttled, mvShuttledExprToMvScanExprQueryBased,
rewriteMode, materializationContext, summaryIfFail, detailIfFail);
}
/**
* Rewrite the shuttled query expression by view, contains both group by dimension and aggregate
* function. The query expression is expected to be shuttled by lineage already, so the top plan is
* not traversed again here and the batched lineage result can be reused.
*/
private Expression rewriteShuttledExpression(StructInfo queryStructInfo, Expression queryShuttledExpression,
Map<Expression, Expression> mvShuttledExprToMvScanExprQueryBased, ExpressionRewriteMode rewriteMode,
MaterializationContext materializationContext, String summaryIfFail, Supplier<String> detailIfFail) {
AggregateExpressionRewriteContext expressionRewriteContext = new AggregateExpressionRewriteContext(
rewriteMode, mvShuttledExprToMvScanExprQueryBased, queryStructInfo.getTopPlan(),
queryStructInfo.getGroupingId());
Expression rewrittenExpression = queryShuttledExpression.accept(AGGREGATE_EXPRESSION_REWRITER,
expressionRewriteContext);
if (!expressionRewriteContext.isValid()) {
materializationContext.recordFailReason(queryStructInfo, summaryIfFail, detailIfFail);
return null;
}
return rewrittenExpression;
}
/**
* Not all query after rewritten successfully can compensate union all
* Such as:
* mv def sql is as following, partition column is a
* select a, b, count(*) from t1 group by a, b
* Query is as following:
* select b, count(*) from t1 group by b, after rewritten by materialized view successfully
* If mv part partition is invalid, can not compensate union all, because result is wrong after
* compensate union all.
*/
@Override
protected boolean canUnionRewrite(Plan queryPlan, AsyncMaterializationContext context,
CascadesContext cascadesContext) {
// Check query plan is contain the partition column
// Query plan in the current rule must contain aggregate node, because the rule pattern is
Optional<LogicalAggregate<Plan>> logicalAggregateOptional =
queryPlan.collectFirst(planTreeNode -> planTreeNode instanceof LogicalAggregate);
if (!logicalAggregateOptional.isPresent()) {
return true;
}
List<Expression> groupByExpressions = logicalAggregateOptional.get().getGroupByExpressions();
if (groupByExpressions.isEmpty()) {
// Scalar aggregate can not compensate union all
return false;
}
MTMVPartitionInfo mvPartitionInfo = context.getMtmv().getMvPartitionInfo();
List<BaseColInfo> pctInfos = mvPartitionInfo.getPctInfos();
boolean canUnionRewrite = false;
// Check the query plan group by expression contains partition col or not
List<? extends Expression> groupByShuttledExpressions =
ExpressionUtils.shuttleExpressionWithLineage(groupByExpressions, queryPlan);
for (Expression expression : groupByShuttledExpressions) {
canUnionRewrite = !expression.collectToSet(expr -> {
if (!(expr instanceof SlotReference) || !((SlotReference) expr).isColumnFromTable()) {
return false;
}
String columnName = ((SlotReference) expr).getOriginalColumn().map(Column::getName).orElse(null);
BaseTableInfo baseTableInfo = ((SlotReference) expr).getOriginalTable().map(BaseTableInfo::new)
.orElse(null);
return pctInfos.stream().anyMatch(colInfo -> colInfo.getTableInfo().equals(baseTableInfo)
&& colInfo.getColName().equals(columnName));
}
).isEmpty();
if (canUnionRewrite) {
break;
}
}
return canUnionRewrite;
}
/**
* Check query and view aggregate compatibility
*/
private static boolean checkCompatibility(
StructInfo queryStructInfo,
LogicalAggregate<Plan> queryAggregate, LogicalAggregate<Plan> viewAggregate,
MaterializationContext materializationContext) {
// if view is scalar aggregate but query is not. Or if query is scalar aggregate but view is not
// Should not rewrite
List<Expression> queryGroupByExpressions = queryAggregate.getGroupByExpressions();
List<Expression> viewGroupByExpressions = viewAggregate.getGroupByExpressions();
if (!queryGroupByExpressions.isEmpty() && viewGroupByExpressions.isEmpty()) {
materializationContext.recordFailReason(queryStructInfo,
"only one the of query or view is scalar aggregate and "
+ "can not rewrite expression meanwhile",
() -> String.format("query aggregate = %s,\n view aggregate = %s,\n",
queryAggregate.treeString(),
viewAggregate.treeString()));
return false;
}
boolean viewHasGroupSets = viewAggregate.getSourceRepeat()
.map(repeat -> repeat.getGroupingSets().size()).orElse(0) > 0;
// if both query and view has group sets, or query doesn't hava, mv have, not supported
if (viewHasGroupSets) {
materializationContext.recordFailReason(queryStructInfo,
"both query and view have group sets, or query doesn't have but view have, not supported",
() -> String.format("query aggregate = %s,\n view aggregate = %s,\n",
queryAggregate.treeString(),
viewAggregate.treeString()));
return false;
}
return true;
}
private boolean isGroupByEquals(Pair<Plan, LogicalAggregate<Plan>> queryTopPlanAndAggPair,
Pair<Plan, LogicalAggregate<Plan>> viewTopPlanAndAggPair,
SlotMapping viewToQuerySlotMapping,
StructInfo queryStructInfo,
StructInfo viewStructInfo,
Plan tempRewrittenPlan,
MaterializationContext materializationContext,
CascadesContext cascadesContext) {
if (materializationContext instanceof SyncMaterializationContext) {
// For data correctness, should always add aggregate node if rewritten by sync materialized view
return false;
}
Plan queryTopPlan = queryTopPlanAndAggPair.key();
Plan viewTopPlan = viewTopPlanAndAggPair.key();
LogicalAggregate<Plan> queryAggregate = queryTopPlanAndAggPair.value();
LogicalAggregate<Plan> viewAggregate = viewTopPlanAndAggPair.value();
Set<Expression> queryGroupByShuttledExpression = new HashSet<>(ExpressionUtils.shuttleExpressionWithLineage(
queryAggregate.getGroupByExpressions(), queryTopPlan));
// try to eliminate group by dimension by function dependency if group by expression is not in query
Map<Expression, Expression> viewShuttledExpressionQueryBasedToGroupByExpressionMap = new HashMap<>();
List<Expression> viewGroupByExpressions = viewAggregate.getGroupByExpressions();
List<? extends Expression> viewGroupByShuttledExpressions = ExpressionUtils.shuttleExpressionWithLineage(
viewGroupByExpressions, viewTopPlan);
for (int index = 0; index < viewGroupByExpressions.size(); index++) {
Expression viewExpression = viewGroupByExpressions.get(index);
Expression viewGroupExpressionQueryBased = ExpressionUtils.replace(
viewGroupByShuttledExpressions.get(index),
viewToQuerySlotMapping.toSlotReferenceMap());
viewShuttledExpressionQueryBasedToGroupByExpressionMap.put(viewGroupExpressionQueryBased,
viewExpression);
}
if (queryGroupByShuttledExpression.equals(viewShuttledExpressionQueryBasedToGroupByExpressionMap.keySet())) {
// return true, if equals directly
return true;
}
// Check is equals by equal filter eliminate
return isGroupByEqualsByFunctionDependency(
(LogicalPlan) tempRewrittenPlan,
queryGroupByShuttledExpression,
viewShuttledExpressionQueryBasedToGroupByExpressionMap,
materializationContext);
}
/**
* Check group by is equals by uniform function dependency
* For example query is:
* select
* a, b, c from t1
* where a = 1 and d = 'xx'
* group by a, b, c;
* mv is :
* select a, b, c, d
* from t1
* group by a, b, c, d;
* After group by key eliminate, the query group by is b, c
* but mv is group by a, b, c, d, the group by a and d of mv is more dimensions than the query
* But in tempRewrittenPlan is as following:
* select *
* from mv
* where a = 1 and d = 'xx'
* We can get group by a and d is uniform by function dependency info,
* so the group by expression between query and view is equals, should not aggregate roll up
* */
private static boolean isGroupByEqualsByFunctionDependency(
LogicalPlan tempRewrittenPlan,
Set<Expression> queryGroupShuttledExpression,
Map<Expression, Expression> viewShuttledExprQueryBasedToViewGroupByExprMap,
MaterializationContext materializationContext) {
Map<Expression, Expression> viewShuttledExprToScanExprMapping =
materializationContext.getShuttledExprToScanExprMapping().flattenMap().get(0);
Set<Expression> viewShuttledExprQueryBasedSet = viewShuttledExprQueryBasedToViewGroupByExprMap.keySet();
// view group by expr can not cover query group by expr
if (!viewShuttledExprQueryBasedSet.containsAll(queryGroupShuttledExpression)) {
return false;
}
Set<Expression> viewShouldRemovedExpressionSet = new HashSet<>();
Set<Expression> viewScanShouldReservedExpressionSet = new HashSet<>();
// calc the group by expr which is needed to roll up and should be uniform
for (Map.Entry<Expression, Expression> expressionMappingEntry :
viewShuttledExprQueryBasedToViewGroupByExprMap.entrySet()) {
if (queryGroupShuttledExpression.contains(expressionMappingEntry.getKey())) {
// the group expr which query has, do not require eliminate
viewScanShouldReservedExpressionSet.add(
viewShuttledExprToScanExprMapping.get(expressionMappingEntry.getValue()));
} else {
// the view expression which is more than query's expression, should try to eliminate
viewShouldRemovedExpressionSet.add(expressionMappingEntry.getValue());
}
}
DataTrait dataTrait = tempRewrittenPlan.computeDataTrait();
for (Expression viewShouldRemovedExpr : viewShouldRemovedExpressionSet) {
Expression viewScanExpression = viewShuttledExprToScanExprMapping.get(viewShouldRemovedExpr);
if (viewScanExpression == null) {
return false;
}
if (!(viewScanExpression instanceof Slot)) {
return false;
}
if (!dataTrait.isUniform((Slot) viewScanExpression)
&& Sets.intersection(dataTrait.calEqualSet((Slot) viewScanExpression),
viewScanShouldReservedExpressionSet).isEmpty()) {
// Such as query is l_orderkey#0, l_linenumber#1, o_custkey#17, l_partkey#2
// view is ps_partkey#25, o_orderkey#16, l_orderkey#0, l_linenumber#1, o_custkey#17, l_partkey#2
// If want to check the group by equality, ps_partkey#25, o_orderkey#16 should be uniform
// or should be equal any of [ l_orderkey#0, l_linenumber#1, o_custkey#17, l_partkey#2]
return false;
}
}
return true;
}
/**
* Roll up query aggregate function when query dimension num is less than mv dimension num,
*
* @param queryAggregateFunction query aggregate function to roll up.
* @param queryAggregateFunctionShuttled query aggregate function shuttled by lineage.
* @param mvExprToMvScanExprQueryBased mv def sql output expressions to mv result data output mapping.
* <p>
* Such as query is
* select max(a) + 1 from table group by b.
* mv is
* select max(a) from table group by a, b.
* the queryAggregateFunction is max(a), queryAggregateFunctionShuttled is max(a) + 1
* mvExprToMvScanExprQueryBased is { max(a) : MTMVScan(output#0) }
*/
private static Function rollup(AggregateFunction queryAggregateFunction,
Expression queryAggregateFunctionShuttled,
Map<Expression, Expression> mvExprToMvScanExprQueryBased) {
for (Map.Entry<Expression, Expression> expressionEntry : mvExprToMvScanExprQueryBased.entrySet()) {
Pair<Expression, Expression> mvExprToMvScanExprQueryBasedPair = Pair.of(expressionEntry.getKey(),
expressionEntry.getValue());
for (AggFunctionRollUpHandler rollUpHandler : ROLL_UP_HANDLERS) {
if (!rollUpHandler.canRollup(queryAggregateFunction, queryAggregateFunctionShuttled,
mvExprToMvScanExprQueryBasedPair, mvExprToMvScanExprQueryBased)) {
continue;
}
Function rollupFunction = rollUpHandler.doRollup(queryAggregateFunction,
queryAggregateFunctionShuttled, mvExprToMvScanExprQueryBasedPair,
mvExprToMvScanExprQueryBased);
if (rollupFunction != null) {
return rollupFunction;
}
}
}
return null;
}
protected Pair<Set<? extends Expression>, Set<? extends Expression>> topPlanSplitToGroupAndFunction(
Pair<Plan, LogicalAggregate<Plan>> topPlanAndAggPair, StructInfo queryStructInfo) {
LogicalAggregate<Plan> bottomQueryAggregate = topPlanAndAggPair.value();
Set<Expression> groupByExpressionSet = new HashSet<>(bottomQueryAggregate.getGroupByExpressions());
// when query is bitmap_count(bitmap_union), the plan is as following:
// project(bitmap_count()#1)
// aggregate(bitmap_union()#2)
// we should use exprId which query top plan used to decide the query top plan is use the
// bottom agg function or not
Set<ExprId> bottomAggregateFunctionExprIdSet = bottomQueryAggregate.getOutput().stream()
.filter(expr -> !groupByExpressionSet.contains(expr))
.map(NamedExpression::getExprId)
.collect(Collectors.toSet());
Plan queryTopPlan = topPlanAndAggPair.key();
Set<Expression> topGroupByExpressions = new HashSet<>();
Set<Expression> topFunctionExpressions = new HashSet<>();
queryTopPlan.getOutput().forEach(expression -> {
ExpressionLineageReplacer.ExpressionReplaceContext replaceContext =
new ExpressionLineageReplacer.ExpressionReplaceContext(ImmutableList.of(expression));
queryTopPlan.accept(ExpressionLineageReplacer.INSTANCE, replaceContext);
if (Collections.disjoint(bottomAggregateFunctionExprIdSet, replaceContext.getUsedExprIdSet())) {
topGroupByExpressions.add(expression);
} else {
// if query top plan expression use any aggregate function, then consider it is aggregate function
topFunctionExpressions.add(expression);
}
});
return Pair.of(topGroupByExpressions, topFunctionExpressions);
}
protected Pair<Plan, LogicalAggregate<Plan>> splitToTopPlanAndAggregate(StructInfo structInfo) {
Plan topPlan = structInfo.getTopPlan();
PlanSplitContext splitContext = new PlanSplitContext(Sets.newHashSet(LogicalAggregate.class));
topPlan.accept(StructInfo.PLAN_SPLITTER, splitContext);
if (!(splitContext.getBottomPlan() instanceof LogicalAggregate)) {
return null;
} else {
return Pair.of(topPlan, (LogicalAggregate<Plan>) splitContext.getBottomPlan());
}
}
/**
* Check Aggregate is simple or not and check join is whether valid or not.
* Support project, filter, join, logical relation node and join condition should only contain
* slot reference equals currently.
*/
@Override
protected boolean checkQueryPattern(StructInfo structInfo, CascadesContext cascadesContext) {
PlanCheckContext checkContext = PlanCheckContext.of(SUPPORTED_JOIN_TYPE_SET);
// if query or mv contains more than one top aggregate, should fail
return structInfo.getTopPlan().accept(StructInfo.PLAN_PATTERN_CHECKER, checkContext)
&& checkContext.isContainsTopAggregate() && checkContext.getTopAggregateNum() <= 1
&& !checkContext.isContainsTopLimit() && !checkContext.isContainsTopTopN()
&& !checkContext.isContainsTopWindow();
}
/**
* Aggregate expression rewriter which is responsible for rewriting group by and
* aggregate function expression
*/
protected static class AggregateExpressionRewriter
extends DefaultExpressionRewriter<AggregateExpressionRewriteContext> {
@Override
public Expression visitAggregateFunction(AggregateFunction aggregateFunction,
AggregateExpressionRewriteContext rewriteContext) {
if (!rewriteContext.isValid()) {
return aggregateFunction;
}
if (ExpressionRewriteMode.EXPRESSION_DIRECT.equals(rewriteContext.getExpressionRewriteMode())) {
rewriteContext.setValid(false);
return aggregateFunction;
}
Function rewrittenFunction;
if (ExpressionRewriteMode.EXPRESSION_ROLL_UP.equals(rewriteContext.getExpressionRewriteMode())) {
Expression queryFunctionShuttled = ExpressionUtils.shuttleExpressionWithLineage(
aggregateFunction,
rewriteContext.getQueryTopPlan());
rewrittenFunction = rollup(aggregateFunction, queryFunctionShuttled,
rewriteContext.getMvExprToMvScanExprQueryBasedMapping());
if (rewrittenFunction == null) {
rewriteContext.setValid(false);
return aggregateFunction;
}
return rewrittenFunction;
}
if (ExpressionRewriteMode.EXPRESSION_DIRECT_ALL.equals(rewriteContext.getExpressionRewriteMode())) {
List<Expression> children = aggregateFunction.children();
List<Expression> rewrittenChildren = new ArrayList<>();
for (Expression child : children) {
Expression rewrittenExpression = child.accept(this, rewriteContext);
if (!rewriteContext.isValid()) {
return aggregateFunction;
}
rewrittenChildren.add(rewrittenExpression);
}
return aggregateFunction.withChildren(rewrittenChildren);
}
rewriteContext.setValid(false);
return aggregateFunction;
}
@Override
public Expression visitSlot(Slot slot, AggregateExpressionRewriteContext rewriteContext) {
if (!rewriteContext.isValid()) {
return slot;
}
if (rewriteContext.getGroupingId().isPresent() && slot.equals(rewriteContext.getGroupingId().get())) {
return slot;
}
if (rewriteContext.getMvExprToMvScanExprQueryBasedMapping().containsKey(slot)) {
return rewriteContext.getMvExprToMvScanExprQueryBasedMapping().get(slot);
}
rewriteContext.setValid(false);
return slot;
}
@Override
public Expression visit(Expression expr, AggregateExpressionRewriteContext rewriteContext) {
if (!rewriteContext.isValid()) {
return expr;
}
// for group by expression try to get corresponding expression directly
if ((ExpressionRewriteMode.EXPRESSION_DIRECT.equals(rewriteContext.getExpressionRewriteMode())
|| ExpressionRewriteMode.EXPRESSION_DIRECT_ALL.equals(rewriteContext.getExpressionRewriteMode()))
&& rewriteContext.getMvExprToMvScanExprQueryBasedMapping().containsKey(expr)) {
return rewriteContext.getMvExprToMvScanExprQueryBasedMapping().get(expr);
}
List<Expression> newChildren = new ArrayList<>(expr.arity());
boolean hasNewChildren = false;
for (Expression child : expr.children()) {
Expression newChild = child.accept(this, rewriteContext);
if (!rewriteContext.isValid()) {
return expr;
}
if (newChild != child) {
hasNewChildren = true;
}
newChildren.add(newChild);
}
return hasNewChildren ? expr.withChildren(newChildren) : expr;
}
}
/**
* AggregateExpressionRewriteContext
*/
public static class AggregateExpressionRewriteContext {
private boolean valid = true;
private final ExpressionRewriteMode expressionRewriteMode;
private final Map<Expression, Expression> mvExprToMvScanExprQueryBasedMapping;
private final Plan queryTopPlan;
private final Optional<SlotReference> groupingId;
public AggregateExpressionRewriteContext(ExpressionRewriteMode expressionRewriteMode,
Map<Expression, Expression> mvExprToMvScanExprQueryBasedMapping, Plan queryTopPlan,
Optional<SlotReference> groupingId) {
this.expressionRewriteMode = expressionRewriteMode;
this.mvExprToMvScanExprQueryBasedMapping = mvExprToMvScanExprQueryBasedMapping;
this.queryTopPlan = queryTopPlan;
this.groupingId = groupingId;
}
public boolean isValid() {
return valid;
}
public void setValid(boolean valid) {
this.valid = valid;
}
public ExpressionRewriteMode getExpressionRewriteMode() {
return expressionRewriteMode;
}
public Map<Expression, Expression> getMvExprToMvScanExprQueryBasedMapping() {
return mvExprToMvScanExprQueryBasedMapping;
}
public Plan getQueryTopPlan() {
return queryTopPlan;
}
public Optional<SlotReference> getGroupingId() {
return groupingId;
}
/**
* The expression rewrite mode, which decide how the expression in query is rewritten by mv
*/
public enum ExpressionRewriteMode {
/**
* Try to use the expression in mv directly, and doesn't handle aggregate function
*/
EXPRESSION_DIRECT,
/**
* Try to use the expression in mv directly, and try to rewrite the arguments in aggregate function except
* the aggregate function
*/
EXPRESSION_DIRECT_ALL,
/**
* Try to roll up aggregate function
*/
EXPRESSION_ROLL_UP
}
}
}