snowflake: add support for LATERAL FLATTEN and similar (#1026)

src/ast/query.rs

@@ -695,6 +695,13 @@ pub enum TableFactor {
         expr: Expr,
         alias: Option<TableAlias>,
     },
+    /// `e.g. LATERAL FLATTEN(<args>)[ AS <alias> ]`
+    Function {
+        lateral: bool,
+        name: ObjectName,
+        args: Vec<FunctionArg>,
+        alias: Option<TableAlias>,
+    },
     /// ```sql
     /// SELECT * FROM UNNEST ([10,20,30]) as numbers WITH OFFSET;
     /// +---------+--------+
@@ -793,6 +800,22 @@ impl fmt::Display for TableFactor {
                 }
                 Ok(())
             }
+            TableFactor::Function {
+                lateral,
+                name,
+                args,
+                alias,
+            } => {
+                if *lateral {
+                    write!(f, "LATERAL ")?;
+                }
+                write!(f, "{name}")?;
+                write!(f, "({})", display_comma_separated(args))?;
+                if let Some(alias) = alias {
+                    write!(f, " AS {alias}")?;
+                }
+                Ok(())
+            }
             TableFactor::TableFunction { expr, alias } => {
                 write!(f, "TABLE({expr})")?;
                 if let Some(alias) = alias {

src/parser/mod.rs

@@ -6555,11 +6555,21 @@ impl<'a> Parser<'a> {
     /// A table name or a parenthesized subquery, followed by optional `[AS] alias`
     pub fn parse_table_factor(&mut self) -> Result<TableFactor, ParserError> {
         if self.parse_keyword(Keyword::LATERAL) {
-            // LATERAL must always be followed by a subquery.
-            if !self.consume_token(&Token::LParen) {
-                self.expected("subquery after LATERAL", self.peek_token())?;
-            }
+            // LATERAL must always be followed by a subquery or table function.
+            if self.consume_token(&Token::LParen) {
                 self.parse_derived_table_factor(Lateral)
+            } else {
+                let name = self.parse_object_name()?;
+                self.expect_token(&Token::LParen)?;
+                let args = self.parse_optional_args()?;
+                let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
+                Ok(TableFactor::Function {
+                    lateral: true,
+                    name,
+                    args,
+                    alias,
+                })
+            }
         } else if self.parse_keyword(Keyword::TABLE) {
             // parse table function (SELECT * FROM TABLE (<expr>) [ AS <alias> ])
             self.expect_token(&Token::LParen)?;
@@ -6638,6 +6648,7 @@ impl<'a> Parser<'a> {
                     match &mut table_and_joins.relation {
                         TableFactor::Derived { alias, .. }
                         | TableFactor::Table { alias, .. }
+                        | TableFactor::Function { alias, .. }
                         | TableFactor::UNNEST { alias, .. }
                         | TableFactor::TableFunction { alias, .. }
                         | TableFactor::Pivot { alias, .. }

tests/sqlparser_common.rs

@@ -18,6 +18,8 @@
 //! sqlparser regardless of the chosen dialect (i.e. it doesn't conflict with
 //! dialect-specific parsing rules).
 
+extern crate core;
+
 use matches::assert_matches;
 use sqlparser::ast::SelectItem::UnnamedExpr;
 use sqlparser::ast::TableFactor::{Pivot, Unpivot};
@@ -37,6 +39,9 @@ use test_utils::{
 #[macro_use]
 mod test_utils;
 
+#[cfg(test)]
+use pretty_assertions::assert_eq;
+
 #[test]
 fn parse_insert_values() {
     let row = vec![
@@ -5976,12 +5981,10 @@ fn lateral_derived() {
     chk(false);
     chk(true);
 
-    let sql = "SELECT * FROM customer LEFT JOIN LATERAL generate_series(1, customer.id)";
+    let sql = "SELECT * FROM LATERAL UNNEST ([10,20,30]) as numbers WITH OFFSET;";
     let res = parse_sql_statements(sql);
     assert_eq!(
-        ParserError::ParserError(
-            "Expected subquery after LATERAL, found: generate_series".to_string()
-        ),
+        ParserError::ParserError("Expected end of statement, found: WITH".to_string()),
         res.unwrap_err()
     );
 
@@ -5995,6 +5998,60 @@ fn lateral_derived() {
     );
 }
 
+#[test]
+fn lateral_function() {
+    let sql = "SELECT * FROM customer LEFT JOIN LATERAL generate_series(1, customer.id)";
+    let actual_select_only = verified_only_select(sql);
+    let expected = Select {
+        distinct: None,
+        top: None,
+        projection: vec![SelectItem::Wildcard(WildcardAdditionalOptions {
+            opt_exclude: None,
+            opt_except: None,
+            opt_rename: None,
+            opt_replace: None,
+        })],
+        into: None,
+        from: vec![TableWithJoins {
+            relation: TableFactor::Table {
+                name: ObjectName(vec![Ident {
+                    value: "customer".to_string(),
+                    quote_style: None,
+                }]),
+                alias: None,
+                args: None,
+                with_hints: vec![],
+                version: None,
+                partitions: vec![],
+            },
+            joins: vec![Join {
+                relation: TableFactor::Function {
+                    lateral: true,
+                    name: ObjectName(vec!["generate_series".into()]),
+                    args: vec![
+                        FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(number("1")))),
+                        FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::CompoundIdentifier(
+                            vec![Ident::new("customer"), Ident::new("id")],
+                        ))),
+                    ],
+                    alias: None,
+                },
+                join_operator: JoinOperator::LeftOuter(JoinConstraint::None),
+            }],
+        }],
+        lateral_views: vec![],
+        selection: None,
+        group_by: GroupByExpr::Expressions(vec![]),
+        cluster_by: vec![],
+        distribute_by: vec![],
+        sort_by: vec![],
+        having: None,
+        named_window: vec![],
+        qualify: None,
+    };
+    assert_eq!(actual_select_only, expected);
+}
+
 #[test]
 fn parse_start_transaction() {
     match verified_stmt("START TRANSACTION READ ONLY, READ WRITE, ISOLATION LEVEL SERIALIZABLE") {

tests/sqlparser_snowflake.rs

@@ -176,6 +176,12 @@ fn parse_array() {
     );
 }
 
+#[test]
+fn parse_lateral_flatten() {
+    snowflake().verified_only_select(r#"SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88]}'), outer => true)) AS f"#);
+    snowflake().verified_only_select(r#"SELECT emp.employee_ID, emp.last_name, index, value AS project_name FROM employees AS emp, LATERAL FLATTEN(INPUT => emp.project_names) AS proj_names"#);
+}
+
 #[test]
 fn parse_json_using_colon() {
     let sql = "SELECT a:b FROM t";