From f2dca848dbb613ed5d7681e0bc9f4bc17f756d04 Mon Sep 17 00:00:00 2001
From: jakevin <jakevingoo@gmail.com>
Date: Tue, 8 Aug 2023 16:56:34 +0800
Subject: [PATCH] [chore](Nereids): optimize to handle enforcer in MergeGroup()
 (#22709)

---
 .../org/apache/doris/nereids/memo/Group.java     | 16 ++++++----------
 .../java/org/apache/doris/nereids/memo/Memo.java |  8 ++++----
 2 files changed, 10 insertions(+), 14 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 7d499ae42b..c816151731 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.Plan;
 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;
@@ -217,6 +216,10 @@ public class Group {
         enforcers.add(enforcer);
     }
 
+    public List<GroupExpression> getEnforcers() {
+        return enforcers;
+    }
+
     /**
      * Set or update lowestCostPlans: properties --> Pair.of(cost, expression)
      */
@@ -316,10 +319,11 @@ public class Group {
         // move parentExpressions Ownership
         parentExpressions.keySet().forEach(parent -> target.addParentExpression(parent));
 
-        // TODO: dedup?
         // 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<GroupExpression, GroupExpression> logicalSet = target.getLogicalExpressions().stream()
@@ -351,15 +355,7 @@ public class Group {
         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);
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 3ae0459c02..4204f33ca1 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.Plan;
 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;
 
@@ -47,6 +46,7 @@ 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;
@@ -530,9 +530,9 @@ public class Memo {
                 // 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<GroupExpression> enforcers = new HashSet<>(parentOwnerGroup.getEnforcers());
+            if (enforcers.contains(parent)) {
                 continue;
             }
             needReplaceChild.add(parent);