From 648528b271157ec0f1d6b0489de6bf82bc6eeb3e Mon Sep 17 00:00:00 2001
From: Lukasz Osipiuk <lukasz@osipiuk.net>
Date: Tue, 31 Jan 2017 18:39:20 +0100
Subject: [PATCH] Use PushProjectionThrough{Union,Exchange} optimizers

---
 .../presto/sql/planner/PlanOptimizers.java        | 15 ++++++++++++---
 .../planner/optimizations/ProjectionPushDown.java |  1 +
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java
index 397d66b597f6a..8d6d2bdff019a 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java
@@ -37,6 +37,8 @@
 import com.facebook.presto.sql.planner.iterative.rule.PushLimitThroughMarkDistinct;
 import com.facebook.presto.sql.planner.iterative.rule.PushLimitThroughProject;
 import com.facebook.presto.sql.planner.iterative.rule.PushLimitThroughSemiJoin;
+import com.facebook.presto.sql.planner.iterative.rule.PushProjectionThroughExchange;
+import com.facebook.presto.sql.planner.iterative.rule.PushProjectionThroughUnion;
 import com.facebook.presto.sql.planner.iterative.rule.PushTopNThroughUnion;
 import com.facebook.presto.sql.planner.iterative.rule.RemoveEmptyDelete;
 import com.facebook.presto.sql.planner.iterative.rule.RemoveFullSample;
@@ -126,6 +128,13 @@ public PlanOptimizers(Metadata metadata, SqlParser sqlParser, FeaturesConfig fea
                         new InlineProjections(),
                         new RemoveRedundantIdentityProjections()));
 
+        IterativeOptimizer projectionPushDown = new IterativeOptimizer(
+                stats,
+                ImmutableList.of(new ProjectionPushDown()),
+                ImmutableSet.of(
+                        new PushProjectionThroughUnion(),
+                        new PushProjectionThroughExchange()));
+
         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(),
@@ -182,7 +191,7 @@ public PlanOptimizers(Metadata metadata, SqlParser sqlParser, FeaturesConfig fea
                 ),
                 inlineProjections,
                 new SimplifyExpressions(metadata, sqlParser), // Re-run the SimplifyExpressions to simplify any recomposed expressions from other optimizations
-                new ProjectionPushDown(),
+                projectionPushDown,
                 new UnaliasSymbolReferences(), // Run again because predicate pushdown and projection pushdown might add more projections
                 new PruneUnreferencedOutputs(), // Make sure to run this before index join. Filtered projections may not have all the columns.
                 new IndexJoinOptimizer(metadata), // Run this after projections and filters have been fully simplified and pushed down
@@ -206,7 +215,7 @@ public PlanOptimizers(Metadata metadata, SqlParser sqlParser, FeaturesConfig fea
                 new MetadataQueryOptimizer(metadata),
                 new EliminateCrossJoins(), // This can pull up Filter and Project nodes from between Joins, so we need to push them down again
                 new PredicatePushDown(metadata, sqlParser),
-                new ProjectionPushDown());
+                projectionPushDown);
 
         if (featuresConfig.isOptimizeSingleDistinct()) {
             builder.add(
@@ -238,7 +247,7 @@ public PlanOptimizers(Metadata metadata, SqlParser sqlParser, FeaturesConfig fea
                 ));
 
         builder.add(new PredicatePushDown(metadata, sqlParser)); // Run predicate push down one more time in case we can leverage new information from layouts' effective predicate
-        builder.add(new ProjectionPushDown());
+        builder.add(projectionPushDown);
         builder.add(inlineProjections);
         builder.add(new UnaliasSymbolReferences()); // Run unalias after merging projections to simplify projections more efficiently
         builder.add(new PruneUnreferencedOutputs());
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/ProjectionPushDown.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/ProjectionPushDown.java
index d5d373352d00d..9f3459598bc3e 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/ProjectionPushDown.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/ProjectionPushDown.java
@@ -38,6 +38,7 @@
 import static com.facebook.presto.sql.planner.plan.ChildReplacer.replaceChildren;
 import static java.util.Objects.requireNonNull;
 
+@Deprecated
 public class ProjectionPushDown
         implements PlanOptimizer
 {