Introduce new Keyword FETCH FIRST N ROWS ONLY in Presto Query

prestodb · Jan 16, 2023 · ea64fb6 · ea64fb6
1 parent 1387be7
commit ea64fb6
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 5 deletions.
diff --git a/presto-parser/src/main/antlr4/com/facebook/presto/sql/parser/SqlBase.g4 b/presto-parser/src/main/antlr4/com/facebook/presto/sql/parser/SqlBase.g4
@@ -223,6 +223,7 @@ externalRoutineName
 queryNoWith:
       queryTerm
       (ORDER BY sortItem (',' sortItem)*)?
+      (FETCH FIRST fetchFirstNRows=INTEGER_VALUE ROWS ONLY)?
       (OFFSET offset=INTEGER_VALUE (ROW | ROWS)?)?
       (LIMIT limit=(INTEGER_VALUE | ALL))?
     ;
@@ -566,7 +567,7 @@ nonReserved
     | CALL | CALLED | CASCADE | CATALOGS | COLUMN | COLUMNS | COMMENT | COMMIT | COMMITTED | CURRENT | CURRENT_ROLE
     | DATA | DATE | DAY | DEFINER | DESC | DETERMINISTIC | DISTRIBUTED
     | EXCLUDING | EXPLAIN | EXTERNAL
-    | FILTER | FIRST | FOLLOWING | FORMAT | FUNCTION | FUNCTIONS
+    | FETCH | FILTER | FIRST | FOLLOWING | FORMAT | FUNCTION | FUNCTIONS
     | GRANT | GRANTED | GRANTS | GRAPHVIZ
     | HOUR
     | IF | IGNORE | INCLUDING | INPUT | INTERVAL | INVOKER | IO | ISOLATION
@@ -645,6 +646,7 @@ EXPLAIN: 'EXPLAIN';
 EXTRACT: 'EXTRACT';
 EXTERNAL: 'EXTERNAL';
 FALSE: 'FALSE';
+FETCH: 'FETCH';
 FILTER: 'FILTER';
 FIRST: 'FIRST';
 FOLLOWING: 'FOLLOWING';

diff --git a/presto-parser/src/main/java/com/facebook/presto/sql/parser/AstBuilder.java b/presto-parser/src/main/java/com/facebook/presto/sql/parser/AstBuilder.java
@@ -734,6 +734,14 @@ public Node visitQueryNoWith(SqlBaseParser.QueryNoWithContext context)
             offset = Optional.of(new Offset(Optional.of(getLocation(context.OFFSET())), getTextIfPresent(context.offset).orElseThrow(() -> new IllegalStateException("Missing OFFSET row count"))));
         }
 
+        Optional<String> limit = Optional.empty();
+        if (context.LIMIT() != null) {
+            limit = getTextIfPresent(context.limit);
+        }
+        else if (context.FETCH() != null) {
+            limit = getTextIfPresent(context.fetchFirstNRows);
+        }
+
         if (term instanceof QuerySpecification) {
             // When we have a simple query specification
             // followed by order by, offset, limit, fold the order by and limit
@@ -754,7 +762,7 @@ public Node visitQueryNoWith(SqlBaseParser.QueryNoWithContext context)
                             query.getHaving(),
                             orderBy,
                             offset,
-                            getTextIfPresent(context.limit)),
+                            limit),
                     Optional.empty(),
                     Optional.empty(),
                     Optional.empty());
@@ -766,7 +774,7 @@ public Node visitQueryNoWith(SqlBaseParser.QueryNoWithContext context)
                 term,
                 orderBy,
                 offset,
-                getTextIfPresent(context.limit));
+                limit);
     }
 
     @Override

diff --git a/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParser.java b/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParser.java
@@ -961,6 +961,22 @@ public void testSelectWithOffset()
                         Optional.empty(),
                         Optional.empty()));
 
+        assertStatement("SELECT * FROM table1 order by x FETCH FIRST 10 ROWS ONLY",
+                new Query(
+                        Optional.empty(),
+                        new QuerySpecification(
+                                selectList(new AllColumns()),
+                                Optional.of(new Table(QualifiedName.of("table1"))),
+                                Optional.empty(),
+                                Optional.empty(),
+                                Optional.empty(),
+                                Optional.of(new OrderBy(ImmutableList.of(new SortItem(new Identifier("x"), ASCENDING, UNDEFINED)))),
+                                Optional.empty(),
+                                Optional.of("10")),
+                        Optional.empty(),
+                        Optional.empty(),
+                        Optional.empty()));
+
         Query valuesQuery = query(values(
                 row(new LongLiteral("1"), new StringLiteral("1")),
                 row(new LongLiteral("2"), new StringLiteral("2"))));

diff --git a/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParserErrorHandling.java b/presto-parser/src/test/java/com/facebook/presto/sql/parser/TestSqlParserErrorHandling.java
@@ -49,7 +49,7 @@ public Object[][] getStatements()
                 {"select * from 'oops",
                  "line 1:15: mismatched input '''. Expecting: '(', 'LATERAL', 'UNNEST', <identifier>"},
                 {"select *\nfrom x\nfrom",
-                 "line 3:1: mismatched input 'from'. Expecting: ',', '.', 'AS', 'CROSS', 'EXCEPT', 'FULL', 'GROUP', 'HAVING', 'INNER', 'INTERSECT', 'JOIN', 'LEFT', 'LIMIT', 'NATURAL', 'OFFSET', " +
+                 "line 3:1: mismatched input 'from'. Expecting: ',', '.', 'AS', 'CROSS', 'EXCEPT', 'FETCH', 'FULL', 'GROUP', 'HAVING', 'INNER', 'INTERSECT', 'JOIN', 'LEFT', 'LIMIT', 'NATURAL', 'OFFSET', " +
                          "'ORDER', 'RIGHT', 'TABLESAMPLE', 'UNION', 'WHERE', <EOF>, <identifier>"},
                 {"select *\nfrom x\nwhere from",
                  "line 3:7: mismatched input 'from'. Expecting: <expression>"},
@@ -112,7 +112,7 @@ public Object[][] getStatements()
                 {"SELECT foo(*) filter (",
                  "line 1:23: mismatched input '<EOF>'. Expecting: 'WHERE'"},
                 {"SELECT * FROM t t x",
-                 "line 1:19: mismatched input 'x'. Expecting: '(', ',', 'CROSS', 'EXCEPT', 'FULL', 'GROUP', 'HAVING', 'INNER', 'INTERSECT', 'JOIN', 'LEFT', 'LIMIT', 'NATURAL', 'OFFSET', 'ORDER', " +
+                 "line 1:19: mismatched input 'x'. Expecting: '(', ',', 'CROSS', 'EXCEPT', 'FETCH', 'FULL', 'GROUP', 'HAVING', 'INNER', 'INTERSECT', 'JOIN', 'LEFT', 'LIMIT', 'NATURAL', 'OFFSET', 'ORDER', " +
                          "'RIGHT', 'TABLESAMPLE', 'UNION', 'WHERE', <EOF>"},
                 {"SELECT * FROM t WHERE EXISTS (",
                  "line 1:31: mismatched input '<EOF>'. Expecting: <query>"},