feat(query): top k syntax support

src/query/ast/src/ast/query.rs

@@ -143,6 +143,8 @@ pub struct SelectStmt {
     pub hints: Option<Hint>,
     #[drive(skip)]
     pub distinct: bool,
+    #[drive(skip)]
+    pub top_n: Option<u64>,
     // Result set of current subquery
     pub select_list: Vec<SelectTarget>,
     // `FROM` clause, a list of table references.
@@ -171,6 +173,9 @@ impl Display for SelectStmt {
         if self.distinct {
             write!(f, "DISTINCT ")?;
         }
+        if let Some(topn) = &self.top_n {
+            write!(f, "TOP {} ", topn)?;
+        }
         write_comma_separated_list(f, &self.select_list)?;
 
         // FROM clause

src/query/ast/src/parser/query.rs

@@ -30,6 +30,7 @@ use crate::parser::expr::*;
 use crate::parser::input::Input;
 use crate::parser::input::WithSpan;
 use crate::parser::statement::hint;
+use crate::parser::statement::top_n;
 use crate::parser::token::*;
 use crate::parser::ErrorKind;
 use crate::rule;
@@ -53,6 +54,7 @@ pub enum SetOperationElement {
     SelectStmt {
         hints: Option<Hint>,
         distinct: bool,
+        top_n: Option<u64>,
         select_list: Vec<SelectTarget>,
         from: Vec<TableReference>,
         selection: Option<Expr>,
@@ -101,7 +103,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
     let select_stmt = map_res(
         rule! {
             ( FROM ~ ^#comma_separated_list1(table_reference) )?
-            ~ SELECT ~ #hint? ~ DISTINCT? ~ ^#comma_separated_list1(select_target)
+            ~ SELECT ~ #hint? ~ DISTINCT? ~ #top_n? ~ ^#comma_separated_list1(select_target)
             ~ ( FROM ~ ^#comma_separated_list1(table_reference) )?
             ~ ( WHERE ~ ^#expr )?
             ~ ( GROUP ~ ^BY ~ ^#group_by_items )?
@@ -114,6 +116,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
             _select,
             opt_hints,
             opt_distinct,
+            opt_top_n,
             select_list,
             opt_from_block_second,
             opt_where_block,
@@ -131,6 +134,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
             Ok(SetOperationElement::SelectStmt {
                 hints: opt_hints,
                 distinct: opt_distinct.is_some(),
+                top_n: opt_top_n,
                 select_list,
                 from: opt_from_block_first
                     .or(opt_from_block_second)
@@ -229,6 +233,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, SetOperationElement>>> PrattParser<I>
             SetOperationElement::SelectStmt {
                 hints,
                 distinct,
+                top_n,
                 select_list,
                 from,
                 selection,
@@ -239,6 +244,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, SetOperationElement>>> PrattParser<I>
             } => SetExpr::Select(Box::new(SelectStmt {
                 span: transform_span(input.span.tokens),
                 hints,
+                top_n,
                 distinct,
                 select_list,
                 from,

src/query/ast/src/parser/statement.rs

@@ -29,6 +29,7 @@ use minitrace::func_name;
 use nom::branch::alt;
 use nom::combinator::consumed;
 use nom::combinator::map;
+use nom::combinator::not;
 use nom::combinator::value;
 use nom::Slice;
 
@@ -2602,6 +2603,22 @@ pub fn hint(i: Input) -> IResult<Hint> {
     rule!(#hint|#invalid_hint)(i)
 }
 
+pub fn top_n(i: Input) -> IResult<u64> {
+    map(
+        rule! {
+            TOP
+            ~ ^#error_hint(
+                not(literal_u64),
+                "expecting a literal number after keyword `TOP`, if you were refering a column with name `top`, \
+                        please quote it like `\"top\"`"
+            )
+            ~ ^#literal_u64
+            : "TOP <limit>"
+        },
+        |(_, _, n)| n,
+    )(i)
+}
+
 pub fn rest_str(i: Input) -> IResult<(String, usize)> {
     // It's safe to unwrap because input must contain EOI.
     let first_token = i.tokens.first().unwrap();

src/query/ast/src/parser/token.rs

@@ -157,6 +157,9 @@ pub enum TokenKind {
     #[regex(r#"@([^\s`;'"()]|\\\s|\\'|\\"|\\\\)+"#)]
     AtString,
 
+    #[token("TOP", ignore(ascii_case))]
+    TOP,
+
     #[regex(r"[xX]'[a-fA-F0-9]*'")]
     PGLiteralHex,
     #[regex(r"0[xX][a-fA-F0-9]+")]

src/query/ast/tests/it/parser.rs

@@ -193,6 +193,7 @@ fn test_statement() {
         r#"DROP database if exists db1;"#,
         r#"select distinct a, count(*) from t where a = 1 and b - 1 < a group by a having a = 1;"#,
         r#"select * from t4;"#,
+        r#"select top 2 * from t4;"#,
         r#"select * from aa.bb;"#,
         r#"select * from a, b, c;"#,
         r#"select * from a, b, c order by "db"."a"."c1";"#,
@@ -847,6 +848,8 @@ fn test_statement_error() {
         r#"PRESIGN INVALID @my_stage/path/to/file"#,
         r#"SELECT c a as FROM t"#,
         r#"SELECT c a as b FROM t"#,
+        r#"SELECT top -1 c a as b FROM t"#,
+        r#"SELECT top c a as b FROM t"#,
         r#"SELECT * FROM t GROUP BY GROUPING SETS a, b"#,
         r#"SELECT * FROM t GROUP BY GROUPING SETS ()"#,
         r#"select * from aa.bb limit 10 order by bb;"#,