DuckDB: support LIMIT 10% and TABLESAMPLE without a method (#2802)

* DuckDB: support LIMIT 10% and TABLESAMPLE without a method

- `LIMIT 10%` and `LIMIT 10 PERCENT` set the new `Limit.percentUnit`.
  Inside a LIMIT, a `%` followed by the end of the clause is the unit,
  not a modulo; `LIMIT 10 % 3` and `LIMIT a % offset` are unchanged.
- TABLESAMPLE takes DuckDB's optional method and reservoir, as USING
  SAMPLE already does: `TABLESAMPLE 10%`, `TABLESAMPLE reservoir(10 ROWS)`.
- A table function keeps and prints its sample clause: `range(10) r
  USING SAMPLE 10%` used to lose it.

* Find the percent of LIMIT 10% once and keep columns named offset

- LIMIT scans its clause once for the top-level % followed by its end,
  and the Pratt loop stops at that token only, so `LIMIT a % offset + 1`
  and modulos nested in sub-selects are unchanged.
- OFFSET ends the clause only before a value, FETCH only before FIRST or
  NEXT.
- After DuckDB's TABLESAMPLE n, OFFSET belongs to the query.
- The unit is printed after NULL and ALL too.

Author: guizmaii

src/main/java/net/sf/jsqlparser/statement/select/Limit.java

@@ -32,6 +32,8 @@ public class Limit extends ASTNodeAccessImpl {
      */
     private ExpressionList<Expression> byExpressions;
 
+    private String percentUnit;
+
     public Expression getOffset() {
         return offset;
     }
@@ -72,19 +74,39 @@ public void setLimitNull(boolean b) {
         }
     }
 
+    /**
+     * DuckDB's {@code %} or {@code PERCENT} after the row count, which then is a percentage of the
+     * rows, or {@code null}.
+     */
+    public String getPercentUnit() {
+        return percentUnit;
+    }
+
+    public Limit setPercentUnit(String percentUnit) {
+        this.percentUnit = percentUnit;
+        return this;
+    }
+
+    String percentUnitToString() {
+        if (percentUnit == null) {
+            return "";
+        }
+        return "%".equals(percentUnit) ? "%" : " " + percentUnit;
+    }
+
     @Override
     public String toString() {
         String retVal = " LIMIT ";
 
         if (rowCount instanceof AllValue || rowCount instanceof NullValue) {
             // no offset allowed
-            retVal += rowCount;
+            retVal += rowCount + percentUnitToString();
         } else {
             if (null != offset) {
                 retVal += offset + ", ";
             }
             if (null != rowCount) {
-                retVal += rowCount;
+                retVal += rowCount + percentUnitToString();
             }
         }
 

src/main/java/net/sf/jsqlparser/statement/select/TableFunction.java

@@ -18,6 +18,8 @@
 
 @SuppressWarnings({"PMD.UncommentedEmptyMethodBody"})
 public class TableFunction extends Function implements FromItem {
+    private SampleClause sampleClause;
+
     private String prefix = null;
     private Alias alias = null;
     private Pivot pivot = null;
@@ -220,12 +222,13 @@ public TableFunction withUnPivot(UnPivot unpivot) {
 
     @Override
     public SampleClause getSampleClause() {
-        return null;
+        return sampleClause;
     }
 
     @Override
-    public FromItem setSampleClause(SampleClause sampleClause) {
-        return null;
+    public TableFunction setSampleClause(SampleClause sampleClause) {
+        this.sampleClause = sampleClause;
+        return this;
     }
 
     public StringBuilder appendTo(StringBuilder builder) {
@@ -263,6 +266,9 @@ public StringBuilder appendTo(StringBuilder builder, Consumer<Expression> expres
                 builder.append(offsetAlias);
             }
         }
+        if (sampleClause != null) {
+            sampleClause.appendTo(builder);
+        }
         return builder;
     }
 

src/main/java/net/sf/jsqlparser/util/deparser/LimitDeparser.java

@@ -27,20 +27,18 @@ public void deParse(Limit limit) {
 
     public <S> void deParse(Limit limit, S context) {
         builder.append(" LIMIT ");
-        if (limit.isLimitNull()) {
-            builder.append("NULL");
+        if (limit.isLimitNull() || limit.isLimitAll()) {
+            builder.append(limit.isLimitNull() ? "NULL" : "ALL");
+            appendPercentUnit(limit);
         } else {
-            if (limit.isLimitAll()) {
-                builder.append("ALL");
-            } else {
-                if (null != limit.getOffset()) {
-                    limit.getOffset().accept(expressionVisitor, context);
-                    builder.append(", ");
-                }
-
-                if (null != limit.getRowCount()) {
-                    limit.getRowCount().accept(expressionVisitor, context);
-                }
+            if (null != limit.getOffset()) {
+                limit.getOffset().accept(expressionVisitor, context);
+                builder.append(", ");
+            }
+
+            if (null != limit.getRowCount()) {
+                limit.getRowCount().accept(expressionVisitor, context);
+                appendPercentUnit(limit);
             }
         }
 
@@ -50,6 +48,12 @@ public <S> void deParse(Limit limit, S context) {
         }
     }
 
+    private void appendPercentUnit(Limit limit) {
+        if (limit.getPercentUnit() != null) {
+            builder.append("%".equals(limit.getPercentUnit()) ? "%" : " " + limit.getPercentUnit());
+        }
+    }
+
     public ExpressionVisitor<StringBuilder> getExpressionVisitor() {
         return expressionVisitor;
     }

src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt

@@ -687,6 +687,8 @@ public class CCJSqlParser extends AbstractJSqlParser<CCJSqlParser> {
             else {
                 // Handle unnamed tokens by image
                 String img = t.image;
+                // DuckDB's LIMIT 10%: the % ends the row count instead of starting a modulo
+                if (t == limitPercentToken) break;
                 if      ("*".equals(img) || "%".equals(img))                  prec = 6;
                 else if ("+".equals(img) || "-".equals(img) || "&".equals(img)) prec = 5;
                 else break; // not an arithmetic operator
@@ -1530,6 +1532,51 @@ public class CCJSqlParser extends AbstractJSqlParser<CCJSqlParser> {
         }
     }
 
+    // the % of DuckDB's LIMIT 10% in the LIMIT clause being parsed, which ends its row count
+    private Token limitPercentToken = null;
+
+    /**
+     * Finds the % that ends the row count of the LIMIT clause starting after the current token:
+     * at its top level and directly followed by the end of the clause, as in LIMIT 10% OFFSET 2.
+     */
+    private Token findLimitPercentToken() {
+        int depth = 0;
+        for (Token token = nextWithItemToken(getToken(0));; token = nextWithItemToken(token)) {
+            if (token.kind == EOF || token.kind == ST_SEMICOLON) {
+                return null;
+            }
+            if ("(".equals(token.image) || "[".equals(token.image)) {
+                depth++;
+            } else if (")".equals(token.image) || "]".equals(token.image)) {
+                if (--depth < 0) {
+                    return null;
+                }
+            } else if (depth == 0 && "%".equals(token.image)
+                    && isLimitClauseEnd(nextWithItemToken(token))) {
+                return token;
+            }
+        }
+    }
+
+    /** Whether the token ends a LIMIT clause; OFFSET and FETCH only when they start a clause. */
+    private boolean isLimitClauseEnd(Token token) {
+        switch (token.kind) {
+            case EOF: case ST_SEMICOLON: case CLOSING_BRACKET:
+            case K_UNION: case K_EXCEPT: case K_INTERSECT: case K_MINUS:
+                return true;
+            case K_OFFSET:
+                // not a column named offset, as in LIMIT a % offset + 1
+                Token value = nextWithItemToken(token);
+                return value.kind == S_LONG || value.kind == OPENING_BRACKET
+                        || "?".equals(value.image) || ":".equals(value.image);
+            case K_FETCH:
+                Token next = nextWithItemToken(token);
+                return next.kind == K_FIRST || next.kind == K_NEXT;
+            default:
+                return false;
+        }
+    }
+
     /** DuckDB sampling methods, which are ordinary identifiers except for SYSTEM/BERNOULLI. */
     private boolean isSampleMethod(Token token) {
         return token.kind == K_SYSTEM || token.kind == K_BERNOULLI
@@ -7146,9 +7193,13 @@ SampleClause SampleClause():
         )
         |
         (
-            // SQL:2016 compliant
+            // SQL:2016 compliant; DuckDB: the method is optional and may be reservoir
             token = <K_TABLESAMPLE> { keyword = token.image; }
-            ( token = <K_SYSTEM> | token = <K_BERNOULLI> ) { method = token.image; }
+            [
+                LOOKAHEAD({ isSampleMethod(getToken(1)) })
+                ( token = <K_SYSTEM> | token = <K_BERNOULLI> | token = <S_IDENTIFIER> )
+                { method = token.image; }
+            ]
         )
         |
         (
@@ -7189,7 +7240,8 @@ SampleClause SampleClause():
                 <K_ROWS> { percentageUnit="ROWS"; }
             )
         ]
-        [ LOOKAHEAD(2) <K_OFFSET> offsetArgument = Number() ]
+        // ClickHouse's SAMPLE k OFFSET m; after DuckDB's TABLESAMPLE n, OFFSET belongs to the query
+        [ LOOKAHEAD(2, { !"TABLESAMPLE".equalsIgnoreCase(keyword) }) <K_OFFSET> offsetArgument = Number() ]
     )
 
     // DuckDB: USING SAMPLE 10 PERCENT (bernoulli[, seed])
@@ -10311,11 +10363,17 @@ Limit LimitWithOffset() #LimitWithOffset:
 }
 {
     <K_LIMIT>
-    (
-        LOOKAHEAD(3) firstExpression = ParenthesedSelect()
-        |
-        firstExpression = Expression()
-    )
+    { Token outerLimitPercentToken = limitPercentToken; }
+    try {
+        { limitPercentToken = findLimitPercentToken(); }
+        (
+            LOOKAHEAD(3) firstExpression = ParenthesedSelect()
+            |
+            firstExpression = Expression()
+        )
+    } finally {
+        limitPercentToken = outerLimitPercentToken;
+    }
     (
         // MySQL: LIMIT offset, row_count
         ","
@@ -10325,10 +10383,15 @@ Limit LimitWithOffset() #LimitWithOffset:
             limit.setRowCount(secondExpression);
         }
         |
-        // PostgreSQL: LIMIT row_count
+        // PostgreSQL: LIMIT row_count; DuckDB: LIMIT 10% or LIMIT 10 PERCENT
         {
             limit.setRowCount(firstExpression);
         }
+        [
+            LOOKAHEAD({ getToken(1).kind == K_PERCENT
+                    || "%".equals(getToken(1).image) && isLimitClauseEnd(getToken(2)) })
+            ( "%" { limit.setPercentUnit("%"); } | <K_PERCENT> { limit.setPercentUnit("PERCENT"); } )
+        ]
     )
     {
         linkAST(limit, jjtThis);

src/test/java/net/sf/jsqlparser/statement/select/DuckDBTest.java

@@ -865,4 +865,42 @@ void testTableQueryKeepsTableColumnsAndTrailingClauses() throws JSQLParserExcept
         Assertions.assertInstanceOf(net.sf.jsqlparser.statement.piped.FromQuery.class,
                 ((ParenthesedSelect) CCJSqlParserUtil.parse("(FROM t)")).getSelect());
     }
+
+    @Test
+    void testLimitPercent() throws JSQLParserException {
+        for (String sqlStr : new String[] {"SELECT a FROM t LIMIT 10%",
+                "SELECT a FROM t LIMIT 10.5%", "SELECT a FROM t LIMIT 40 PERCENT",
+                "SELECT a FROM t LIMIT 50% OFFSET 2", "SELECT a FROM t LIMIT (SELECT 10)%",
+                "SELECT * FROM (SELECT a FROM t LIMIT 10%) x",
+                "SELECT a FROM t LIMIT (SELECT a FROM u LIMIT 5%)", "SELECT a FROM t LIMIT NULL%",
+                // modulo and a column named offset are unchanged
+                "SELECT a FROM t LIMIT 10 % 3", "SELECT a FROM t LIMIT a % offset",
+                "SELECT a FROM t LIMIT a % offset + 1", "SELECT a FROM t LIMIT a % offset OFFSET 2",
+                "SELECT a FROM t LIMIT (SELECT a % offset + 1 FROM u)"}) {
+            TestUtils.assertSqlCanBeParsedAndDeparsed(sqlStr, true);
+        }
+        PlainSelect select = (PlainSelect) CCJSqlParserUtil.parse("SELECT a FROM t LIMIT 10%");
+        Assertions.assertEquals("%", select.getLimit().getPercentUnit());
+        Assertions.assertEquals("10", select.getLimit().getRowCount().toString());
+
+        select = (PlainSelect) CCJSqlParserUtil.parse("SELECT a FROM t LIMIT a % offset + 1");
+        Assertions.assertNull(select.getLimit().getPercentUnit());
+        Assertions.assertNull(select.getOffset());
+    }
+
+    @Test
+    void testTableSampleWithoutOrWithReservoirMethod() throws JSQLParserException {
+        for (String sqlStr : new String[] {"SELECT a FROM t TABLESAMPLE 10%",
+                "SELECT a FROM t TABLESAMPLE 10 ROWS", "SELECT a FROM t t1 TABLESAMPLE 10 PERCENT",
+                "SELECT a FROM t TABLESAMPLE RESERVOIR (10 ROWS) REPEATABLE (1)",
+                "SELECT a FROM (SELECT 1) TABLESAMPLE 10%",
+                "SELECT a FROM read_csv('x') TABLESAMPLE 10%",
+                "SELECT a FROM range(10) r USING SAMPLE 10%"}) {
+            TestUtils.assertSqlCanBeParsedAndDeparsed(sqlStr, true);
+        }
+        // after DuckDB's TABLESAMPLE n, OFFSET belongs to the query
+        PlainSelect select = (PlainSelect) CCJSqlParserUtil
+                .parse("SELECT a FROM t TABLESAMPLE 10 OFFSET 3");
+        Assertions.assertEquals("3", select.getOffset().getOffset().toString());
+    }
 }