keep group id in plan

apache · Aug 9, 2023 · 877858b · 877858b
1 parent 4e05753
commit 877858b
Show file tree

Hide file tree

Showing 21 changed files with 50 additions and 37 deletions.
diff --git a/...ain/java/org/apache/doris/nereids/processor/post/RecomputeLogicalPropertiesProcessor.java b/...ain/java/org/apache/doris/nereids/processor/post/RecomputeLogicalPropertiesProcessor.java
@@ -20,6 +20,7 @@
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalPlan;
+import org.apache.doris.nereids.util.MutableState;
 
 /**
  * merge consecutive projects
@@ -28,6 +29,8 @@ public class RecomputeLogicalPropertiesProcessor extends PlanPostProcessor {
     @Override
     public Plan visit(Plan plan, CascadesContext ctx) {
         PhysicalPlan physicalPlan = (PhysicalPlan) visitChildren(this, plan, ctx);
-        return physicalPlan.resetLogicalProperties();
+        physicalPlan.resetLogicalProperties();
+        physicalPlan.setMutableState(MutableState.KEY_GROUP, plan.getGroupIdAsString());
+        return physicalPlan;
     }
 }
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/AbstractPlan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/AbstractPlan.java
@@ -209,23 +209,6 @@ public void setMutableState(String key, Object state) {
         this.mutableState = this.mutableState.set(key, state);
     }
 
-    /**
-     * used in treeString()
-     *
-     * @return "" if groupExpression is empty, o.w. string format of group id
-     */
-    public String getGroupIdAsString() {
-        String groupId;
-        if (getGroupExpression().isPresent()) {
-            groupId = "@" + groupExpression.get().getOwnerGroup().getGroupId().asInt();
-        } else if (getMutableState("group").isPresent()) {
-            groupId = "@" + getMutableState("group").get();
-        } else {
-            groupId = "";
-        }
-        return groupId;
-    }
-
     @Override
     public boolean deepEquals(TreeNode o) {
         AbstractPlan that = (AbstractPlan) o;

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/Plan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/Plan.java
@@ -26,6 +26,7 @@
 import org.apache.doris.nereids.trees.expressions.NamedExpression;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
+import org.apache.doris.nereids.util.MutableState;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -164,4 +165,25 @@ default String shape(String prefix) {
     default String shapeInfo() {
         return this.getClass().getSimpleName();
     }
+
+    /**
+     * used in treeString()
+     *
+     * @return "" if groupExpression is empty, o.w. string format of group id
+     */
+    default String getGroupIdAsString() {
+        String groupId;
+        if (getGroupExpression().isPresent()) {
+            groupId = getGroupExpression().get().getOwnerGroup().getGroupId().asInt() + "";
+        } else if (getMutableState(MutableState.KEY_GROUP).isPresent()) {
+            groupId = getMutableState(MutableState.KEY_GROUP).get().toString();
+        } else {
+            groupId = "";
+        }
+        return groupId;
+    }
+
+    default String getGroupIdWithPrefix() {
+        return "@" + getGroupIdAsString();
+    }
 }
diff --git a/...re/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalAssertNumRows.java b/...re/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalAssertNumRows.java
@@ -70,7 +70,7 @@ public AssertNumRowsElement getAssertNumRowsElement() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalAssertNumRows" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalAssertNumRows" + getGroupIdWithPrefix(),
                 "assertNumRowsElement", assertNumRowsElement);
     }
 

diff --git a/...-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalDistribute.java b/...-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalDistribute.java
@@ -72,7 +72,7 @@ public PhysicalDistribute(DistributionSpec spec, Optional<GroupExpression> group
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalDistribute[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalDistribute[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "distributionSpec", distributionSpec,
                 "stats", statistics
         );

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalFilter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalFilter.java
@@ -74,7 +74,7 @@ public List<Expression> getExpressions() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalFilter[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalFilter[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "predicates", getPredicate(),
                 "stats", statistics
         );

diff --git a/...re/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java b/...re/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashAggregate.java
@@ -196,7 +196,7 @@ public List<? extends Expression> getExpressions() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalHashAggregate[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalHashAggregate[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "aggPhase", aggregateParam.aggPhase,
                 "aggMode", aggregateParam.aggMode,
                 "maybeUseStreaming", maybeUsingStream,

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashJoin.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalHashJoin.java
@@ -36,6 +36,7 @@
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.PlanType;
 import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
+import org.apache.doris.nereids.util.MutableState;
 import org.apache.doris.nereids.util.Utils;
 import org.apache.doris.planner.RuntimeFilterId;
 import org.apache.doris.statistics.Statistics;
@@ -161,7 +162,7 @@ public String toString() {
             args.add("runtimeFilters");
             args.add(runtimeFilters.stream().map(rf -> rf.toString() + " ").collect(Collectors.toList()));
         }
-        return Utils.toSqlString("PhysicalHashJoin[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalHashJoin[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 args.toArray());
     }
 
@@ -173,7 +174,7 @@ public PhysicalHashJoin<Plan, Plan> withChildren(List<Plan> children) {
                 Optional.empty(), getLogicalProperties(), physicalProperties, statistics,
                 children.get(0), children.get(1));
         if (groupExpression.isPresent()) {
-            newJoin.setMutableState("group", groupExpression.get().getOwnerGroup().getGroupId().asInt());
+            newJoin.setMutableState(MutableState.KEY_GROUP, groupExpression.get().getOwnerGroup().getGroupId().asInt());
         }
         return newJoin;
     }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalLimit.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalLimit.java
@@ -156,7 +156,7 @@ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalLimit[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalLimit[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "limit", limit,
                 "offset", offset,
                 "phase", phase,

diff --git a/...e/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalNestedLoopJoin.java b/...e/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalNestedLoopJoin.java
@@ -27,6 +27,7 @@
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.PlanType;
 import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
+import org.apache.doris.nereids.util.MutableState;
 import org.apache.doris.nereids.util.Utils;
 import org.apache.doris.statistics.Statistics;
 
@@ -115,7 +116,7 @@ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
     @Override
     public String toString() {
         // TODO: Maybe we could pull up this to the abstract class in the future.
-        return Utils.toSqlString("PhysicalNestedLoopJoin[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalNestedLoopJoin[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "type", joinType,
                 "otherJoinCondition", otherJoinConjuncts,
                 "isMarkJoin", markJoinSlotReference.isPresent(),
@@ -131,7 +132,7 @@ public PhysicalNestedLoopJoin<Plan, Plan> withChildren(List<Plan> children) {
                 hashJoinConjuncts, otherJoinConjuncts, markJoinSlotReference, Optional.empty(),
                 getLogicalProperties(), physicalProperties, statistics, children.get(0), children.get(1));
         if (groupExpression.isPresent()) {
-            newJoin.setMutableState("group", groupExpression.get().getOwnerGroup().getGroupId().asInt());
+            newJoin.setMutableState(MutableState.KEY_GROUP, groupExpression.get().getOwnerGroup().getGroupId().asInt());
         }
         return newJoin;
     }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java
@@ -114,7 +114,7 @@ public List<Slot> getBaseOutputs() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalOlapScan[" + relationId.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalOlapScan[" + relationId.asInt() + "]" + getGroupIdWithPrefix(),
                 "qualified", Utils.qualifiedName(qualifier, table.getName()),
                 "stats", statistics, "fr", getMutableState(AbstractPlan.FRAGMENT_ID)
         );

diff --git a/...e/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOneRowRelation.java b/...e/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOneRowRelation.java
@@ -108,7 +108,7 @@ public int hashCode() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalOneRowRelation[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalOneRowRelation[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "expressions", projects
         );
     }

diff --git a/...re/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalPartitionTopN.java b/...re/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalPartitionTopN.java
@@ -178,7 +178,7 @@ public PhysicalPartitionTopN<CHILD_TYPE> withPhysicalPropertiesAndStats(Physical
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalPartitionTopN[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalPartitionTopN[" + id.asInt() + "]" + getGroupIdWithPrefix(),
             "function", function,
             "partitionKeys", partitionKeys,
             "orderKeys", orderKeys,

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalProject.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalProject.java
@@ -77,7 +77,7 @@ public List<NamedExpression> getProjects() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalProject[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalProject[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "projects", projects,
                 "stats", statistics
         );

diff --git a/...e-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalQuickSort.java b/...e-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalQuickSort.java
@@ -101,7 +101,7 @@ public PhysicalQuickSort<CHILD_TYPE> withPhysicalPropertiesAndStats(PhysicalProp
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalQuickSort[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalQuickSort[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "orderKeys", orderKeys,
                 "phase", phase.toString()
         );

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalRepeat.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalRepeat.java
@@ -97,7 +97,7 @@ public List<NamedExpression> getOutputs() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalRepeat[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalRepeat[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "groupingSets", groupingSets,
                 "outputExpressions", outputExpressions,
                 "stats", statistics

diff --git a/...ain/java/org/apache/doris/nereids/trees/plans/physical/PhysicalStorageLayerAggregate.java b/...ain/java/org/apache/doris/nereids/trees/plans/physical/PhysicalStorageLayerAggregate.java
@@ -73,7 +73,7 @@ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalStorageLayerAggregate[" + relationId.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalStorageLayerAggregate[" + relationId.asInt() + "]" + getGroupIdWithPrefix(),
                 "pushDownAggOp", aggOp,
                 "relation", relation,
                 "stats", statistics

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalTopN.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalTopN.java
@@ -138,7 +138,7 @@ public PhysicalTopN<CHILD_TYPE> withPhysicalPropertiesAndStats(PhysicalPropertie
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalTopN[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalTopN[" + id.asInt() + "]" + getGroupIdWithPrefix(),
                 "limit", limit,
                 "offset", offset,
                 "orderKeys", orderKeys,

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalUnion.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalUnion.java
@@ -84,7 +84,7 @@ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalUnion" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalUnion" + getGroupIdWithPrefix(),
                 "qualifier", qualifier,
                 "constantExprsList", constantExprsList,
                 "stats", statistics);

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalWindow.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalWindow.java
@@ -104,7 +104,7 @@ public List<? extends Expression> getExpressions() {
 
     @Override
     public String toString() {
-        return Utils.toSqlString("PhysicalWindow[" + id.asInt() + "]" + getGroupIdAsString(),
+        return Utils.toSqlString("PhysicalWindow[" + id.asInt() + "]" + getGroupIdWithPrefix(),
             "windowFrameGroup", windowFrameGroup,
             "requiredProperties", requireProperties
         );

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/util/MutableState.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/util/MutableState.java
@@ -23,6 +23,9 @@
 
 /** MutableState */
 public interface MutableState {
+    String KEY_GROUP = "group";
+    String KEY_FRAGMENT = "fragment";
+
     <T> Optional<T> get(String key);
 
     MutableState set(String key, Object value);