Fix predicate pushdown for identity projections

trinodb · Sep 25, 2019 · 7195e05 · 7195e05
1 parent 91f7679
commit 7195e05
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 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
@@ -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
@@ -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()
     {