Repository: prestodb/presto
Issue: row_number() with limit filter incorrect when filter contains NOT

The following query:

``` sql
SELECT * FROM (
    SELECT row_number() OVER () rn 
    FROM (VALUES (0),(0),(0),(0))
) 
WHERE NOT rn < 3;
```

produces the following (incorrect) output:

```
 rn 
----
(0 rows)
```

The correct output should be:

```
 rn 
----
  3 
  4 
(2 rows)
```


Changed files:
presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java
index 85601e9b50..86a9c5e7a3 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java
@@ -111,7 +111,8 @@ public class WindowFilterPushDown
             if (filter.getLimit().isPresent()) {
                 return filter.getLimit();
             }
-            if (filterContainsWindowFunctions(node, filter.getFilterExpression().get())) {
+            if (filterContainsWindowFunctions(node, filter.getFilterExpression().get()) &&
+                    filter.getFilterExpression().get() instanceof ComparisonExpression) {
                 Symbol rowNumberSymbol = Iterables.getOnlyElement(node.getWindowFunctions().entrySet()).getKey();
                 return WindowLimitExtractor.extract(filter.getFilterExpression().get(), rowNumberSymbol);
             }
@@ -213,52 +214,56 @@ public class WindowFilterPushDown
             @Override
             protected Long visitComparisonExpression(ComparisonExpression node, Symbol rowNumberSymbol)
             {
-                QualifiedNameReference reference = extractReference(node);
-                Literal literal = extractLiteral(node);
-                if (!Symbol.fromQualifiedName(reference.getName()).equals(rowNumberSymbol)) {
+                Optional<QualifiedNameReference> reference = extractReference(node);
+                Optional<Literal> literal = extractLiteral(node);
+                if (!reference.isPresent() || !literal.isPresent()) {
+                    return null;
+                }
+                if (!Symbol.fromQualifiedName(reference.get().getName()).equals(rowNumberSymbol)) {
                     return null;
                 }
 
+                long literalValue = extractValue(literal.get());
                 if (node.getLeft() instanceof QualifiedNameReference && node.getRight() instanceof Literal) {
                     if (node.getType() == ComparisonExpression.Type.LESS_THAN_OR_EQUAL) {
-                        return extractValue(literal);
+                        return literalValue;
                     }
                     if (node.getType() == ComparisonExpression.Type.LESS_THAN) {
-                        return extractValue(literal) - 1;
+                        return literalValue - 1;
                     }
                 }
                 else if (node.getLeft() instanceof Literal && node.getRight() instanceof QualifiedNameReference) {
                     if (node.getType() == ComparisonExpression.Type.GREATER_THAN_OR_EQUAL) {
-                        return extractValue(literal);
+                        return literalValue;
                     }
                     if (node.getType() == ComparisonExpression.Type.GREATER_THAN) {
-                        return extractValue(literal) - 1;
+                        return literalValue - 1;
                     }
                 }
                 return null;
             }
         }
 
-        private static QualifiedNameReference extractReference(ComparisonExpression expression)
+        private static Optional<QualifiedNameReference> extractReference(ComparisonExpression expression)
         {
             if (expression.getLeft() instanceof QualifiedNameReference) {
-                return (QualifiedNameReference) expression.getLeft();
+                return Optional.of((QualifiedNameReference) expression.getLeft());
             }
             if (expression.getRight() instanceof QualifiedNameReference) {
-                return (QualifiedNameReference) expression.getRight();
+                return Optional.of((QualifiedNameReference) expression.getRight());
             }
-            throw new IllegalArgumentException("Comparison does not have a child of type QualifiedNameReference");
+            return Optional.absent();
         }
 
-        private static Literal extractLiteral(ComparisonExpression expression)
+        private static Optional<Literal> extractLiteral(ComparisonExpression expression)
         {
             if (expression.getLeft() instanceof Literal) {
-                return (Literal) expression.getLeft();
+                return Optional.of((Literal) expression.getLeft());
             }
             if (expression.getRight() instanceof Literal) {
-                return (Literal) expression.getRight();
+                return Optional.of((Literal) expression.getRight());
             }
-            throw new IllegalArgumentException("Comparison does not have a child of type Literal");
+            return Optional.absent();
         }
 
         private static long extractValue(Literal literal)
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
index 1e130e2abd..16c35dc736 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
@@ -1873,6 +1873,29 @@ public abstract class AbstractTestQueries
         assertEquals(actual, expected);
     }
 
+    @Test
+    public void testRowNumberNoOptimization()
+            throws Exception
+    {
+        MaterializedResult actual = computeActual("" +
+                "SELECT orderkey, orderstatus FROM (\\n" +
+                "   SELECT row_number() OVER () rn, orderkey, orderstatus\\n" +
+                "   FROM orders\\n" +
+                ") WHERE NOT rn <= 10\\n");
+        MaterializedResult all = computeExpected("SELECT orderkey, orderstatus FROM ORDERS", actual.getTypes());
+        assertEquals(actual.getMaterializedRows().size(), all.getMaterializedRows().size() - 10);
+        assertTrue(all.getMaterializedRows().containsAll(actual.getMaterializedRows()));
+
+        actual = computeActual("" +
+                "SELECT orderkey, orderstatus FROM (\\n" +
+                "   SELECT row_number() OVER () rn, orderkey, orderstatus\\n" +
+                "   FROM orders\\n" +
+                ") WHERE rn - 5 <= 10\\n");
+        all = computeExpected("SELECT orderkey, orderstatus FROM ORDERS", actual.getTypes());
+        assertEquals(actual.getMaterializedRows().size(), 15);
+        assertTrue(all.getMaterializedRows().containsAll(actual.getMaterializedRows()));
+    }
+
     @Test
     public void testRowNumberUnpartitionedLimit()
             throws Exception