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

apache · Oct 2, 2023 · 40e2ecb · 40e2ecb
1 parent 2786c7e
commit 40e2ecb
Show file tree

Hide file tree

Showing 4 changed files with 231 additions and 47 deletions.
diff --git a/src/ast/query.rs b/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"))]
-        name: ObjectName,
-        table_alias: Option<TableAlias>,
+        #[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
+        table: Box<TableFactor>,
         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_alias,
+                table,
                 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 {
-                    write!(f, "{}", value)?;
-                    if !value.eq(pivot_values.last().unwrap()) {
-                        write!(f, ", ")?;
-                    }
+                if alias.is_some() {
+                    write!(f, " AS {}", alias.as_ref().unwrap())?;
                 }
-                write!(f, "))")?;
-                if pivot_alias.is_some() {
-                    write!(f, " AS {}", pivot_alias.as_ref().unwrap())?;
+                Ok(())
+            }
+            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(())
             }

diff --git a/src/keywords.rs b/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,

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -6276,9 +6276,8 @@ impl<'a> Parser<'a> {
                         | TableFactor::Table { alias, .. }
                         | TableFactor::UNNEST { alias, .. }
                         | TableFactor::TableFunction { alias, .. }
-                        | TableFactor::Pivot {
-                            pivot_alias: alias, ..
-                        }
+                        | TableFactor::Pivot { 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_alias: Option<TableAlias>,
+        table: TableFactor,
     ) -> 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_alias,
+            table: Box::new(table),
             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,
         })
     }