feat: support multi value columns and aliases in unpivot (#1969)

Author: chenkovsky

src/ast/query.rs

@@ -1349,11 +1349,12 @@ pub enum TableFactor {
     /// ```
     ///
     /// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
+    /// See <https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-qry-select-unpivot>.
     Unpivot {
         table: Box<TableFactor>,
-        value: Ident,
+        value: Expr,
         name: Ident,
-        columns: Vec<Ident>,
+        columns: Vec<ExprWithAlias>,
         null_inclusion: Option<NullInclusion>,
         alias: Option<TableAlias>,
     },

src/ast/spans.rs

@@ -2002,9 +2002,9 @@ impl Spanned for TableFactor {
                 alias,
             } => union_spans(
                 core::iter::once(table.span())
-                    .chain(core::iter::once(value.span))
+                    .chain(core::iter::once(value.span()))
                     .chain(core::iter::once(name.span))
-                    .chain(columns.iter().map(|i| i.span))
+                    .chain(columns.iter().map(|ilist| ilist.span()))
                     .chain(alias.as_ref().map(|alias| alias.span())),
             ),
             TableFactor::MatchRecognize {

src/parser/mod.rs

@@ -14033,11 +14033,13 @@ impl<'a> Parser<'a> {
             None
         };
         self.expect_token(&Token::LParen)?;
-        let value = self.parse_identifier()?;
+        let value = self.parse_expr()?;
         self.expect_keyword_is(Keyword::FOR)?;
         let name = self.parse_identifier()?;
         self.expect_keyword_is(Keyword::IN)?;
-        let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
+        let columns = self.parse_parenthesized_column_list_inner(Mandatory, false, |p| {
+            p.parse_expr_with_alias()
+        })?;
         self.expect_token(&Token::RParen)?;
         let alias = self.maybe_parse_table_alias()?;
         Ok(TableFactor::Unpivot {

tests/sqlparser_common.rs

@@ -11009,20 +11009,14 @@ fn parse_unpivot_table() {
             index_hints: vec![],
         }),
         null_inclusion: None,
-        value: Ident {
-            value: "quantity".to_string(),
-            quote_style: None,
-            span: Span::empty(),
-        },
-
-        name: Ident {
-            value: "quarter".to_string(),
-            quote_style: None,
-            span: Span::empty(),
-        },
+        value: Expr::Identifier(Ident::new("quantity")),
+        name: Ident::new("quarter"),
         columns: ["Q1", "Q2", "Q3", "Q4"]
             .into_iter()
-            .map(Ident::new)
+            .map(|col| ExprWithAlias {
+                expr: Expr::Identifier(Ident::new(col)),
+                alias: None,
+            })
             .collect(),
         alias: Some(TableAlias {
             name: Ident::new("u"),
@@ -11084,6 +11078,129 @@ fn parse_unpivot_table() {
         verified_stmt(sql_unpivot_include_nulls).to_string(),
         sql_unpivot_include_nulls
     );
+
+    let sql_unpivot_with_alias = concat!(
+        "SELECT * FROM sales AS s ",
+        "UNPIVOT INCLUDE NULLS ",
+        "(quantity FOR quarter IN ",
+        "(Q1 AS Quater1, Q2 AS Quater2, Q3 AS Quater3, Q4 AS Quater4)) ",
+        "AS u (product, quarter, quantity)"
+    );
+
+    if let Unpivot { value, columns, .. } =
+        &verified_only_select(sql_unpivot_with_alias).from[0].relation
+    {
+        assert_eq!(
+            *columns,
+            vec![
+                ExprWithAlias {
+                    expr: Expr::Identifier(Ident::new("Q1")),
+                    alias: Some(Ident::new("Quater1")),
+                },
+                ExprWithAlias {
+                    expr: Expr::Identifier(Ident::new("Q2")),
+                    alias: Some(Ident::new("Quater2")),
+                },
+                ExprWithAlias {
+                    expr: Expr::Identifier(Ident::new("Q3")),
+                    alias: Some(Ident::new("Quater3")),
+                },
+                ExprWithAlias {
+                    expr: Expr::Identifier(Ident::new("Q4")),
+                    alias: Some(Ident::new("Quater4")),
+                },
+            ]
+        );
+        assert_eq!(*value, Expr::Identifier(Ident::new("quantity")));
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_with_alias).to_string(),
+        sql_unpivot_with_alias
+    );
+
+    let sql_unpivot_with_alias_and_multi_value = concat!(
+        "SELECT * FROM sales AS s ",
+        "UNPIVOT INCLUDE NULLS ((first_quarter, second_quarter) ",
+        "FOR half_of_the_year IN (",
+        "(Q1, Q2) AS H1, ",
+        "(Q3, Q4) AS H2",
+        "))"
+    );
+
+    if let Unpivot { value, columns, .. } =
+        &verified_only_select(sql_unpivot_with_alias_and_multi_value).from[0].relation
+    {
+        assert_eq!(
+            *columns,
+            vec![
+                ExprWithAlias {
+                    expr: Expr::Tuple(vec![
+                        Expr::Identifier(Ident::new("Q1")),
+                        Expr::Identifier(Ident::new("Q2")),
+                    ]),
+                    alias: Some(Ident::new("H1")),
+                },
+                ExprWithAlias {
+                    expr: Expr::Tuple(vec![
+                        Expr::Identifier(Ident::new("Q3")),
+                        Expr::Identifier(Ident::new("Q4")),
+                    ]),
+                    alias: Some(Ident::new("H2")),
+                },
+            ]
+        );
+        assert_eq!(
+            *value,
+            Expr::Tuple(vec![
+                Expr::Identifier(Ident::new("first_quarter")),
+                Expr::Identifier(Ident::new("second_quarter")),
+            ])
+        );
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_with_alias_and_multi_value).to_string(),
+        sql_unpivot_with_alias_and_multi_value
+    );
+
+    let sql_unpivot_with_alias_and_multi_value_and_qualifier = concat!(
+        "SELECT * FROM sales AS s ",
+        "UNPIVOT INCLUDE NULLS ((first_quarter, second_quarter) ",
+        "FOR half_of_the_year IN (",
+        "(sales.Q1, sales.Q2) AS H1, ",
+        "(sales.Q3, sales.Q4) AS H2",
+        "))"
+    );
+
+    if let Unpivot { columns, .. } =
+        &verified_only_select(sql_unpivot_with_alias_and_multi_value_and_qualifier).from[0].relation
+    {
+        assert_eq!(
+            *columns,
+            vec![
+                ExprWithAlias {
+                    expr: Expr::Tuple(vec![
+                        Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q1"),]),
+                        Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q2"),]),
+                    ]),
+                    alias: Some(Ident::new("H1")),
+                },
+                ExprWithAlias {
+                    expr: Expr::Tuple(vec![
+                        Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q3"),]),
+                        Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q4"),]),
+                    ]),
+                    alias: Some(Ident::new("H2")),
+                },
+            ]
+        );
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_with_alias_and_multi_value_and_qualifier).to_string(),
+        sql_unpivot_with_alias_and_multi_value_and_qualifier
+    );
 }
 
 #[test]
@@ -11181,20 +11298,14 @@ fn parse_pivot_unpivot_table() {
                     index_hints: vec![],
                 }),
                 null_inclusion: None,
-                value: Ident {
-                    value: "population".to_string(),
-                    quote_style: None,
-                    span: Span::empty()
-                },
-
-                name: Ident {
-                    value: "year".to_string(),
-                    quote_style: None,
-                    span: Span::empty()
-                },
+                value: Expr::Identifier(Ident::new("population")),
+                name: Ident::new("year"),
                 columns: ["population_2000", "population_2010"]
                     .into_iter()
-                    .map(Ident::new)
+                    .map(|col| ExprWithAlias {
+                        expr: Expr::Identifier(Ident::new(col)),
+                        alias: None,
+                    })
                     .collect(),
                 alias: Some(TableAlias {
                     name: Ident::new("u"),