From d4826664b97bc5bab9f08c283751449e84e42ba8 Mon Sep 17 00:00:00 2001
From: Karol Sobczak <karol.sobczak@karolsobczak.com>
Date: Fri, 13 Sep 2019 13:54:00 +0200
Subject: [PATCH] Fix predicate pushdown for identity projections

---
 .../planner/optimizations/PredicatePushDown.java    |  6 ++++--
 .../sql/planner/TestPredicatePushdown.java          | 13 +++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/presto-main/src/main/java/io/prestosql/sql/planner/optimizations/PredicatePushDown.java b/presto-main/src/main/java/io/prestosql/sql/planner/optimizations/PredicatePushDown.java
index 2a6240702935..606a79e224ab 100644
--- a/presto-main/src/main/java/io/prestosql/sql/planner/optimizations/PredicatePushDown.java
+++ b/presto-main/src/main/java/io/prestosql/sql/planner/optimizations/PredicatePushDown.java
@@ -277,7 +277,7 @@ private boolean isInliningCandidate(Expression expression, ProjectNode node)
             verify(AstUtils.preOrder(expression).noneMatch(TryExpression.class::isInstance));
 
             // candidate symbols for inlining are
-            //   1. references to simple constants
+            //   1. references to simple constants or symbol references
             //   2. references to complex expressions that appear only once
             // which come from the node, as opposed to an enclosing scope.
             Set<Symbol> childOutputSet = ImmutableSet.copyOf(node.getOutputSymbols());
@@ -286,7 +286,9 @@ private boolean isInliningCandidate(Expression expression, ProjectNode node)
                     .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
 
             return dependencies.entrySet().stream()
-                    .allMatch(entry -> entry.getValue() == 1 || node.getAssignments().get(entry.getKey()) instanceof Literal);
+                    .allMatch(entry -> entry.getValue() == 1
+                            || node.getAssignments().get(entry.getKey()) instanceof Literal
+                            || node.getAssignments().get(entry.getKey()) instanceof SymbolReference);
         }
 
         @Override
diff --git a/presto-main/src/test/java/io/prestosql/sql/planner/TestPredicatePushdown.java b/presto-main/src/test/java/io/prestosql/sql/planner/TestPredicatePushdown.java
index 7ebd6e61fe89..3f16dcc59921 100644
--- a/presto-main/src/test/java/io/prestosql/sql/planner/TestPredicatePushdown.java
+++ b/presto-main/src/test/java/io/prestosql/sql/planner/TestPredicatePushdown.java
@@ -384,6 +384,19 @@ public void testPredicatePushDownOverProjection()
                                                 "ORDERKEY", "orderkey"))))));
     }
 
+    @Test
+    public void testPredicatePushDownOverSymbolReferences()
+    {
+        // Identities should be pushed down
+        assertPlan(
+                "WITH t AS (SELECT orderkey x, (orderkey + 1) x2 FROM orders) " +
+                        "SELECT * FROM t WHERE x > 1 OR x < 0",
+                anyTree(
+                        filter("orderkey < BIGINT '0' OR orderkey > BIGINT '1'",
+                                tableScan("orders", ImmutableMap.of(
+                                        "orderkey", "orderkey")))));
+    }
+
     @Test
     public void testConjunctsOrder()
     {