chore(query): add python style list comprehension (#13887)

* chore(query): add python style list comprehension

* chore(query): add python style list comprehension

Author: sundy-li

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

@@ -296,6 +296,13 @@ pub enum ExprElement {
         args: Vec<Expr>,
         lambda: Option<Lambda>,
     },
+    /// python/rust list comprehension
+    ListComprehension {
+        source: Expr,
+        param: Identifier,
+        filter: Option<Expr>,
+        result: Expr,
+    },
     /// An expression between parentheses
     Group(Expr),
     /// `[1, 2, 3]`
@@ -518,6 +525,44 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
                 span: transform_span(elem.span.0),
                 exprs,
             },
+            ExprElement::ListComprehension {
+                source,
+                param,
+                filter,
+                result,
+            } => {
+                let span = transform_span(elem.span.0);
+                let mut source = source;
+
+                // array_filter(source, filter)
+                if let Some(filter) = filter {
+                    source = Expr::FunctionCall {
+                        span,
+                        distinct: false,
+                        name: Identifier::from_name("array_filter"),
+                        args: vec![source],
+                        params: vec![],
+                        window: None,
+                        lambda: Some(Lambda {
+                            params: vec![param.clone()],
+                            expr: Box::new(filter),
+                        }),
+                    };
+                }
+                // array_map(source, result)
+                Expr::FunctionCall {
+                    span,
+                    distinct: false,
+                    name: Identifier::from_name("array_map"),
+                    args: vec![source],
+                    params: vec![],
+                    window: None,
+                    lambda: Some(Lambda {
+                        params: vec![param.clone()],
+                        expr: Box::new(result),
+                    }),
+                }
+            }
             ExprElement::Map { kvs } => Expr::Map {
                 span: transform_span(elem.span.0),
                 kvs,
@@ -1022,6 +1067,28 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
         )),
     );
 
+    // python style list comprehensions
+    // python: [i for i in range(10) if i%2==0 ]
+    // sql: [i for i in range(10) if i%2 = 0 ]
+    let list_comprehensions = check_experimental_chain_function(
+        true,
+        map(
+            rule! {
+                "[" ~ #subexpr(0) ~ FOR ~ #ident ~ IN
+                ~ #subexpr(0) ~ (IF ~ #subexpr(2))? ~ "]"
+            },
+            |(_, result, _, param, _, source, opt_filter, _)| {
+                let filter = opt_filter.map(|(_, filter)| filter);
+                ExprElement::ListComprehension {
+                    source,
+                    param,
+                    filter,
+                    result,
+                }
+            },
+        ),
+    );
+
     // Floating point literal with leading dot will be parsed as a period map access,
     // and then will be converted back to a floating point literal if the map access
     // is not following a primary element nor a postfix element.
@@ -1154,6 +1221,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
             | #trim_from : "`TRIM([(BOTH | LEADEING | TRAILING) ... FROM ...)`"
             | #is_distinct_from: "`... IS [NOT] DISTINCT FROM ...`"
             | #chain_function_call : "x.func(...)"
+            | #list_comprehensions: "[expr for x in ... [if ...]]"
             | #count_all_with_window : "`COUNT(*) OVER ...`"
             | #function_call : "<function>"
             | #case : "`CASE ... END`"
@@ -1684,6 +1752,11 @@ pub fn parse_float(text: &str) -> Result<Literal, ErrorKind> {
 
 pub fn parse_uint(text: &str, radix: u32) -> Result<Literal, ErrorKind> {
     let text = text.trim_start_matches('0');
+    let contains_underscore = text.contains('_');
+    if contains_underscore {
+        let text = text.replace(|p| p == '_', "");
+        return parse_uint(&text, radix);
+    }
 
     if text.is_empty() {
         return Ok(Literal::UInt64(0));

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

@@ -159,7 +159,7 @@ pub enum TokenKind {
     #[regex(r"0[xX][a-fA-F0-9]+")]
     MySQLLiteralHex,
 
-    #[regex(r"[0-9]+")]
+    #[regex(r"[0-9]+(_|[0-9])*")]
     LiteralInteger,
 
     #[regex(r"[0-9]+[eE][+-]?[0-9]+")]
@@ -1358,6 +1358,7 @@ impl TokenKind {
             | TokenKind::FALSE
             // | TokenKind::FOREIGN
             // | TokenKind::FREEZE
+            | TokenKind::FOR
             | TokenKind::FULL
             // | TokenKind::ILIKE
             | TokenKind::IN
@@ -1407,7 +1408,6 @@ impl TokenKind {
             | TokenKind::ATTACH
             | TokenKind::EXCEPT
             // | TokenKind::FETCH
-            | TokenKind::FOR
             | TokenKind::FROM
             // | TokenKind::GRANT
             | TokenKind::GROUP

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

@@ -731,6 +731,8 @@ fn test_expr() {
         r#"123456789012345678901234567890"#,
         r#"x'123456789012345678901234567890'"#,
         r#"1e100000000000000"#,
+        r#"100_100_000"#,
+        r#"1_12200_00"#,
         r#".1"#,
         r#"-1"#,
         r#"(1)"#,

src/query/ast/tests/it/testdata/expr.txt

@@ -257,6 +257,36 @@ Literal {
 }
 
 
+---------- Input ----------
+100_100_000
+---------- Output ---------
+100100000
+---------- AST ------------
+Literal {
+    span: Some(
+        0..11,
+    ),
+    lit: UInt64(
+        100100000,
+    ),
+}
+
+
+---------- Input ----------
+1_12200_00
+---------- Output ---------
+11220000
+---------- AST ------------
+Literal {
+    span: Some(
+        0..10,
+    ),
+    lit: UInt64(
+        11220000,
+    ),
+}
+
+
 ---------- Input ----------
 .1
 ---------- Output ---------

tests/sqllogictests/suites/query/02_function/02_0069_chain_function.test renamed to tests/sqllogictests/suites/query/02_function/02_0069_experimental_expr.test

@@ -21,3 +21,18 @@ with t(f) as (select '11|open|22|ai|33|is nothing without sam'.split('|')
 
 statement error 1008
 SELECT t.a::String.lowe() FROM numbers(1) t(a)
+
+## List Comprehension
+
+query T
+select  [ x * 100  FOR x in [1,2,3] if x % 2 = 0 ];
+----
+[200]
+
+
+query IT
+SELECT  12_000_111_222, [x.split(' ')[2]
+     FOR x IN ['OpenAI', 'I LOVE', 'YOU ALL']
+     IF x.INSTR('LOVE') > 0][1];
+----
+12000111222 LOVE