Join reordering stats11 speedup

presto-main/src/main/java/com/facebook/presto/cost/CachingCostCalculator.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.presto.cost;
+
+import com.facebook.presto.Session;
+import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.sql.planner.Symbol;
+import com.facebook.presto.sql.planner.iterative.Lookup;
+import com.facebook.presto.sql.planner.plan.PlanNode;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static com.google.common.base.Preconditions.checkState;
+import static java.util.Objects.requireNonNull;
+
+public class CachingCostCalculator
+        implements CostCalculator
+{
+    private final CostCalculator costCalculator;
+    private final Map<PlanNode, PlanNodeCostEstimate> costs = new HashMap<>();
+
+    public CachingCostCalculator(CostCalculator costCalculator)
+    {
+        this.costCalculator = requireNonNull(costCalculator, "costCalculator is null");
+    }
+
+    @Override
+    public PlanNodeCostEstimate calculateCost(PlanNode planNode, Lookup lookup, Session session, Map<Symbol, Type> types)
+    {
+        if (!costs.containsKey(planNode)) {
+            // cannot use Map.computeIfAbsent due to costs map modification in the mappingFunction callback
+            PlanNodeCostEstimate cost = costCalculator.calculateCumulativeCost(planNode, lookup, session, types);
+            requireNonNull(costs, "computed cost can not be null");
+            checkState(costs.put(planNode, cost) == null, "cost for " + planNode + " already computed");
+        }
+        return costs.get(planNode);
+    }
+}

presto-main/src/main/java/com/facebook/presto/cost/CachingStatsCalculator.java

@@ -0,0 +1,50 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.cost;
+
+import com.facebook.presto.Session;
+import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.sql.planner.Symbol;
+import com.facebook.presto.sql.planner.iterative.Lookup;
+import com.facebook.presto.sql.planner.plan.PlanNode;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static com.google.common.base.Preconditions.checkState;
+import static java.util.Objects.requireNonNull;
+
+public class CachingStatsCalculator
+        implements StatsCalculator
+{
+    private final StatsCalculator statsCalculator;
+    private final Map<PlanNode, PlanNodeStatsEstimate> stats = new HashMap<>();
+
+    public CachingStatsCalculator(StatsCalculator statsCalculator)
+    {
+        this.statsCalculator = statsCalculator;
+    }
+
+    @Override
+    public PlanNodeStatsEstimate calculateStats(PlanNode planNode, Lookup lookup, Session session, Map<Symbol, Type> types)
+    {
+        if (!stats.containsKey(planNode)) {
+            // cannot use Map.computeIfAbsent due to stats map modification in the mappingFunction callback
+            PlanNodeStatsEstimate statsEstimate = statsCalculator.calculateStats(planNode, lookup, session, types);
+            requireNonNull(stats, "computed stats can not be null");
+            checkState(stats.put(planNode, statsEstimate) == null, "statistics for " + planNode + " already computed");
+        }
+        return stats.get(planNode);
+    }
+}

presto-main/src/main/java/com/facebook/presto/sql/planner/iterative/IterativeOptimizer.java

@@ -15,6 +15,8 @@
 
 import com.facebook.presto.Session;
 import com.facebook.presto.SystemSessionProperties;
+import com.facebook.presto.cost.CachingCostCalculator;
+import com.facebook.presto.cost.CachingStatsCalculator;
 import com.facebook.presto.cost.CostCalculator;
 import com.facebook.presto.cost.StatsCalculator;
 import com.facebook.presto.matching.MatchingEngine;
@@ -81,7 +83,7 @@ public PlanNode optimize(PlanNode plan, Session session, Map<Symbol, Type> types
         }
 
         Memo memo = new Memo(idAllocator, plan);
-        Lookup lookup = new MemoBasedLookup(memo, statsCalculator, costCalculator);
+        Lookup lookup = Lookup.from(memo::resolve, new CachingStatsCalculator(statsCalculator), new CachingCostCalculator(costCalculator));
 
         Duration timeout = SystemSessionProperties.getOptimizerTimeout(session);
         exploreGroup(memo.getRootGroup(), new Context(memo, lookup, idAllocator, symbolAllocator, System.nanoTime(), timeout.toMillis(), session));

presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/TestMemoBasedLookup.java

@@ -24,7 +24,6 @@
 import com.facebook.presto.sql.planner.iterative.GroupReference;
 import com.facebook.presto.sql.planner.iterative.Lookup;
 import com.facebook.presto.sql.planner.iterative.Memo;
-import com.facebook.presto.sql.planner.iterative.MemoBasedLookup;
 import com.facebook.presto.sql.planner.plan.PlanNode;
 import com.facebook.presto.sql.planner.plan.PlanNodeId;
 import com.facebook.presto.testing.LocalQueryRunner;
@@ -60,7 +59,7 @@ public void testResolvesGroupReferenceNode()
         PlanNode plan = node(source);
         Memo memo = new Memo(idAllocator, plan);
 
-        MemoBasedLookup lookup = new MemoBasedLookup(memo, new NodeCountingStatsCalculator(), new CostCalculatorUsingExchanges(1));
+        Lookup lookup = Lookup.from(memo::resolve, new NodeCountingStatsCalculator(), new CostCalculatorUsingExchanges(1));
         PlanNode memoSource = Iterables.getOnlyElement(memo.getNode(memo.getRootGroup()).getSources());
         checkState(memoSource instanceof GroupReference, "expected GroupReference");
         assertEquals(lookup.resolve(memoSource), source);
@@ -71,7 +70,7 @@ public void testComputesStatsAndResolvesNodes()
     {
         PlanNode plan = node(node(node()));
         Memo memo = new Memo(idAllocator, plan);
-        MemoBasedLookup lookup = new MemoBasedLookup(memo, new NodeCountingStatsCalculator(), new CostCalculatorUsingExchanges(1));
+        Lookup lookup = Lookup.from(memo::resolve, new NodeCountingStatsCalculator(), new CostCalculatorUsingExchanges(1));
 
         PlanNodeStatsEstimate actualStats = lookup.getStats(memo.getNode(memo.getRootGroup()), queryRunner.getDefaultSession(), ImmutableMap.of());
         PlanNodeStatsEstimate expectedStats = PlanNodeStatsEstimate.builder().setOutputRowCount(3).build();