Snowflake: Add support for Lambda functions

Author: yoavcloud

src/keywords.rs

@@ -827,6 +827,7 @@ define_keywords!(
     RECEIVE,
     RECLUSTER,
     RECURSIVE,
+    REDUCE,
     REF,
     REFERENCES,
     REFERENCING,
@@ -1051,6 +1052,7 @@ define_keywords!(
     TRACE,
     TRAILING,
     TRANSACTION,
+    TRANSFORM,
     TRANSIENT,
     TRANSLATE,
     TRANSLATE_REGEX,

src/parser/mod.rs

@@ -1556,6 +1556,43 @@ impl<'a> Parser<'a> {
             Keyword::LAMBDA if self.dialect.supports_lambda_functions() => {
                 Ok(Some(self.parse_lambda_expr()?))
             }
+            Keyword::FILTER
+            | Keyword::TRANSFORM
+            | Keyword::REDUCE if self.dialect.supports_lambda_functions() => {
+                self.expect_token(&Token::LParen)?;
+                let array = self.parse_expr()?;
+                self.expect_token(&Token::Comma)?;
+                let initial_value = if w.keyword == Keyword::REDUCE {
+                    let initial_value = self.parse_expr()?;
+                    self.expect_token(&Token::Comma)?;
+                    Some(initial_value)
+                } else {
+                    None
+                };
+                let lambda = self.parse_lambda_expr_with_typed_args()?;
+                let mut args = vec![
+                    FunctionArg::Unnamed(FunctionArgExpr::Expr(array)),
+                    FunctionArg::Unnamed(FunctionArgExpr::Expr(lambda)),
+                ];
+                if let Some(initial_value) = initial_value {
+                    args.insert(1, FunctionArg::Unnamed(FunctionArgExpr::Expr(initial_value)));
+                }
+                self.expect_token(&Token::RParen)?;
+                Ok(Some(Expr::Function(Function {
+                    name: ObjectName::from(vec![w.to_ident(w_span)]),
+                    uses_odbc_syntax: false,
+                    parameters: FunctionArguments::None,
+                    args: FunctionArguments::List(FunctionArgumentList {
+                        duplicate_treatment: None,
+                        clauses: vec![],
+                        args,
+                    }),
+                    filter: None,
+                    null_treatment: None,
+                    over: None,
+                    within_group: vec![],
+                })))
+            }
             _ if self.dialect.supports_geometric_types() => match w.keyword {
                 Keyword::CIRCLE => Ok(Some(self.parse_geometric_type(GeometricTypeKind::Circle)?)),
                 Keyword::BOX => Ok(Some(self.parse_geometric_type(GeometricTypeKind::GeometricBox)?)),
@@ -1614,21 +1651,6 @@ impl<'a> Parser<'a> {
                     syntax: LambdaSyntax::Arrow,
                 }))
             }
-            Token::Word(_)
-                if self.peek_nth_token(1).token == Token::Arrow
-                    && self.dialect.supports_lambda_functions() =>
-            {
-                let first_param = LambdaFunctionParameter {
-                    name: w.to_ident(self.peek_nth_token_ref(1).span),
-                    data_type: self.maybe_parse(|parser| parser.parse_data_type())?,
-                };
-                self.expect_token(&Token::Arrow)?;
-                Ok(Expr::Lambda(LambdaFunction {
-                    params: OneOrManyWithParens::One(first_param),
-                    body: Box::new(self.parse_expr()?),
-                    syntax: LambdaSyntax::Arrow,
-                }))
-            }
             _ => Ok(Expr::Identifier(w.to_ident(w_span))),
         }
     }
@@ -2210,12 +2232,7 @@ impl<'a> Parser<'a> {
             return Ok(None);
         }
         self.maybe_parse(|p| {
-            let params = p.parse_comma_separated(|p| {
-                Ok(LambdaFunctionParameter {
-                    name: p.parse_identifier()?,
-                    data_type: None,
-                })
-            })?;
+            let params = p.parse_comma_separated(|p| p.parse_lambda_function_parameter(false))?;
             p.expect_token(&Token::RParen)?;
             p.expect_token(&Token::Arrow)?;
             let expr = p.parse_expr()?;
@@ -2227,6 +2244,18 @@ impl<'a> Parser<'a> {
         })
     }
 
+    /// Parses a lambda expression using the arrow syntax with optionally typed arguments.
+    fn parse_lambda_expr_with_typed_args(&mut self) -> Result<Expr, ParserError> {
+        let params = self.parse_lambda_function_parameters(true)?;
+        self.expect_token(&Token::Arrow)?;
+        let body = self.parse_expr()?;
+        Ok(Expr::Lambda(LambdaFunction {
+            params,
+            body: Box::new(body),
+            syntax: LambdaSyntax::Arrow,
+        }))
+    }
+
     /// Parses a lambda expression using the `LAMBDA` keyword syntax.
     ///
     /// Syntax: `LAMBDA <params> : <expr>`
@@ -2237,40 +2266,55 @@ impl<'a> Parser<'a> {
     ///
     /// See <https://duckdb.org/docs/stable/sql/functions/lambda>
     fn parse_lambda_expr(&mut self) -> Result<Expr, ParserError> {
+        // Parse the parameters: either a single identifier or comma-separated identifiers
+        let params = self.parse_lambda_function_parameters(false)?;
+        // Expect the colon separator
+        self.expect_token(&Token::Colon)?;
+        // Parse the body expression
+        let body = self.parse_expr()?;
+        Ok(Expr::Lambda(LambdaFunction {
+            params,
+            body: Box::new(body),
+            syntax: LambdaSyntax::LambdaKeyword,
+        }))
+    }
+
+    /// Parses the parameters of a lambda function with optional typing.
+    fn parse_lambda_function_parameters(
+        &mut self,
+        typed: bool,
+    ) -> Result<OneOrManyWithParens<LambdaFunctionParameter>, ParserError> {
         // Parse the parameters: either a single identifier or comma-separated identifiers
         let params = if self.consume_token(&Token::LParen) {
             // Parenthesized parameters: (x, y)
-            let params = self.parse_comma_separated(|p| {
-                Ok(LambdaFunctionParameter {
-                    name: p.parse_identifier()?,
-                    data_type: None,
-                })
-            })?;
+            let params =
+                self.parse_comma_separated(|p| p.parse_lambda_function_parameter(typed))?;
             self.expect_token(&Token::RParen)?;
             OneOrManyWithParens::Many(params)
         } else {
             // Unparenthesized parameters: x or x, y
-            let params = self.parse_comma_separated(|p| {
-                Ok(LambdaFunctionParameter {
-                    name: p.parse_identifier()?,
-                    data_type: None,
-                })
-            })?;
+            let params =
+                self.parse_comma_separated(|p| p.parse_lambda_function_parameter(typed))?;
             if params.len() == 1 {
                 OneOrManyWithParens::One(params.into_iter().next().unwrap())
             } else {
                 OneOrManyWithParens::Many(params)
             }
         };
-        // Expect the colon separator
-        self.expect_token(&Token::Colon)?;
-        // Parse the body expression
-        let body = self.parse_expr()?;
-        Ok(Expr::Lambda(LambdaFunction {
-            params,
-            body: Box::new(body),
-            syntax: LambdaSyntax::LambdaKeyword,
-        }))
+        Ok(params)
+    }
+
+    /// Parses a single parameter of a lambda function, with optional typing.
+    fn parse_lambda_function_parameter(
+        &mut self,
+        typed: bool,
+    ) -> Result<LambdaFunctionParameter, ParserError> {
+        let name = self.parse_identifier()?;
+        let data_type = match self.peek_token().token {
+            Token::Word(_) if typed => self.maybe_parse(|p| p.parse_data_type())?,
+            _ => None,
+        };
+        Ok(LambdaFunctionParameter { name, data_type })
     }
 
     /// Tries to parse the body of an [ODBC escaping sequence]

tests/sqlparser_snowflake.rs

@@ -4562,4 +4562,6 @@ fn test_truncate_table_if_exists() {
 fn test_snowflake_lambda() {
     snowflake().verified_expr("TRANSFORM([1, 2, 3], a -> a * 2)");
     snowflake().verified_expr("TRANSFORM([1, 2, 3], a INT -> a * 2)");
+    snowflake().verified_expr("TRANSFORM([1, 2, 3], (x INT, y INT) -> (x + y))");
+    snowflake().verified_expr("REDUCE([1, 2, 3], 0, (acc, val) -> acc + val)");
 }