diff --git a/src/ast/mod.rs b/src/ast/mod.rs index d937b7275..1a1a2d42d 100644 --- a/src/ast/mod.rs +++ b/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, diff --git a/src/ast/query.rs b/src/ast/query.rs index 08a0bc5af..d43692059 100644 --- a/src/ast/query.rs +++ b/src/ast/query.rs @@ -946,7 +946,8 @@ pub enum TableFactor { table: Box, aggregate_functions: Vec, // Function expression value_column: Vec, - pivot_values: Vec, + value_source: PivotValueSource, + default_on_null: Option, alias: Option, }, /// An UNPIVOT operation on a table. @@ -987,6 +988,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 . + List(Vec), + /// Pivot on all distinct values of the pivot column. + /// + /// See . + Any(Vec), + /// Pivot on all values returned by a subquery. + /// + /// See . + 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 . @@ -1313,17 +1349,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())?; } diff --git a/src/parser/mod.rs b/src/parser/mod.rs index a7ec4d093..61d41c281 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -9068,16 +9068,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, }) } diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 6668ce8f4..7aa69e979 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -8612,7 +8612,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")) @@ -8625,7 +8625,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(), @@ -8763,7 +8764,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 @@ -8772,7 +8773,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![] diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs index 30f2cc601..e1dba252d 100644 --- a/tests/sqlparser_snowflake.rs +++ b/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", + )); +}