Iterative optimizer

This optimizer decouples the traversal of the plan tree (IterativeOptimizer)
from the transformation logic (Rule). The optimization loop applies rules
recursively until a fixpoint is reached.

It's implemented as PlanOptimizer so that it fits right into the existing
framework.

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -67,6 +67,7 @@ public final class SystemSessionProperties
     public static final String OPTIMIZE_DISTINCT_AGGREGATIONS = "optimize_mixed_distinct_aggregations";
     public static final String LEGACY_ORDER_BY = "legacy_order_by";
     public static final String REORDER_WINDOWS = "reorder_windows";
+    public static final String ITERATIVE_OPTIMIZER = "iterative_optimizer_enabled";
 
     private final List<PropertyMetadata<?>> sessionProperties;
 
@@ -271,6 +272,11 @@ public SystemSessionProperties(
                         REORDER_WINDOWS,
                         "Allow reordering window functions in query",
                         featuresConfig.isReorderWindows(),
+                        false),
+                booleanSessionProperty(
+                        ITERATIVE_OPTIMIZER,
+                        "Experimental: enable iterative optimizer",
+                        featuresConfig.isIterativeOptimizerEnabled(),
                         false));
     }
 
@@ -427,4 +433,9 @@ public static boolean isLegacyOrderByEnabled(Session session)
     {
         return session.getSystemProperty(LEGACY_ORDER_BY, Boolean.class);
     }
+
+    public static boolean isNewOptimizerEnabled(Session session)
+    {
+        return session.getSystemProperty(ITERATIVE_OPTIMIZER, Boolean.class);
+    }
 }

presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java

@@ -68,6 +68,7 @@ public static class ProcessingOptimization
     private DataSize operatorMemoryLimitBeforeSpill = new DataSize(4, DataSize.Unit.MEGABYTE);
     private Path spillerSpillPath = Paths.get(System.getProperty("java.io.tmpdir"), "presto", "spills");
     private int spillerThreads = 4;
+    private boolean iterativeOptimizerEnabled;
 
     public boolean isResourceGroupsEnabled()
     {
@@ -304,6 +305,18 @@ public FeaturesConfig setSpillEnabled(boolean spillEnabled)
         return this;
     }
 
+    public boolean isIterativeOptimizerEnabled()
+    {
+        return iterativeOptimizerEnabled;
+    }
+
+    @Config("experimental.iterative-optimizer-enabled")
+    public FeaturesConfig setIterativeOptimizerEnabled(boolean value)
+    {
+        this.iterativeOptimizerEnabled = value;
+        return this;
+    }
+
     public DataSize getOperatorMemoryLimitBeforeSpill()
     {
         return operatorMemoryLimitBeforeSpill;

presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java

@@ -16,6 +16,7 @@
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.sql.analyzer.FeaturesConfig;
 import com.facebook.presto.sql.parser.SqlParser;
+import com.facebook.presto.sql.planner.iterative.IterativeOptimizer;
 import com.facebook.presto.sql.planner.optimizations.AddExchanges;
 import com.facebook.presto.sql.planner.optimizations.AddLocalExchanges;
 import com.facebook.presto.sql.planner.optimizations.BeginTableWrite;
@@ -56,6 +57,7 @@
 import com.facebook.presto.sql.planner.optimizations.UnaliasSymbolReferences;
 import com.facebook.presto.sql.planner.optimizations.WindowFilterPushDown;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.inject.Inject;
 
 import java.util.List;
@@ -77,6 +79,7 @@ public PlanOptimizers(Metadata metadata, SqlParser sqlParser, FeaturesConfig fea
         builder.add(
                 new DesugaringOptimizer(metadata, sqlParser), // Clean up all the sugar in expressions, e.g. AtTimeZone, must be run before all the other optimizers
                 new CanonicalizeExpressions(),
+                new IterativeOptimizer(ImmutableSet.of()),
                 new ImplementFilteredAggregations(),
                 new ImplementSampleAsFilter(),
                 new SimplifyExpressions(metadata, sqlParser),

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

@@ -0,0 +1,61 @@
+/*
+ * 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.sql.planner.iterative;
+
+import com.facebook.presto.sql.planner.Symbol;
+import com.facebook.presto.sql.planner.plan.PlanNode;
+import com.facebook.presto.sql.planner.plan.PlanNodeId;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+public class GroupReference
+        extends PlanNode
+{
+    private final int groupId;
+    private final List<Symbol> outputs;
+
+    public GroupReference(PlanNodeId id, int groupId, List<Symbol> outputs)
+    {
+        super(id);
+        this.groupId = groupId;
+        this.outputs = ImmutableList.copyOf(outputs);
+    }
+
+    public int getGroupId()
+    {
+        return groupId;
+    }
+
+    @Override
+    public List<PlanNode> getSources()
+    {
+        return ImmutableList.of();
+    }
+
+    @Override
+    public List<Symbol> getOutputSymbols()
+    {
+        return outputs;
+    }
+
+    @Override
+    public PlanNode replaceChildren(List<PlanNode> newChildren)
+    {
+        checkArgument(newChildren.isEmpty(), "newChildren is not empty");
+        return this;
+    }
+}

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

@@ -0,0 +1,160 @@
+/*
+ * 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.sql.planner.iterative;
+
+import com.facebook.presto.Session;
+import com.facebook.presto.SystemSessionProperties;
+import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.sql.planner.PlanNodeIdAllocator;
+import com.facebook.presto.sql.planner.Symbol;
+import com.facebook.presto.sql.planner.SymbolAllocator;
+import com.facebook.presto.sql.planner.optimizations.PlanOptimizer;
+import com.facebook.presto.sql.planner.plan.PlanNode;
+import com.google.common.collect.ImmutableSet;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+
+import static com.google.common.base.Preconditions.checkState;
+
+public class IterativeOptimizer
+        implements PlanOptimizer
+{
+    private final Set<Rule> rules;
+
+    public IterativeOptimizer(Set<Rule> rules)
+    {
+        this.rules = ImmutableSet.copyOf(rules);
+    }
+
+    @Override
+    public PlanNode optimize(PlanNode plan, Session session, Map<Symbol, Type> types, SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator)
+    {
+        if (!SystemSessionProperties.isNewOptimizerEnabled(session)) {
+            return plan;
+        }
+
+        Memo memo = new Memo(idAllocator, plan);
+
+        Lookup lookup = node -> {
+            if (node instanceof GroupReference) {
+                return memo.getNode(((GroupReference) node).getGroupId());
+            }
+
+            return node;
+        };
+
+        exploreGroup(memo.getRootGroup(), new Context(memo, lookup, idAllocator, symbolAllocator));
+
+        return memo.extract();
+    }
+
+    private boolean exploreGroup(int group, Context context)
+    {
+        // tracks whether this group or any children groups change as
+        // this method executes
+        boolean progress = exploreNode(group, context);
+
+        while (exploreChildren(group, context)) {
+            progress = true;
+
+            // if children changed, try current group again
+            // in case we can match additional rules
+            if (!exploreNode(group, context)) {
+                // no additional matches, so bail out
+                break;
+            }
+        }
+
+        return progress;
+    }
+
+    private boolean exploreNode(int group, Context context)
+    {
+        PlanNode node = context.getMemo().getNode(group);
+
+        boolean done = false;
+        boolean progress = false;
+
+        while (!done) {
+            done = true;
+            for (Rule rule : rules) {
+                Optional<PlanNode> transformed = rule.apply(node, context.getLookup(), context.getIdAllocator(), context.getSymbolAllocator());
+
+                if (transformed.isPresent()) {
+                    context.getMemo().replace(group, transformed.get(), rule.getClass().getName());
+                    node = transformed.get();
+
+                    done = false;
+                    progress = true;
+                }
+            }
+        }
+
+        return progress;
+    }
+
+    private boolean exploreChildren(int group, Context context)
+    {
+        boolean progress = false;
+
+        PlanNode expression = context.getMemo().getNode(group);
+        for (PlanNode child : expression.getSources()) {
+            checkState(child instanceof GroupReference, "Expected child to be a group reference. Found: " + child.getClass().getName());
+
+            if (exploreGroup(((GroupReference) child).getGroupId(), context)) {
+                progress = true;
+            }
+        }
+
+        return progress;
+    }
+
+    private static class Context
+    {
+        private final Memo memo;
+        private final Lookup lookup;
+        private final PlanNodeIdAllocator idAllocator;
+        private final SymbolAllocator symbolAllocator;
+
+        public Context(Memo memo, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator)
+        {
+            this.memo = memo;
+            this.lookup = lookup;
+            this.idAllocator = idAllocator;
+            this.symbolAllocator = symbolAllocator;
+        }
+
+        public Memo getMemo()
+        {
+            return memo;
+        }
+
+        public Lookup getLookup()
+        {
+            return lookup;
+        }
+
+        public PlanNodeIdAllocator getIdAllocator()
+        {
+            return idAllocator;
+        }
+
+        public SymbolAllocator getSymbolAllocator()
+        {
+            return symbolAllocator;
+        }
+    }
+}

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

@@ -0,0 +1,28 @@
+/*
+ * 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.sql.planner.iterative;
+
+import com.facebook.presto.sql.planner.plan.PlanNode;
+
+public interface Lookup
+{
+    /**
+     * Resolves a node by materializing GroupReference nodes
+     * representing symbolic references to other nodes.
+     *
+     * If the node is not a GroupReference, it returns the
+     * argument as is.
+     */
+    PlanNode resolve(PlanNode node);
+}