database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feat](nereids) add merge aggregate rule (#31811)

Browse Source
This commit is contained in:
feiniaofeiafei
2024-03-13 12:05:47 +08:00
committed by yiguolei
parent fd1e0e933e
commit 99830b77a8
8 changed files with 695 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java
View File

@ -81,6 +81,7 @@ import org.apache.doris.nereids.rules.rewrite.InferPredicates;
import org.apache.doris.nereids.rules.rewrite.InferSetOperatorDistinct;
import org.apache.doris.nereids.rules.rewrite.InlineLogicalView;
import org.apache.doris.nereids.rules.rewrite.LimitSortToTopN;
import org.apache.doris.nereids.rules.rewrite.MergeAggregate;
import org.apache.doris.nereids.rules.rewrite.MergeFilters;
import org.apache.doris.nereids.rules.rewrite.MergeOneRowRelationIntoUnion;
import org.apache.doris.nereids.rules.rewrite.MergeProjects;
@ -341,7 +342,8 @@ public class Rewriter extends AbstractBatchJobExecutor {
),
topic("Eliminate GroupBy",
topDown(new EliminateGroupBy())
topDown(new EliminateGroupBy(),
new MergeAggregate())
),
topic("Eager aggregation",

1
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java
View File

@ -196,6 +196,7 @@ public enum RuleType {
MERGE_LIMITS(RuleTypeClass.REWRITE),
MERGE_GENERATES(RuleTypeClass.REWRITE),
// Eliminate plan
MERGE_AGGREGATE(RuleTypeClass.REWRITE),
ELIMINATE_AGGREGATE(RuleTypeClass.REWRITE),
ELIMINATE_LIMIT(RuleTypeClass.REWRITE),
ELIMINATE_LIMIT_ON_ONE_ROW_RELATION(RuleTypeClass.REWRITE),

46
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ColumnPruning.java
View File

@ -151,9 +151,7 @@ public class ColumnPruning extends DefaultPlanRewriter<PruneContext> implements
if (union.getQualifier() == Qualifier.DISTINCT) {
return skipPruneThisAndFirstLevelChildren(union);
}
LogicalUnion prunedOutputUnion = pruneOutput(union, union.getOutputs(), union::pruneOutputs, context);
LogicalUnion prunedOutputUnion = pruneUnionOutput(union, context);
// start prune children of union
List<Slot> originOutput = union.getOutput();
Set<Slot> prunedOutput = prunedOutputUnion.getOutputSet();
@ -303,6 +301,48 @@ public class ColumnPruning extends DefaultPlanRewriter<PruneContext> implements
}
}
private LogicalUnion pruneUnionOutput(LogicalUnion union, PruneContext context) {
List<NamedExpression> originOutput = union.getOutputs();
if (originOutput.isEmpty()) {
return union;
}
List<NamedExpression> prunedOutputs = Lists.newArrayList();
List<List<NamedExpression>> constantExprsList = union.getConstantExprsList();
List<List<NamedExpression>> prunedConstantExprsList = Lists.newArrayList();
List<Integer> extractColumnIndex = Lists.newArrayList();
for (int i = 0; i < originOutput.size(); i++) {
NamedExpression output = originOutput.get(i);
if (context.requiredSlots.contains(output.toSlot())) {
prunedOutputs.add(output);
extractColumnIndex.add(i);
}
}
int len = extractColumnIndex.size();
for (List<NamedExpression> row : constantExprsList) {
ArrayList<NamedExpression> newRow = new ArrayList<>(len);
for (int idx : extractColumnIndex) {
newRow.add(row.get(idx));
}
prunedConstantExprsList.add(newRow);
}
if (prunedOutputs.isEmpty()) {
List<NamedExpression> candidates = Lists.newArrayList(originOutput);
candidates.retainAll(keys);
if (candidates.isEmpty()) {
candidates = originOutput;
}
NamedExpression minimumColumn = ExpressionUtils.selectMinimumColumn(candidates);
prunedOutputs = ImmutableList.of(minimumColumn);
}
if (prunedOutputs.equals(originOutput)) {
return union;
} else {
return union.withNewOutputsAndConstExprsList(prunedOutputs, prunedConstantExprsList);
}
}
private <P extends Plan> P pruneChildren(P plan) {
return pruneChildren(plan, ImmutableSet.of());
}

211
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/MergeAggregate.java Normal file
View File

@ -0,0 +1,211 @@
// 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.rewrite;
import org.apache.doris.nereids.rules.Rule;
import org.apache.doris.nereids.rules.RuleType;
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.SlotReference;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
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.LogicalProject;
import org.apache.doris.nereids.util.ExpressionUtils;
import org.apache.doris.nereids.util.PlanUtils;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
/**MergeAggregate*/
public class MergeAggregate implements RewriteRuleFactory {
private static final ImmutableSet<String> ALLOW_MERGE_AGGREGATE_FUNCTIONS =
ImmutableSet.of("min", "max", "sum", "any_value");
private Map<ExprId, AggregateFunction> innerAggExprIdToAggFunc = new HashMap<>();
@Override
public List<Rule> buildRules() {
return ImmutableList.of(
logicalAggregate(logicalAggregate()).when(this::canMergeAggregateWithoutProject)
.then(this::mergeTwoAggregate)
.toRule(RuleType.MERGE_AGGREGATE),
logicalAggregate(logicalProject(logicalAggregate()))
.when(this::canMergeAggregateWithProject)
.then(this::mergeAggProjectAgg)
.toRule(RuleType.MERGE_AGGREGATE));
}
/**
* before:
* LogicalAggregate
* +--LogicalAggregate
* after:
* LogicalAggregate
*/
private Plan mergeTwoAggregate(LogicalAggregate<LogicalAggregate<Plan>> outerAgg) {
LogicalAggregate<Plan> innerAgg = outerAgg.child();
List<NamedExpression> newOutputExpressions = outerAgg.getOutputExpressions().stream()
.map(e -> rewriteAggregateFunction(e, innerAggExprIdToAggFunc))
.collect(Collectors.toList());
return outerAgg.withAggOutput(newOutputExpressions).withChildren(innerAgg.children());
}
/**
* before:
* LogicalAggregate (outputExpressions = [col2, sum(col1)], groupByKeys = [col2])
* +--LogicalProject (projects = [a as col2, col1])
* +--LogicalAggregate (outputExpressions = [a, b, sum(c) as col1], groupByKeys = [a,b])
* after:
* LogicalProject (projects = [a as col2, sum(col1) as sum(col1)]
* +--LogicalAggregate (outputExpression = [a, sum(c) as sum(col1)], groupByKeys = [a])
*/
private Plan mergeAggProjectAgg(LogicalAggregate<LogicalProject<LogicalAggregate<Plan>>> outerAgg) {
LogicalProject<LogicalAggregate<Plan>> project = outerAgg.child();
LogicalAggregate<Plan> innerAgg = project.child();
// rewrite agg function. e.g. max(max)
List<NamedExpression> aggFunc = outerAgg.getOutputExpressions().stream()
.filter(expr -> (expr instanceof Alias) && (expr.child(0) instanceof AggregateFunction))
.map(e -> rewriteAggregateFunction(e, innerAggExprIdToAggFunc))
.collect(Collectors.toList());
// rewrite agg function directly refer to the slot below the project
List<Expression> replacedAggFunc = PlanUtils.replaceExpressionByProjections(project.getProjects(),
(List) aggFunc);
// replace groupByKeys directly refer to the slot below the project
List<Expression> replacedGroupBy = PlanUtils.replaceExpressionByProjections(project.getProjects(),
outerAgg.getGroupByExpressions());
List<NamedExpression> newOutputExpressions = ImmutableList.<NamedExpression>builder()
.addAll(replacedGroupBy.stream().map(slot -> (NamedExpression) slot).iterator())
.addAll(replacedAggFunc.stream().map(alias -> (NamedExpression) alias).iterator()).build();
// construct agg
LogicalAggregate<Plan> resAgg = outerAgg.withGroupByAndOutput(replacedGroupBy, newOutputExpressions)
.withChildren(innerAgg.children());
// construct upper project
Map<SlotReference, Alias> childToAlias = project.getProjects().stream()
.filter(expr -> (expr instanceof Alias) && (expr.child(0) instanceof SlotReference))
.collect(Collectors.toMap(alias -> (SlotReference) alias.child(0), alias -> (Alias) alias));
List<Expression> projectGroupBy = ExpressionUtils.replace(replacedGroupBy, childToAlias);
List<NamedExpression> upperProjects = ImmutableList.<NamedExpression>builder()
.addAll(projectGroupBy.stream().map(namedExpr -> (NamedExpression) namedExpr).iterator())
.addAll(replacedAggFunc.stream().map(expr -> ((NamedExpression) expr).toSlot()).iterator())
.build();
return new LogicalProject<Plan>(upperProjects, resAgg);
}
private NamedExpression rewriteAggregateFunction(NamedExpression e,
Map<ExprId, AggregateFunction> innerAggExprIdToAggFunc) {
return (NamedExpression) e.rewriteDownShortCircuit(expr -> {
if (expr instanceof Alias && ((Alias) expr).child() instanceof AggregateFunction) {
Alias alias = (Alias) expr;
AggregateFunction aggFunc = (AggregateFunction) alias.child();
ExprId childExprId = ((SlotReference) aggFunc.child(0)).getExprId();
if (innerAggExprIdToAggFunc.containsKey(childExprId)) {
return new Alias(alias.getExprId(), innerAggExprIdToAggFunc.get(childExprId),
alias.getName());
} else {
return expr;
}
} else {
return expr;
}
});
}
boolean commonCheck(LogicalAggregate<? extends Plan> outerAgg, LogicalAggregate<Plan> innerAgg,
boolean sameGroupBy) {
innerAggExprIdToAggFunc = innerAgg.getOutputExpressions().stream()
.filter(expr -> (expr instanceof Alias) && (expr.child(0) instanceof AggregateFunction))
.collect(Collectors.toMap(NamedExpression::getExprId, value -> (AggregateFunction) value.child(0),
(existValue, newValue) -> existValue));
Set<AggregateFunction> aggregateFunctions = outerAgg.getAggregateFunctions();
for (AggregateFunction outerFunc : aggregateFunctions) {
if (!(ALLOW_MERGE_AGGREGATE_FUNCTIONS.contains(outerFunc.getName()))) {
return false;
}
if (outerFunc.isDistinct() && !sameGroupBy) {
return false;
}
// not support outerAggFunc: sum(a+1),sum(a+b)
if (!(outerFunc.child(0) instanceof SlotReference)) {
return false;
}
ExprId childExprId = ((SlotReference) outerFunc.child(0)).getExprId();
if (innerAggExprIdToAggFunc.containsKey(childExprId)) {
AggregateFunction innerFunc = innerAggExprIdToAggFunc.get(childExprId);
if (innerFunc.isDistinct() && !sameGroupBy) {
return false;
}
// support sum(sum),min(min),max(max),any_value(any_value),sum(count)
// sum(count) -> count() need outerAgg having group by keys (reason: nullable)
if (!(outerFunc.getName().equals("sum") && innerFunc.getName().equals("count")
&& !outerAgg.getGroupByExpressions().isEmpty())
&& !innerFunc.getName().equals(outerFunc.getName())) {
return false;
}
} else {
// select a, max(b), min(b), any_value(b) from (select a,b from t1 group by a, b) group by a;
// equals select a, max(b), min(b), any_value(b) from t1 group by a;
if (!outerFunc.getName().equals("max")
&& !outerFunc.getName().equals("min")
&& !outerFunc.getName().equals("any_value")) {
return false;
}
}
}
return true;
}
private boolean canMergeAggregateWithoutProject(LogicalAggregate<LogicalAggregate<Plan>> outerAgg) {
LogicalAggregate<Plan> innerAgg = outerAgg.child();
if (!new HashSet<>(innerAgg.getGroupByExpressions()).containsAll(outerAgg.getGroupByExpressions())) {
return false;
}
boolean sameGroupBy = (innerAgg.getGroupByExpressions().size() == outerAgg.getGroupByExpressions().size());
return commonCheck(outerAgg, innerAgg, sameGroupBy);
}
private boolean canMergeAggregateWithProject(LogicalAggregate<LogicalProject<LogicalAggregate<Plan>>> outerAgg) {
LogicalProject<LogicalAggregate<Plan>> project = outerAgg.child();
LogicalAggregate<Plan> innerAgg = project.child();
List<Expression> outerAggGroupByKeys = PlanUtils.replaceExpressionByProjections(project.getProjects(),
outerAgg.getGroupByExpressions());
if (!new HashSet<>(innerAgg.getGroupByExpressions()).containsAll(outerAggGroupByKeys)) {
return false;
}
// project cannot have expressions like a+1
if (ExpressionUtils.anyMatch(project.getProjects(),
expr -> !(expr instanceof SlotReference) && !(expr instanceof Alias))) {
return false;
}
boolean sameGroupBy = (innerAgg.getGroupByExpressions().size() == outerAgg.getGroupByExpressions().size());
return commonCheck(outerAgg, innerAgg, sameGroupBy);
}
}

6
fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalUnion.java
View File

@ -165,6 +165,12 @@ public class LogicalUnion extends LogicalSetOperation implements Union, OutputPr
hasPushedFilter, Optional.empty(), Optional.empty(), children);
}
public LogicalUnion withNewOutputsAndConstExprsList(List<NamedExpression> newOutputs,
List<List<NamedExpression>> constantExprsList) {
return new LogicalUnion(qualifier, newOutputs, regularChildrenOutputs, constantExprsList,
hasPushedFilter, Optional.empty(), Optional.empty(), children);
}
public LogicalUnion withChildrenAndConstExprsList(List<Plan> children,
List<List<SlotReference>> childrenOutputs, List<List<NamedExpression>> constantExprsList) {
return new LogicalUnion(qualifier, outputs, childrenOutputs, constantExprsList, hasPushedFilter, children);

6
fe/fe-core/src/main/java/org/apache/doris/nereids/util/PlanUtils.java
View File

@ -100,6 +100,12 @@ public class PlanUtils {
return ExpressionUtils.replaceNamedExpressions(parentProjects, replaceMap);
}
public static List<Expression> replaceExpressionByProjections(List<NamedExpression> childProjects,
List<Expression> targetExpression) {
Map<Slot, Expression> replaceMap = ExpressionUtils.generateReplaceMap(childProjects);
return ExpressionUtils.replace(targetExpression, replaceMap);
}
public static Plan skipProjectFilterLimit(Plan plan) {
if (plan instanceof LogicalProject && ((LogicalProject<?>) plan).isAllSlots()
|| plan instanceof LogicalFilter || plan instanceof LogicalLimit) {