snowflake: support for UNPIVOT and a fix for chained PIVOTs (#983)

src/ast/query.rs

@@ -720,13 +720,28 @@ pub enum TableFactor {
     /// For example `FROM monthly_sales PIVOT(sum(amount) FOR MONTH IN ('JAN', 'FEB'))`
     /// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot>
     Pivot {
-        #[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
+        #[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
-        name: ObjectName,
+        table: Box<TableFactor>,
-        table_alias: Option<TableAlias>,
         aggregate_function: Expr, // Function expression
         value_column: Vec<Ident>,
         pivot_values: Vec<Value>,
-        pivot_alias: Option<TableAlias>,
+        alias: Option<TableAlias>,
+    },
+    /// An UNPIVOT operation on a table.
+    ///
+    /// Syntax:
+    /// ```sql
+    /// table UNPIVOT(value FOR name IN (column1, [ column2, ... ])) [ alias ]
+    /// ```
+    ///
+    /// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
+    Unpivot {
+        #[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
+        table: Box<TableFactor>,
+        value: Ident,
+        name: Ident,
+        columns: Vec<Ident>,
+        alias: Option<TableAlias>,
     },
 }
 
@@ -810,32 +825,42 @@ impl fmt::Display for TableFactor {
                 Ok(())
             }
             TableFactor::Pivot {
-                name,
+                table,
-                table_alias,
                 aggregate_function,
                 value_column,
                 pivot_values,
-                pivot_alias,
+                alias,
             } => {
-                write!(f, "{}", name)?;
-                if table_alias.is_some() {
-                    write!(f, " AS {}", table_alias.as_ref().unwrap())?;
-                }
                 write!(
                     f,
-                    " PIVOT({} FOR {} IN (",
+                    "{} PIVOT({} FOR {} IN ({}))",
+                    table,
                     aggregate_function,
-                    Expr::CompoundIdentifier(value_column.to_vec())
+                    Expr::CompoundIdentifier(value_column.to_vec()),
+                    display_comma_separated(pivot_values)
                 )?;
-                for value in pivot_values {
+                if alias.is_some() {
-                    write!(f, "{}", value)?;
+                    write!(f, " AS {}", alias.as_ref().unwrap())?;
-                    if !value.eq(pivot_values.last().unwrap()) {
-                        write!(f, ", ")?;
-                    }
                 }
-                write!(f, "))")?;
+                Ok(())
-                if pivot_alias.is_some() {
+            }
-                    write!(f, " AS {}", pivot_alias.as_ref().unwrap())?;
+            TableFactor::Unpivot {
+                table,
+                value,
+                name,
+                columns,
+                alias,
+            } => {
+                write!(
+                    f,
+                    "{} UNPIVOT({} FOR {} IN ({}))",
+                    table,
+                    value,
+                    name,
+                    display_comma_separated(columns)
+                )?;
+                if alias.is_some() {
+                    write!(f, " AS {}", alias.as_ref().unwrap())?;
                 }
                 Ok(())
             }

src/keywords.rs

@@ -635,6 +635,7 @@ define_keywords!(
     UNKNOWN,
     UNLOGGED,
     UNNEST,
+    UNPIVOT,
     UNSIGNED,
     UNTIL,
     UPDATE,
@@ -693,6 +694,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     Keyword::HAVING,
     Keyword::ORDER,
     Keyword::PIVOT,
+    Keyword::UNPIVOT,
     Keyword::TOP,
     Keyword::LATERAL,
     Keyword::VIEW,

src/parser/mod.rs

@@ -6276,9 +6276,8 @@ impl<'a> Parser<'a> {
                         | TableFactor::Table { alias, .. }
                         | TableFactor::UNNEST { alias, .. }
                         | TableFactor::TableFunction { alias, .. }
-                        | TableFactor::Pivot {
+                        | TableFactor::Pivot { alias, .. }
-                            pivot_alias: alias, ..
+                        | TableFactor::Unpivot { alias, .. }
-                        }
                         | TableFactor::NestedJoin { alias, .. } => {
                             // but not `FROM (mytable AS alias1) AS alias2`.
                             if let Some(inner_alias) = alias {
@@ -6357,11 +6356,6 @@ impl<'a> Parser<'a> {
 
             let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
 
-            // Pivot
-            if self.parse_keyword(Keyword::PIVOT) {
-                return self.parse_pivot_table_factor(name, alias);
-            }
-
             // MSSQL-specific table hints:
             let mut with_hints = vec![];
             if self.parse_keyword(Keyword::WITH) {
@@ -6373,14 +6367,25 @@ impl<'a> Parser<'a> {
                     self.prev_token();
                 }
             };
-            Ok(TableFactor::Table {
+
+            let mut table = TableFactor::Table {
                 name,
                 alias,
                 args,
                 with_hints,
                 version,
                 partitions,
-            })
+            };
+
+            while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {
+                table = match kw {
+                    Keyword::PIVOT => self.parse_pivot_table_factor(table)?,
+                    Keyword::UNPIVOT => self.parse_unpivot_table_factor(table)?,
+                    _ => unreachable!(),
+                }
+            }
+
+            Ok(table)
         }
     }
 
@@ -6417,8 +6422,7 @@ impl<'a> Parser<'a> {
 
     pub fn parse_pivot_table_factor(
         &mut self,
-        name: ObjectName,
+        table: TableFactor,
-        table_alias: Option<TableAlias>,
     ) -> Result<TableFactor, ParserError> {
         self.expect_token(&Token::LParen)?;
         let function_name = match self.next_token().token {
@@ -6435,12 +6439,32 @@ impl<'a> Parser<'a> {
         self.expect_token(&Token::RParen)?;
         let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
         Ok(TableFactor::Pivot {
-            name,
+            table: Box::new(table),
-            table_alias,
             aggregate_function: function,
             value_column,
             pivot_values,
-            pivot_alias: alias,
+            alias,
+        })
+    }
+
+    pub fn parse_unpivot_table_factor(
+        &mut self,
+        table: TableFactor,
+    ) -> Result<TableFactor, ParserError> {
+        self.expect_token(&Token::LParen)?;
+        let value = self.parse_identifier()?;
+        self.expect_keyword(Keyword::FOR)?;
+        let name = self.parse_identifier()?;
+        self.expect_keyword(Keyword::IN)?;
+        let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
+        self.expect_token(&Token::RParen)?;
+        let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
+        Ok(TableFactor::Unpivot {
+            table: Box::new(table),
+            value,
+            name,
+            columns,
+            alias,
         })
     }
 

tests/sqlparser_common.rs

@@ -20,7 +20,7 @@
 
 use matches::assert_matches;
 use sqlparser::ast::SelectItem::UnnamedExpr;
-use sqlparser::ast::TableFactor::Pivot;
+use sqlparser::ast::TableFactor::{Pivot, Unpivot};
 use sqlparser::ast::*;
 use sqlparser::dialect::{
     AnsiDialect, BigQueryDialect, ClickHouseDialect, DuckDbDialect, GenericDialect, HiveDialect,
@@ -7257,10 +7257,16 @@ fn parse_pivot_table() {
     assert_eq!(
         verified_only_select(sql).from[0].relation,
         Pivot {
-            name: ObjectName(vec![Ident::new("monthly_sales")]),
+            table: Box::new(TableFactor::Table {
-            table_alias: Some(TableAlias {
+                name: ObjectName(vec![Ident::new("monthly_sales")]),
-                name: Ident::new("a"),
+                alias: Some(TableAlias {
-                columns: vec![]
+                    name: Ident::new("a"),
+                    columns: vec![]
+                }),
+                args: None,
+                with_hints: vec![],
+                version: None,
+                partitions: vec![],
             }),
             aggregate_function: Expr::Function(Function {
                 name: ObjectName(vec![Ident::new("SUM")]),
@@ -7279,7 +7285,7 @@ fn parse_pivot_table() {
                 Value::SingleQuotedString("MAR".to_string()),
                 Value::SingleQuotedString("APR".to_string()),
             ],
-            pivot_alias: Some(TableAlias {
+            alias: Some(TableAlias {
                 name: Ident {
                     value: "p".to_string(),
                     quote_style: None
@@ -7290,17 +7296,15 @@ fn parse_pivot_table() {
     );
     assert_eq!(verified_stmt(sql).to_string(), sql);
 
+    // parsing should succeed with empty alias
     let sql_without_table_alias = concat!(
         "SELECT * FROM monthly_sales ",
         "PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
         "ORDER BY EMPID"
     );
     assert_matches!(
-        verified_only_select(sql_without_table_alias).from[0].relation,
+        &verified_only_select(sql_without_table_alias).from[0].relation,
-        Pivot {
+        Pivot { table, .. } if matches!(&**table, TableFactor::Table { alias: None, .. })
-            table_alias: None, // parsing should succeed with empty alias
-            ..
-        }
     );
     assert_eq!(
         verified_stmt(sql_without_table_alias).to_string(),
@@ -7308,6 +7312,135 @@ fn parse_pivot_table() {
     );
 }
 
+#[test]
+fn parse_unpivot_table() {
+    let sql = concat!(
+        "SELECT * FROM sales AS s ",
+        "UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)"
+    );
+
+    pretty_assertions::assert_eq!(
+        verified_only_select(sql).from[0].relation,
+        Unpivot {
+            table: Box::new(TableFactor::Table {
+                name: ObjectName(vec![Ident::new("sales")]),
+                alias: Some(TableAlias {
+                    name: Ident::new("s"),
+                    columns: vec![]
+                }),
+                args: None,
+                with_hints: vec![],
+                version: None,
+                partitions: vec![],
+            }),
+            value: Ident {
+                value: "quantity".to_string(),
+                quote_style: None
+            },
+
+            name: Ident {
+                value: "quarter".to_string(),
+                quote_style: None
+            },
+            columns: ["Q1", "Q2", "Q3", "Q4"]
+                .into_iter()
+                .map(Ident::new)
+                .collect(),
+            alias: Some(TableAlias {
+                name: Ident::new("u"),
+                columns: ["product", "quarter", "quantity"]
+                    .into_iter()
+                    .map(Ident::new)
+                    .collect()
+            }),
+        }
+    );
+    assert_eq!(verified_stmt(sql).to_string(), sql);
+
+    let sql_without_aliases = concat!(
+        "SELECT * FROM sales ",
+        "UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4))"
+    );
+
+    assert_matches!(
+        &verified_only_select(sql_without_aliases).from[0].relation,
+        Unpivot {
+            table,
+            alias: None,
+            ..
+        } if matches!(&**table, TableFactor::Table { alias: None, .. })
+    );
+    assert_eq!(
+        verified_stmt(sql_without_aliases).to_string(),
+        sql_without_aliases
+    );
+}
+
+#[test]
+fn parse_pivot_unpivot_table() {
+    let sql = concat!(
+        "SELECT * FROM census AS c ",
+        "UNPIVOT(population FOR year IN (population_2000, population_2010)) AS u ",
+        "PIVOT(sum(population) FOR year IN ('population_2000', 'population_2010')) AS p"
+    );
+
+    pretty_assertions::assert_eq!(
+        verified_only_select(sql).from[0].relation,
+        Pivot {
+            table: Box::new(Unpivot {
+                table: Box::new(TableFactor::Table {
+                    name: ObjectName(vec![Ident::new("census")]),
+                    alias: Some(TableAlias {
+                        name: Ident::new("c"),
+                        columns: vec![]
+                    }),
+                    args: None,
+                    with_hints: vec![],
+                    version: None,
+                    partitions: vec![],
+                }),
+                value: Ident {
+                    value: "population".to_string(),
+                    quote_style: None
+                },
+
+                name: Ident {
+                    value: "year".to_string(),
+                    quote_style: None
+                },
+                columns: ["population_2000", "population_2010"]
+                    .into_iter()
+                    .map(Ident::new)
+                    .collect(),
+                alias: Some(TableAlias {
+                    name: Ident::new("u"),
+                    columns: vec![]
+                }),
+            }),
+            aggregate_function: Expr::Function(Function {
+                name: ObjectName(vec![Ident::new("sum")]),
+                args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                    Expr::Identifier(Ident::new("population"))
+                ))]),
+                over: None,
+                distinct: false,
+                special: false,
+                order_by: vec![],
+            }),
+            value_column: vec![Ident::new("year")],
+            pivot_values: vec![
+                Value::SingleQuotedString("population_2000".to_string()),
+                Value::SingleQuotedString("population_2010".to_string())
+            ],
+            alias: Some(TableAlias {
+                name: Ident::new("p"),
+                columns: vec![]
+            }),
+        }
+    );
+    assert_eq!(verified_stmt(sql).to_string(), sql);
+}
+
 /// Makes a predicate that looks like ((user_id = $id) OR user_id = $2...)
 fn make_where_clause(num: usize) -> String {
     use std::fmt::Write;