From 5952b05b1bf8119f66275f68460888c8fae07649 Mon Sep 17 00:00:00 2001 From: jakevin Date: Fri, 18 Aug 2023 18:38:12 +0800 Subject: [PATCH] [chore](Nereids): optimize to handle enforcer in MergeGroup() (#22709) (#23183) --- .../org/apache/doris/nereids/memo/Group.java | 36 +++++----- .../org/apache/doris/nereids/memo/Memo.java | 65 ++++++++++++------- 2 files changed, 62 insertions(+), 39 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Group.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Group.java index 7fe2897a87175f5..c816151731d6ec8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Group.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Group.java @@ -27,7 +27,6 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; 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.physical.PhysicalDistribute; import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan; import org.apache.doris.nereids.util.TreeStringUtils; import org.apache.doris.nereids.util.Utils; @@ -59,6 +58,8 @@ public class Group { private final List logicalExpressions = Lists.newArrayList(); private final List physicalExpressions = Lists.newArrayList(); + private final List enforcers = Lists.newArrayList(); + private LogicalProperties logicalProperties; // Map of cost lower bounds @@ -210,6 +211,15 @@ public GroupExpression getBestPlan(PhysicalProperties properties) { return null; } + public void addEnforcer(GroupExpression enforcer) { + enforcer.setOwnerGroup(this); + enforcers.add(enforcer); + } + + public List getEnforcers() { + return enforcers; + } + /** * Set or update lowestCostPlans: properties --> Pair.of(cost, expression) */ @@ -308,12 +318,12 @@ public void removeParentPhysicalExpressions() { public void mergeTo(Group target) { // move parentExpressions Ownership parentExpressions.keySet().forEach(parent -> target.addParentExpression(parent)); - // PhysicalEnforcer isn't in groupExpressions, so mergeGroup() can't replace its children. - // So we need to manually replace the children of PhysicalEnforcer in here. - // TODO: SortEnforcer? - parentExpressions.keySet().stream().filter(ge -> ge.getPlan() instanceof PhysicalDistribute) - .forEach(ge -> ge.children().set(0, target)); - parentExpressions.clear(); + + // move enforcers Ownership + enforcers.forEach(ge -> ge.children().set(0, target)); + // TODO: dedup? + enforcers.forEach(enforcer -> target.addEnforcer(enforcer)); + enforcers.clear(); // move LogicalExpression PhysicalExpression Ownership Map logicalSet = target.getLogicalExpressions().stream() @@ -345,15 +355,7 @@ public void mergeTo(Group target) { physicalExpressions.clear(); // Above we already replaceBestPlanGroupExpr, but we still need to moveLowestCostPlansOwnership. - // Because PhysicalEnforcer don't exist in physicalExpressions, so above `replaceBestPlanGroupExpr` can't - // move PhysicalEnforcer in lowestCostPlans. Following code can move PhysicalEnforcer in lowestCostPlans. lowestCostPlans.forEach((physicalProperties, costAndGroupExpr) -> { - GroupExpression bestGroupExpression = costAndGroupExpr.second; - if (bestGroupExpression.getOwnerGroup() == this || bestGroupExpression.getOwnerGroup() == null) { - // move PhysicalEnforcer into target - Preconditions.checkState(bestGroupExpression.getPlan() instanceof PhysicalDistribute); - bestGroupExpression.setOwnerGroup(target); - } // move lowestCostPlans Ownership if (!target.lowestCostPlans.containsKey(physicalProperties)) { target.lowestCostPlans.put(physicalProperties, costAndGroupExpr); @@ -425,6 +427,10 @@ public String toString() { for (GroupExpression physicalExpression : physicalExpressions) { str.append(" ").append(physicalExpression).append("\n"); } + str.append(" enforcers:\n"); + for (GroupExpression enforcer : enforcers) { + str.append(" ").append(enforcer).append("\n"); + } return str.toString(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java index ad7b3d402d9f1a1..4204f33ca1591be 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java @@ -34,7 +34,6 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; 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.physical.PhysicalDistribute; import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan; import org.apache.doris.qe.ConnectContext; @@ -46,10 +45,11 @@ import org.apache.logging.log4j.Logger; import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; -import java.util.Map.Entry; import java.util.Optional; import java.util.PriorityQueue; import java.util.stream.Collectors; @@ -115,7 +115,7 @@ public int getGroupExpressionsSize() { public void removePhysicalExpression() { groupExpressions.entrySet().removeIf(entry -> entry.getValue().getPlan() instanceof PhysicalPlan); - Iterator> iterator = groups.entrySet().iterator(); + Iterator> iterator = groups.entrySet().iterator(); while (iterator.hasNext()) { Map.Entry entry = iterator.next(); Group group = entry.getValue(); @@ -199,6 +199,20 @@ public Pair countGroupJoin(Group group) { return Pair.of(continuousJoinCount, Math.max(continuousJoinCount, maxJoinCount)); } + /** + * Add plan to Memo. + */ + public CopyInResult copyIn(Plan plan, @Nullable Group target, boolean rewrite, HashMap planTable) { + CopyInResult result; + if (rewrite) { + result = doRewrite(plan, target); + } else { + result = doCopyIn(skipProject(plan, target), target, planTable); + } + maybeAddStateId(result); + return result; + } + /** * Add plan to Memo. * @@ -214,7 +228,7 @@ public CopyInResult copyIn(Plan plan, @Nullable Group target, boolean rewrite) { if (rewrite) { result = doRewrite(plan, target); } else { - result = doCopyIn(skipProject(plan, target), target); + result = doCopyIn(skipProject(plan, target), target, null); } maybeAddStateId(result); return result; @@ -402,7 +416,7 @@ private CopyInResult doRewrite(Plan plan, @Nullable Group targetGroup) { * @return a pair, in which the first element is true if a newly generated groupExpression added into memo, * and the second element is a reference of node in Memo */ - private CopyInResult doCopyIn(Plan plan, @Nullable Group targetGroup) { + private CopyInResult doCopyIn(Plan plan, @Nullable Group targetGroup, @Nullable HashMap planTable) { Preconditions.checkArgument(!(plan instanceof GroupPlan), "plan can not be GroupPlan"); // check logicalproperties, must same output in a Group. if (targetGroup != null && !plan.getLogicalProperties().equals(targetGroup.getLogicalProperties())) { @@ -425,12 +439,12 @@ private CopyInResult doCopyIn(Plan plan, @Nullable Group targetGroup) { } else if (child.getGroupExpression().isPresent()) { childrenGroups.add(child.getGroupExpression().get().getOwnerGroup()); } else { - childrenGroups.add(doCopyIn(child, null).correspondingExpression.getOwnerGroup()); + childrenGroups.add(doCopyIn(child, null, planTable).correspondingExpression.getOwnerGroup()); } } plan = replaceChildrenToGroupPlan(plan, childrenGroups); GroupExpression newGroupExpression = new GroupExpression(plan, childrenGroups); - return insertGroupExpression(newGroupExpression, targetGroup, plan.getLogicalProperties()); + return insertGroupExpression(newGroupExpression, targetGroup, plan.getLogicalProperties(), planTable); // TODO: need to derive logical property if generate new group. currently we not copy logical plan into } @@ -474,12 +488,12 @@ private void validateRewriteChildGroup(Group childGroup, Group targetGroup) { * @return a pair, in which the first element is true if a newly generated groupExpression added into memo, * and the second element is a reference of node in Memo */ - private CopyInResult insertGroupExpression( - GroupExpression groupExpression, Group target, LogicalProperties logicalProperties) { + private CopyInResult insertGroupExpression(GroupExpression groupExpression, Group target, + LogicalProperties logicalProperties, HashMap planTable) { GroupExpression existedGroupExpression = groupExpressions.get(groupExpression); if (existedGroupExpression != null) { if (target != null && !target.getGroupId().equals(existedGroupExpression.getOwnerGroup().getGroupId())) { - mergeGroup(existedGroupExpression.getOwnerGroup(), target); + mergeGroup(target, existedGroupExpression.getOwnerGroup(), planTable); } // When we create a GroupExpression, we will add it into ParentExpression of childGroup. // But if it already exists, we should remove it from ParentExpression of childGroup. @@ -506,7 +520,7 @@ private CopyInResult insertGroupExpression( * @param source source group * @param destination destination group */ - public void mergeGroup(Group source, Group destination) { + public void mergeGroup(Group source, Group destination, HashMap planTable) { if (source.equals(destination)) { return; } @@ -516,9 +530,9 @@ public void mergeGroup(Group source, Group destination) { // cycle, we should not merge return; } - // PhysicalEnforcer don't exist in memo, so we need skip them. - if (parent.getPlan() instanceof PhysicalDistribute) { - // TODO: SortEnforcer. + Group parentOwnerGroup = parent.getOwnerGroup(); + HashSet enforcers = new HashSet<>(parentOwnerGroup.getEnforcers()); + if (enforcers.contains(parent)) { continue; } needReplaceChild.add(parent); @@ -545,25 +559,28 @@ public void mergeGroup(Group source, Group destination) { reinsertGroupExpr.mergeTo(existGroupExpr); } else { // reinsertGroupExpr & existGroupExpr aren't in same group, need to merge their OwnerGroup. - mergeGroup(reinsertGroupExpr.getOwnerGroup(), existGroupExpr.getOwnerGroup()); + mergeGroup(reinsertGroupExpr.getOwnerGroup(), existGroupExpr.getOwnerGroup(), planTable); } } else { groupExpressions.put(reinsertGroupExpr, reinsertGroupExpr); } } + // replace source with destination in groups of planTable + if (planTable != null) { + planTable.forEach((bitset, group) -> { + if (group.equals(source)) { + planTable.put(bitset, destination); + } + }); + } + source.mergeTo(destination); + if (source == root) { + root = destination; + } groups.remove(source.getGroupId()); } - /** - * Add enforcer expression into the target group. - */ - public void addEnforcerPlan(GroupExpression groupExpression, Group group) { - Preconditions.checkArgument(groupExpression != null); - groupExpression.setOwnerGroup(group); - // Don't add groupExpression into group's physicalExpressions, it will cause dead loop; - } - private CopyInResult rewriteByExistedPlan(Group targetGroup, Plan existedPlan) { GroupExpression existedLogicalExpression = existedPlan instanceof GroupPlan ? ((GroupPlan) existedPlan).getGroup().getLogicalExpression() // get first logicalGroupExpression