Support for Snowflake dynamic pivot (apache#1280)

Author: jmhain

src/ast/mod.rs

@@ -46,10 +46,11 @@ pub use self::query::{
     GroupByExpr, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator,
     JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
     MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
-    NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem, RepetitionQuantifier,
-    ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select, SelectInto, SelectItem, SetExpr,
-    SetOperator, SetQuantifier, SymbolDefinition, Table, TableAlias, TableFactor, TableVersion,
-    TableWithJoins, Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With,
+    NonBlock, Offset, OffsetRows, OrderByExpr, PivotValueSource, Query, RenameSelectItem,
+    RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
+    SelectInto, SelectItem, SetExpr, SetOperator, SetQuantifier, SymbolDefinition, Table,
+    TableAlias, TableFactor, TableVersion, TableWithJoins, Top, TopQuantity, ValueTableMode,
+    Values, WildcardAdditionalOptions, With,
 };
 pub use self::value::{
     escape_double_quote_string, escape_quoted_string, DateTimeField, DollarQuotedString,

src/ast/query.rs

@@ -957,7 +957,8 @@ pub enum TableFactor {
         table: Box<TableFactor>,
         aggregate_functions: Vec<ExprWithAlias>, // Function expression
         value_column: Vec<Ident>,
-        pivot_values: Vec<ExprWithAlias>,
+        value_source: PivotValueSource,
+        default_on_null: Option<Expr>,
         alias: Option<TableAlias>,
     },
     /// An UNPIVOT operation on a table.
@@ -998,6 +999,41 @@ pub enum TableFactor {
     },
 }
 
+/// The source of values in a `PIVOT` operation.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum PivotValueSource {
+    /// Pivot on a static list of values.
+    ///
+    /// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot#pivot-on-a-specified-list-of-column-values-for-the-pivot-column>.
+    List(Vec<ExprWithAlias>),
+    /// Pivot on all distinct values of the pivot column.
+    ///
+    /// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot#pivot-on-all-distinct-column-values-automatically-with-dynamic-pivot>.
+    Any(Vec<OrderByExpr>),
+    /// Pivot on all values returned by a subquery.
+    ///
+    /// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot#pivot-on-column-values-using-a-subquery-with-dynamic-pivot>.
+    Subquery(Query),
+}
+
+impl fmt::Display for PivotValueSource {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            PivotValueSource::List(values) => write!(f, "{}", display_comma_separated(values)),
+            PivotValueSource::Any(order_by) => {
+                write!(f, "ANY")?;
+                if !order_by.is_empty() {
+                    write!(f, " ORDER BY {}", display_comma_separated(order_by))?;
+                }
+                Ok(())
+            }
+            PivotValueSource::Subquery(query) => write!(f, "{query}"),
+        }
+    }
+}
+
 /// An item in the `MEASURES` subclause of a `MATCH_RECOGNIZE` operation.
 ///
 /// See <https://docs.snowflake.com/en/sql-reference/constructs/match_recognize#measures-specifying-additional-output-columns>.
@@ -1324,17 +1360,20 @@ impl fmt::Display for TableFactor {
                 table,
                 aggregate_functions,
                 value_column,
-                pivot_values,
+                value_source,
+                default_on_null,
                 alias,
             } => {
                 write!(
                     f,
-                    "{} PIVOT({} FOR {} IN ({}))",
-                    table,
+                    "{table} PIVOT({} FOR {} IN ({value_source})",
                     display_comma_separated(aggregate_functions),
                     Expr::CompoundIdentifier(value_column.to_vec()),
-                    display_comma_separated(pivot_values)
                 )?;
+                if let Some(expr) = default_on_null {
+                    write!(f, " DEFAULT ON NULL ({expr})")?;
+                }
+                write!(f, ")")?;
                 if alias.is_some() {
                     write!(f, " AS {}", alias.as_ref().unwrap())?;
                 }

src/parser/mod.rs

@@ -9191,16 +9191,44 @@ impl<'a> Parser<'a> {
         self.expect_keyword(Keyword::FOR)?;
         let value_column = self.parse_object_name(false)?.0;
         self.expect_keyword(Keyword::IN)?;
+
         self.expect_token(&Token::LParen)?;
-        let pivot_values = self.parse_comma_separated(Self::parse_expr_with_alias)?;
+        let value_source = if self.parse_keyword(Keyword::ANY) {
+            let order_by = if self.parse_keywords(&[Keyword::ORDER, Keyword::BY]) {
+                self.parse_comma_separated(Parser::parse_order_by_expr)?
+            } else {
+                vec![]
+            };
+            PivotValueSource::Any(order_by)
+        } else if self
+            .parse_one_of_keywords(&[Keyword::SELECT, Keyword::WITH])
+            .is_some()
+        {
+            self.prev_token();
+            PivotValueSource::Subquery(self.parse_query()?)
+        } else {
+            PivotValueSource::List(self.parse_comma_separated(Self::parse_expr_with_alias)?)
+        };
         self.expect_token(&Token::RParen)?;
+
+        let default_on_null =
+            if self.parse_keywords(&[Keyword::DEFAULT, Keyword::ON, Keyword::NULL]) {
+                self.expect_token(&Token::LParen)?;
+                let expr = self.parse_expr()?;
+                self.expect_token(&Token::RParen)?;
+                Some(expr)
+            } else {
+                None
+            };
+
         self.expect_token(&Token::RParen)?;
         let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
         Ok(TableFactor::Pivot {
             table: Box::new(table),
             aggregate_functions,
             value_column,
-            pivot_values,
+            value_source,
+            default_on_null,
             alias,
         })
     }

tests/sqlparser_common.rs

@@ -8618,7 +8618,7 @@ fn parse_pivot_table() {
                 expected_function("c", Some("u")),
             ],
             value_column: vec![Ident::new("a"), Ident::new("MONTH")],
-            pivot_values: vec![
+            value_source: PivotValueSource::List(vec![
                 ExprWithAlias {
                     expr: Expr::Value(number("1")),
                     alias: Some(Ident::new("x"))
@@ -8631,7 +8631,8 @@ fn parse_pivot_table() {
                     expr: Expr::Identifier(Ident::new("three")),
                     alias: Some(Ident::new("y"))
                 },
-            ],
+            ]),
+            default_on_null: None,
             alias: Some(TableAlias {
                 name: Ident {
                     value: "p".to_string(),
@@ -8769,7 +8770,7 @@ fn parse_pivot_unpivot_table() {
                 alias: None
             }],
             value_column: vec![Ident::new("year")],
-            pivot_values: vec![
+            value_source: PivotValueSource::List(vec![
                 ExprWithAlias {
                     expr: Expr::Value(Value::SingleQuotedString("population_2000".to_string())),
                     alias: None
@@ -8778,7 +8779,8 @@ fn parse_pivot_unpivot_table() {
                     expr: Expr::Value(Value::SingleQuotedString("population_2010".to_string())),
                     alias: None
                 },
-            ],
+            ]),
+            default_on_null: None,
             alias: Some(TableAlias {
                 name: Ident::new("p"),
                 columns: vec![]

tests/sqlparser_snowflake.rs

@@ -1585,3 +1585,56 @@ fn first_value_ignore_nulls() {
         "FROM some_table"
     ));
 }
+
+#[test]
+fn test_pivot() {
+    // pivot on static list of values with default
+    #[rustfmt::skip]
+    snowflake().verified_only_select(concat!(
+        "SELECT * ",
+        "FROM quarterly_sales ",
+          "PIVOT(SUM(amount) ",
+            "FOR quarter IN (",
+              "'2023_Q1', ",
+              "'2023_Q2', ",
+              "'2023_Q3', ",
+              "'2023_Q4', ",
+              "'2024_Q1') ",
+            "DEFAULT ON NULL (0)",
+          ") ",
+        "ORDER BY empid",
+    ));
+
+    // dynamic pivot from subquery
+    #[rustfmt::skip]
+    snowflake().verified_only_select(concat!(
+        "SELECT * ",
+        "FROM quarterly_sales ",
+          "PIVOT(SUM(amount) FOR quarter IN (",
+            "SELECT DISTINCT quarter ",
+              "FROM ad_campaign_types_by_quarter ",
+              "WHERE television = true ",
+              "ORDER BY quarter)",
+          ") ",
+        "ORDER BY empid",
+    ));
+
+    // dynamic pivot on any value (with order by)
+    #[rustfmt::skip]
+    snowflake().verified_only_select(concat!(
+        "SELECT * ",
+        "FROM quarterly_sales ",
+          "PIVOT(SUM(amount) FOR quarter IN (ANY ORDER BY quarter)) ",
+        "ORDER BY empid",
+    ));
+
+    // dynamic pivot on any value (without order by)
+    #[rustfmt::skip]
+    snowflake().verified_only_select(concat!(
+        "SELECT * ",
+        "FROM sales_data ",
+          "PIVOT(SUM(total_sales) FOR fis_quarter IN (ANY)) ",
+        "WHERE fis_year IN (2023) ",
+        "ORDER BY region",
+    ));
+}