[fix](Nereids): just pull up alias project above join through topn (#32305)
This commit is contained in:
@ -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();
|
||||
|
||||
@ -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);
|
||||
}
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user