Include root node LocalCostEstimate in PlanCostEstimate

presto-main/src/main/java/io/prestosql/cost/CostCalculatorUsingExchanges.java

@@ -305,7 +305,7 @@ public PlanCostEstimate visitUnion(UnionNode node, Void context)
         private PlanCostEstimate costForSource(PlanNode node, LocalCostEstimate localCost)
         {
             verify(node.getSources().isEmpty(), "Unexpected sources for %s: %s", node, node.getSources());
-            return new PlanCostEstimate(localCost.getCpuCost(), localCost.getMaxMemory(), localCost.getMaxMemory(), localCost.getNetworkCost());
+            return new PlanCostEstimate(localCost.getCpuCost(), localCost.getMaxMemory(), localCost.getMaxMemory(), localCost.getNetworkCost(), localCost);
         }
 
         private PlanCostEstimate costForAccumulation(PlanNode node, LocalCostEstimate localCost)
@@ -318,7 +318,8 @@ private PlanCostEstimate costForAccumulation(PlanNode node, LocalCostEstimate lo
                             sourcesCost.getMaxMemory(), // Accumulating operator allocates insignificant amount of memory (usually none) before first input page is received
                             sourcesCost.getMaxMemoryWhenOutputting() + localCost.getMaxMemory()),
                     localCost.getMaxMemory(), // Source freed its memory allocations when finished its output
-                    sourcesCost.getNetworkCost() + localCost.getNetworkCost());
+                    sourcesCost.getNetworkCost() + localCost.getNetworkCost(),
+                    localCost);
         }
 
         private PlanCostEstimate costForStreaming(PlanNode node, LocalCostEstimate localCost)
@@ -331,7 +332,8 @@ private PlanCostEstimate costForStreaming(PlanNode node, LocalCostEstimate local
                             sourcesCost.getMaxMemory(), // Streaming operator allocates insignificant amount of memory (usually none) before first input page is received
                             sourcesCost.getMaxMemoryWhenOutputting() + localCost.getMaxMemory()),
                     sourcesCost.getMaxMemoryWhenOutputting() + localCost.getMaxMemory(),
-                    sourcesCost.getNetworkCost() + localCost.getNetworkCost());
+                    sourcesCost.getNetworkCost() + localCost.getNetworkCost(),
+                    localCost);
         }
 
         private PlanCostEstimate costForLookupJoin(PlanNode node, LocalCostEstimate localCost)
@@ -348,7 +350,8 @@ private PlanCostEstimate costForLookupJoin(PlanNode node, LocalCostEstimate loca
                             probeCost.getMaxMemory() + buildCost.getMaxMemory(), // Probe and build execute independently, so their max memory allocations can be realized at the same time
                             probeCost.getMaxMemory() + buildCost.getMaxMemoryWhenOutputting() + localCost.getMaxMemory()),
                     probeCost.getMaxMemoryWhenOutputting() + localCost.getMaxMemory(), // Build side finished and freed its memory allocations
-                    probeCost.getNetworkCost() + buildCost.getNetworkCost() + localCost.getNetworkCost());
+                    probeCost.getNetworkCost() + buildCost.getNetworkCost() + localCost.getNetworkCost(),
+                    localCost);
         }
 
         private PlanNodeStatsEstimate getStats(PlanNode node)

presto-main/src/main/java/io/prestosql/cost/CostCalculatorWithEstimatedExchanges.java

@@ -80,7 +80,8 @@ private static PlanCostEstimate addExchangeCost(PlanCostEstimate costEstimate, L
                 // exchange memory allocation will actually be freed before node is outputting. Conservatively we assume the exchanges can still
                 // hold the memory when the node is outputting.
                 costEstimate.getMaxMemoryWhenOutputting() + estimatedExchangeCost.getMaxMemory(),
-                costEstimate.getNetworkCost() + estimatedExchangeCost.getNetworkCost());
+                costEstimate.getNetworkCost() + estimatedExchangeCost.getNetworkCost(),
+                addPartialComponents(costEstimate.getRootNodeLocalCostEstimate(), estimatedExchangeCost));
     }
 
     private static class ExchangeCostEstimator

presto-main/src/main/java/io/prestosql/cost/LocalCostEstimate.java

@@ -13,8 +13,11 @@
  */
 package io.prestosql.cost;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import io.prestosql.sql.planner.plan.PlanNode;
 
+import java.util.Objects;
 import java.util.stream.Stream;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
@@ -54,23 +57,30 @@ public static LocalCostEstimate of(double cpuCost, double maxMemory, double netw
         return new LocalCostEstimate(cpuCost, maxMemory, networkCost);
     }
 
-    private LocalCostEstimate(double cpuCost, double maxMemory, double networkCost)
+    @JsonCreator
+    public LocalCostEstimate(
+            @JsonProperty("cpuCost") double cpuCost,
+            @JsonProperty("maxMemory") double maxMemory,
+            @JsonProperty("networkCost") double networkCost)
     {
         this.cpuCost = cpuCost;
         this.maxMemory = maxMemory;
         this.networkCost = networkCost;
     }
 
+    @JsonProperty
     public double getCpuCost()
     {
         return cpuCost;
     }
 
+    @JsonProperty
     public double getMaxMemory()
     {
         return maxMemory;
     }
 
+    @JsonProperty
     public double getNetworkCost()
     {
         return networkCost;
@@ -82,7 +92,7 @@ public double getNetworkCost()
     @Deprecated
     public PlanCostEstimate toPlanCost()
     {
-        return new PlanCostEstimate(cpuCost, maxMemory, maxMemory, networkCost);
+        return new PlanCostEstimate(cpuCost, maxMemory, maxMemory, networkCost, this);
     }
 
     @Override
@@ -95,6 +105,27 @@ public String toString()
                 .toString();
     }
 
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        LocalCostEstimate that = (LocalCostEstimate) o;
+        return Double.compare(that.cpuCost, cpuCost) == 0 &&
+                Double.compare(that.maxMemory, maxMemory) == 0 &&
+                Double.compare(that.networkCost, networkCost) == 0;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(cpuCost, maxMemory, networkCost);
+    }
+
     /**
      * Sums partial cost estimates of some (single) plan node.
      */

presto-main/src/main/java/io/prestosql/cost/PlanCostEstimate.java

@@ -23,6 +23,7 @@
 import static java.lang.Double.NaN;
 import static java.lang.Double.POSITIVE_INFINITY;
 import static java.lang.Double.isNaN;
+import static java.util.Objects.requireNonNull;
 
 public final class PlanCostEstimate
 {
@@ -34,6 +35,7 @@ public final class PlanCostEstimate
     private final double maxMemory;
     private final double maxMemoryWhenOutputting;
     private final double networkCost;
+    private final LocalCostEstimate rootNodeLocalCostEstimate;
 
     public static PlanCostEstimate infinite()
     {
@@ -50,12 +52,22 @@ public static PlanCostEstimate zero()
         return ZERO;
     }
 
+    public PlanCostEstimate(
+            double cpuCost,
+            double maxMemory,
+            double maxMemoryWhenOutputting,
+            double networkCost)
+    {
+        this(cpuCost, maxMemory, maxMemoryWhenOutputting, networkCost, LocalCostEstimate.unknown());
+    }
+
     @JsonCreator
     public PlanCostEstimate(
             @JsonProperty("cpuCost") double cpuCost,
             @JsonProperty("maxMemory") double maxMemory,
             @JsonProperty("maxMemoryWhenOutputting") double maxMemoryWhenOutputting,
-            @JsonProperty("networkCost") double networkCost)
+            @JsonProperty("networkCost") double networkCost,
+            @JsonProperty("rootNodeLocalCostEstimate") LocalCostEstimate rootNodeLocalCostEstimate)
     {
         checkArgument(!(cpuCost < 0), "cpuCost cannot be negative: %s", cpuCost);
         checkArgument(!(maxMemory < 0), "maxMemory cannot be negative: %s", maxMemory);
@@ -66,6 +78,7 @@ public PlanCostEstimate(
         this.maxMemory = maxMemory;
         this.maxMemoryWhenOutputting = maxMemoryWhenOutputting;
         this.networkCost = networkCost;
+        this.rootNodeLocalCostEstimate = requireNonNull(rootNodeLocalCostEstimate, "rootNodeLocalCostEstimate is null");
     }
 
     /**
@@ -114,6 +127,15 @@ public double getNetworkCost()
         return networkCost;
     }
 
+    /**
+     * Returns local cost estimate for the root plan node.
+     */
+    @JsonProperty
+    public LocalCostEstimate getRootNodeLocalCostEstimate()
+    {
+        return rootNodeLocalCostEstimate;
+    }
+
     /**
      * Returns true if this cost has unknown components.
      */
@@ -130,6 +152,7 @@ public String toString()
                 .add("memory", maxMemory)
                 // maxMemoryWhenOutputting is not that useful in toString
                 .add("network", networkCost)
+                .add("rootNodeLocalCostEstimate", rootNodeLocalCostEstimate)
                 .toString();
     }
 
@@ -146,12 +169,13 @@ public boolean equals(Object o)
         return Double.compare(that.cpuCost, cpuCost) == 0 &&
                 Double.compare(that.maxMemory, maxMemory) == 0 &&
                 Double.compare(that.maxMemoryWhenOutputting, maxMemoryWhenOutputting) == 0 &&
-                Double.compare(that.networkCost, networkCost) == 0;
+                Double.compare(that.networkCost, networkCost) == 0 &&
+                Objects.equals(rootNodeLocalCostEstimate, that.rootNodeLocalCostEstimate);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(cpuCost, maxMemory, maxMemoryWhenOutputting, networkCost);
+        return Objects.hash(cpuCost, maxMemory, maxMemoryWhenOutputting, networkCost, rootNodeLocalCostEstimate);
     }
 }