[fix](Nereids): just pull up alias project above join through topn (#32305)

This commit is contained in:
jakevin
2024-03-22 10:00:03 +08:00
committed by yiguolei
parent 66336e59e6
commit 39382a9774
58 changed files with 1233 additions and 1220 deletions

View File

@ -21,6 +21,7 @@ import org.apache.doris.nereids.rules.Rule;
import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalJoin;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.util.PlanUtils;
@ -36,9 +37,10 @@ import java.util.stream.Collectors;
public class PullUpProjectUnderLimit extends OneRewriteRuleFactory {
@Override
public Rule build() {
return logicalLimit(logicalProject().whenNot(p -> p.isAllSlots()))
return logicalLimit(logicalProject(logicalJoin().when(j -> j.getJoinType().isLeftRightOuterOrCrossJoin()))
.whenNot(p -> p.isAllSlots()))
.then(limit -> {
LogicalProject<Plan> project = limit.child();
LogicalProject<LogicalJoin<Plan, Plan>> project = limit.child();
Set<Slot> allUsedSlots = project.getProjects().stream().flatMap(ne -> ne.getInputSlots().stream())
.collect(Collectors.toSet());
Set<Slot> outputSet = project.child().getOutputSet();

View File

@ -20,41 +20,61 @@ package org.apache.doris.nereids.rules.rewrite;
import org.apache.doris.nereids.properties.OrderKey;
import org.apache.doris.nereids.rules.Rule;
import org.apache.doris.nereids.rules.RuleType;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalJoin;
import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
import org.apache.doris.nereids.trees.plans.logical.LogicalTopN;
import org.apache.doris.nereids.util.ExpressionUtils;
import org.apache.doris.nereids.util.PlanUtils;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
/**
* Pull up Project under TopN.
* Pull up Project under TopN for PushDownTopNThroughJoin
*/
public class PullUpProjectUnderTopN extends OneRewriteRuleFactory {
@Override
public Rule build() {
return logicalTopN(logicalProject().whenNot(p -> p.isAllSlots()))
.whenNot(topN -> topN.child().hasPushedDownToProjectionFunctions())
return logicalTopN(
logicalProject(logicalJoin().when(j -> j.getJoinType().isLeftRightOuterOrCrossJoin()))
.whenNot(p -> p.isAllSlots()))
.then(topN -> {
LogicalProject<Plan> project = topN.child();
LogicalProject<LogicalJoin<Plan, Plan>> project = topN.child();
Set<Slot> outputSet = project.child().getOutputSet();
if (!topN.getOrderKeys().stream().map(OrderKey::getExpr).flatMap(e -> e.getInputSlots().stream())
.allMatch(outputSet::contains)) {
return null;
Map<Slot, Expression> slotMap = ExpressionUtils.generateReplaceMap(project.getProjects());
List<OrderKey> newOrderKeys = new ArrayList<>();
for (OrderKey orderKey : topN.getOrderKeys()) {
if (!(orderKey.getExpr() instanceof Slot)) {
return null;
}
Expression expression = slotMap.get((Slot) orderKey.getExpr());
if (expression instanceof Slot) {
newOrderKeys.add(orderKey.withExpression(expression));
} else {
return null;
}
}
Set<Slot> allUsedSlots = project.getProjects().stream().flatMap(ne -> ne.getInputSlots().stream())
.collect(Collectors.toSet());
LogicalTopN<Plan> newTopN = topN.withOrderKeys(newOrderKeys);
if (outputSet.size() == allUsedSlots.size()) {
Preconditions.checkState(outputSet.equals(allUsedSlots));
return project.withChildren(topN.withChildren(project.child()));
return project.withChildren(newTopN.withChildren(project.child()));
} else {
Plan columnProject = PlanUtils.projectOrSelf(ImmutableList.copyOf(allUsedSlots),
project.child());
return project.withChildren(topN.withChildren(columnProject));
return project.withChildren(newTopN.withChildren(columnProject));
}
}).toRule(RuleType.PULL_UP_PROJECT_UNDER_TOPN);
}

View File

@ -152,6 +152,10 @@ public enum JoinType {
return this == RIGHT_OUTER_JOIN;
}
public final boolean isLeftRightOuterOrCrossJoin() {
return this == LEFT_OUTER_JOIN || this == RIGHT_OUTER_JOIN || this == CROSS_JOIN;
}
public final boolean isLeftSemiOrAntiJoin() {
return this == LEFT_SEMI_JOIN || this == LEFT_ANTI_JOIN || this == NULL_AWARE_LEFT_ANTI_JOIN;
}