[Performance](Nereids): pass ConnectContext to avoid ThreadLocal.get() (

#26165)
apache · Nov 1, 2023 · 6838322 · 6838322
1 parent 1770224
commit 6838322
Show file tree

Hide file tree

Showing 24 changed files with 274 additions and 233 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/CascadesContext.java
@@ -199,7 +199,7 @@ public synchronized boolean isTimeout() {
     }
 
     public void toMemo() {
-        this.memo = new Memo(plan);
+        this.memo = new Memo(getConnectContext(), plan);
     }
 
     public Analyzer newAnalyzer() {
@@ -358,7 +358,7 @@ public TableIf getTableInMinidumpCache(String tableName) {
                 return table;
             }
         }
-        if (ConnectContext.get().getSessionVariable().isPlayNereidsDump()) {
+        if (getConnectContext().getSessionVariable().isPlayNereidsDump()) {
             throw new AnalysisException("Minidump cache can not find table:" + tableName);
         }
         return null;

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/PlanContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/PlanContext.java
@@ -19,6 +19,8 @@
 
 import org.apache.doris.nereids.memo.GroupExpression;
 import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.SessionVariable;
 import org.apache.doris.statistics.Statistics;
 
 import java.util.ArrayList;
@@ -31,7 +33,7 @@
  * Inspired by GPORCA-CExpressionHandle.
  */
 public class PlanContext {
-
+    private final ConnectContext connectContext;
     private final List<Statistics> childrenStats;
     private final Statistics planStats;
     private final int arity;
@@ -41,7 +43,8 @@ public class PlanContext {
     /**
      * Constructor for PlanContext.
      */
-    public PlanContext(GroupExpression groupExpression) {
+    public PlanContext(ConnectContext connectContext, GroupExpression groupExpression) {
+        this.connectContext = connectContext;
         this.arity = groupExpression.arity();
         this.planStats = groupExpression.getOwnerGroup().getStatistics();
         this.isStatsReliable = groupExpression.getOwnerGroup().isStatsReliable();
@@ -51,12 +54,8 @@ public PlanContext(GroupExpression groupExpression) {
         }
     }
 
-    // This is used in GraphSimplifier
-    public PlanContext(Statistics planStats, List<Statistics> childrenStats) {
-        this.planStats = planStats;
-        this.childrenStats = childrenStats;
-        this.isStatsReliable = false;
-        this.arity = this.childrenStats.size();
+    public SessionVariable getSessionVariable() {
+        return connectContext.getSessionVariable();
     }
 
     public void setBroadcastJoin() {

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/Cost.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/Cost.java
@@ -17,7 +17,7 @@
 
 package org.apache.doris.nereids.cost;
 
-import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.SessionVariable;
 
 /**
  * Cost encapsulate the real cost with double type.
@@ -27,21 +27,22 @@
 public interface Cost {
     double getValue();
 
-    /**
-     * return zero cost
-     */
-    static Cost zero() {
-        if (ConnectContext.get().getSessionVariable().getEnableNewCostModel()) {
+    static Cost zero(SessionVariable sessionVariable) {
+        if (sessionVariable.getEnableNewCostModel()) {
             return CostV2.zero();
         }
         return CostV1.zero();
     }
 
-    static Cost infinite() {
-        if (ConnectContext.get().getSessionVariable().getEnableNewCostModel()) {
+    static Cost infinite(SessionVariable sessionVariable) {
+        if (sessionVariable.getEnableNewCostModel()) {
             return CostV2.infinite();
         }
         return CostV1.infinite();
     }
+
+    static Cost zeroV1() {
+        return CostV1.zero();
+    }
 }
 
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/CostCalculator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/cost/CostCalculator.java
@@ -18,53 +18,42 @@
 package org.apache.doris.nereids.cost;
 
 import org.apache.doris.nereids.PlanContext;
-import org.apache.doris.nereids.annotation.Developing;
 import org.apache.doris.nereids.memo.GroupExpression;
 import org.apache.doris.nereids.properties.DistributionSpecReplicated;
 import org.apache.doris.nereids.properties.PhysicalProperties;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.SessionVariable;
 
 import java.util.List;
 
 /**
  * Calculate the cost of a plan.
  */
-@Developing
-//TODO: memory cost and network cost should be estimated by byte size.
+// TODO: memory cost and network cost should be estimated by byte size.
 public class CostCalculator {
 
     /**
      * Calculate cost for groupExpression
      */
-    public static Cost calculateCost(GroupExpression groupExpression, List<PhysicalProperties> childrenProperties) {
-        PlanContext planContext = new PlanContext(groupExpression);
+    public static Cost calculateCost(ConnectContext connectContext, GroupExpression groupExpression,
+            List<PhysicalProperties> childrenProperties) {
+        PlanContext planContext = new PlanContext(connectContext, groupExpression);
         if (childrenProperties.size() >= 2
                 && childrenProperties.get(1).getDistributionSpec() instanceof DistributionSpecReplicated) {
             planContext.setBroadcastJoin();
         }
 
-        CostModelV1 costModelV1 = new CostModelV1();
+        CostModelV1 costModelV1 = new CostModelV1(connectContext);
         return groupExpression.getPlan().accept(costModelV1, planContext);
     }
 
-    /**
-     * Calculate cost without groupExpression
-     */
-    public static Cost calculateCost(Plan plan, PlanContext planContext) {
-        if (ConnectContext.get().getSessionVariable().getEnableNewCostModel()) {
-            CostModelV2 costModel = new CostModelV2();
-            return plan.accept(costModel, planContext);
-        } else {
-            CostModelV1 costModel = new CostModelV1();
-            return plan.accept(costModel, planContext);
-        }
-    }
-
-    public static Cost addChildCost(Plan plan, Cost planCost, Cost childCost, int index) {
-        if (ConnectContext.get().getSessionVariable().getEnableNewCostModel()) {
+    public static Cost addChildCost(ConnectContext connectContext, Plan plan, Cost planCost, Cost childCost,
+            int index) {
+        SessionVariable sessionVariable = connectContext.getSessionVariable();
+        if (sessionVariable.getEnableNewCostModel()) {
             return CostModelV2.addChildCost(plan, planCost, childCost, index);
         }
-        return CostModelV1.addChildCost(plan, planCost, childCost, index);
+        return CostModelV1.addChildCost(sessionVariable, planCost, childCost);
     }
 }