feat: support multi value column unpivot

Author: chenkovsky

src/ast/mod.rs

@@ -76,12 +76,12 @@ pub use self::query::{
     AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
     ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, ExprWithAliasAndOrderBy, Fetch, ForClause,
     ForJson, ForXml, FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias,
-    IlikeSelectItem, InputFormatClause, Interpolate, InterpolateExpr, Join, JoinConstraint,
-    JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, JsonTableNamedColumn,
-    JsonTableNestedColumn, LateralView, LimitClause, LockClause, LockType, MatchRecognizePattern,
-    MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
-    OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind, OrderByOptions,
-    PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
+    IdentsWithAlias, IlikeSelectItem, InputFormatClause, Interpolate, InterpolateExpr, Join,
+    JoinConstraint, JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling,
+    JsonTableNamedColumn, JsonTableNestedColumn, LateralView, LimitClause, LockClause, LockType,
+    MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
+    NonBlock, Offset, OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind,
+    OrderByOptions, PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
     RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
     SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SetExpr, SetOperator,
     SetQuantifier, Setting, SymbolDefinition, Table, TableAlias, TableAliasColumnDef, TableFactor,

src/ast/query.rs

@@ -745,6 +745,47 @@ impl fmt::Display for IdentWithAlias {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct IdentsWithAlias {
+    pub idents: Vec<Ident>,
+    pub alias: Option<Ident>,
+}
+
+impl IdentsWithAlias {
+    pub fn new(idents: Vec<Ident>, alias: Option<Ident>) -> Self {
+        Self { idents, alias }
+    }
+}
+
+impl fmt::Display for IdentsWithAlias {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self.idents.len() {
+            0 => Ok(()),
+            1 => {
+                if let Some(alias) = &self.alias {
+                    write!(f, "{} AS {}", self.idents[0], alias)
+                } else {
+                    write!(f, "{}", self.idents[0])
+                }
+            }
+            _ => {
+                if let Some(alias) = &self.alias {
+                    write!(
+                        f,
+                        "({}) AS {}",
+                        display_comma_separated(&self.idents),
+                        alias
+                    )
+                } else {
+                    write!(f, "({})", display_comma_separated(&self.idents))
+                }
+            }
+        }
+    }
+}
+
 /// Additional options for wildcards, e.g. Snowflake `EXCLUDE`/`RENAME` and Bigquery `EXCEPT`.
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1351,9 +1392,9 @@ pub enum TableFactor {
     /// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
     Unpivot {
         table: Box<TableFactor>,
-        value: Ident,
+        value: Vec<Ident>,
         name: Ident,
-        columns: Vec<Ident>,
+        columns: Vec<IdentsWithAlias>,
         null_inclusion: Option<NullInclusion>,
         alias: Option<TableAlias>,
     },
@@ -2035,10 +2076,17 @@ impl fmt::Display for TableFactor {
                 if let Some(null_inclusion) = null_inclusion {
                     write!(f, " {null_inclusion} ")?;
                 }
+                write!(f, "(")?;
+                if value.len() == 1 {
+                    // single value column unpivot
+                    write!(f, "{}", value[0])?;
+                } else {
+                    // multi value column unpivot
+                    write!(f, "({})", display_comma_separated(value))?;
+                }
                 write!(
                     f,
-                    "({} FOR {} IN ({}))",
-                    value,
+                    " FOR {} IN ({}))",
                     name,
                     display_comma_separated(columns)
                 )?;

src/ast/spans.rs

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

src/parser/mod.rs

@@ -10810,6 +10810,29 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_identifiers_with_alias(&mut self) -> Result<IdentsWithAlias, ParserError> {
+        let idents = match self.peek_token_ref().token {
+            Token::LParen => self.parse_parenthesized_column_list(Mandatory, false)?,
+            _ => vec![self.parse_identifier()?],
+        };
+        let alias = if self.parse_keyword(Keyword::AS) {
+            Some(self.parse_identifier()?)
+        } else {
+            None
+        };
+        Ok(IdentsWithAlias { idents, alias })
+    }
+
+    pub fn parse_parenthesized_columns_with_alias_list(
+        &mut self,
+        optional: IsOptional,
+        allow_empty: bool,
+    ) -> Result<Vec<IdentsWithAlias>, ParserError> {
+        self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| {
+            p.parse_identifiers_with_alias()
+        })
+    }
+
     /// Parses a parenthesized comma-separated list of unqualified, possibly quoted identifiers.
     /// For example: `(col1, "col 2", ...)`
     pub fn parse_parenthesized_column_list(
@@ -13882,11 +13905,20 @@ impl<'a> Parser<'a> {
             None
         };
         self.expect_token(&Token::LParen)?;
-        let value = self.parse_identifier()?;
+        let value = match self.peek_token_ref().token {
+            Token::LParen => {
+                // multi value column unpivot
+                self.parse_parenthesized_column_list(Mandatory, false)?
+            }
+            _ => {
+                // single value column unpivot
+                vec![self.parse_identifier()?]
+            }
+        };
         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_columns_with_alias_list(Mandatory, false)?;
         self.expect_token(&Token::RParen)?;
         let alias = self.maybe_parse_table_alias()?;
         Ok(TableFactor::Unpivot {

tests/sqlparser_common.rs

@@ -10947,11 +10947,11 @@ fn parse_unpivot_table() {
             index_hints: vec![],
         }),
         null_inclusion: None,
-        value: Ident {
+        value: vec![Ident {
             value: "quantity".to_string(),
             quote_style: None,
             span: Span::empty(),
-        },
+        }],
 
         name: Ident {
             value: "quarter".to_string(),
@@ -10960,7 +10960,7 @@ fn parse_unpivot_table() {
         },
         columns: ["Q1", "Q2", "Q3", "Q4"]
             .into_iter()
-            .map(Ident::new)
+            .map(|col| IdentsWithAlias::new(vec![Ident::new(col)], None))
             .collect(),
         alias: Some(TableAlias {
             name: Ident::new("u"),
@@ -11022,6 +11022,67 @@ 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![
+                IdentsWithAlias::new(vec![Ident::new("Q1")], Some(Ident::new("Quater1"))),
+                IdentsWithAlias::new(vec![Ident::new("Q2")], Some(Ident::new("Quater2"))),
+                IdentsWithAlias::new(vec![Ident::new("Q3")], Some(Ident::new("Quater3"))),
+                IdentsWithAlias::new(vec![Ident::new("Q4")], Some(Ident::new("Quater4"))),
+            ]
+        );
+        assert_eq!(*value, vec![Ident::new("quantity")]);
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_with_alias).to_string(),
+        sql_unpivot_with_alias
+    );
+
+    let sql_unpivot_with_alias = 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).from[0].relation
+    {
+        assert_eq!(
+            *columns,
+            vec![
+                IdentsWithAlias::new(
+                    vec![Ident::new("Q1"), Ident::new("Q2")],
+                    Some(Ident::new("H1"))
+                ),
+                IdentsWithAlias::new(
+                    vec![Ident::new("Q3"), Ident::new("Q4")],
+                    Some(Ident::new("H2"))
+                ),
+            ]
+        );
+        assert_eq!(
+            *value,
+            vec![Ident::new("first_quarter"), Ident::new("second_quarter")]
+        );
+    }
+
+    assert_eq!(
+        verified_stmt(sql_unpivot_with_alias).to_string(),
+        sql_unpivot_with_alias
+    );
 }
 
 #[test]
@@ -11119,11 +11180,11 @@ fn parse_pivot_unpivot_table() {
                     index_hints: vec![],
                 }),
                 null_inclusion: None,
-                value: Ident {
+                value: vec![Ident {
                     value: "population".to_string(),
                     quote_style: None,
                     span: Span::empty()
-                },
+                }],
 
                 name: Ident {
                     value: "year".to_string(),
@@ -11132,7 +11193,7 @@ fn parse_pivot_unpivot_table() {
                 },
                 columns: ["population_2000", "population_2010"]
                     .into_iter()
-                    .map(Ident::new)
+                    .map(|col| IdentsWithAlias::new(vec![Ident::new(col)], None))
                     .collect(),
                 alias: Some(TableAlias {
                     name: Ident::new("u"),