Extend pivot operator support (#1238)

Author: iffyio

src/ast/mod.rs

@@ -41,14 +41,14 @@ pub use self::dml::{Delete, Insert};
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
-    ExceptSelectItem, ExcludeSelectItem, Fetch, ForClause, ForJson, ForXml, 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,
+    ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
+    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,
 };
 pub use self::value::{
     escape_double_quote_string, escape_quoted_string, DateTimeField, DollarQuotedString,

src/ast/query.rs

@@ -802,6 +802,31 @@ impl fmt::Display for ConnectBy {
     }
 }
 
+/// An expression optionally followed by an alias.
+///
+/// Example:
+/// ```sql
+/// 42 AS myint
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ExprWithAlias {
+    pub expr: Expr,
+    pub alias: Option<Ident>,
+}
+
+impl fmt::Display for ExprWithAlias {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let ExprWithAlias { expr, alias } = self;
+        write!(f, "{expr}")?;
+        if let Some(alias) = alias {
+            write!(f, " AS {alias}")?;
+        }
+        Ok(())
+    }
+}
+
 /// A table name or a parenthesized subquery with an optional alias
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -900,12 +925,14 @@ pub enum TableFactor {
     },
     /// Represents PIVOT operation on a table.
     /// For example `FROM monthly_sales PIVOT(sum(amount) FOR MONTH IN ('JAN', 'FEB'))`
-    /// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot>
+    ///
+    /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#pivot_operator)
+    /// [Snowflake](https://docs.snowflake.com/en/sql-reference/constructs/pivot)
     Pivot {
         table: Box<TableFactor>,
-        aggregate_function: Expr, // Function expression
+        aggregate_functions: Vec<ExprWithAlias>, // Function expression
         value_column: Vec<Ident>,
-        pivot_values: Vec<Value>,
+        pivot_values: Vec<ExprWithAlias>,
         alias: Option<TableAlias>,
     },
     /// An UNPIVOT operation on a table.
@@ -1270,7 +1297,7 @@ impl fmt::Display for TableFactor {
             }
             TableFactor::Pivot {
                 table,
-                aggregate_function,
+                aggregate_functions,
                 value_column,
                 pivot_values,
                 alias,
@@ -1279,7 +1306,7 @@ impl fmt::Display for TableFactor {
                     f,
                     "{} PIVOT({} FOR {} IN ({}))",
                     table,
-                    aggregate_function,
+                    display_comma_separated(aggregate_functions),
                     Expr::CompoundIdentifier(value_column.to_vec()),
                     display_comma_separated(pivot_values)
                 )?;

src/ast/visitor.rs

@@ -850,6 +850,14 @@ mod tests {
                     "PRE: EXPR: a.amount",
                     "POST: EXPR: a.amount",
                     "POST: EXPR: SUM(a.amount)",
+                    "PRE: EXPR: 'JAN'",
+                    "POST: EXPR: 'JAN'",
+                    "PRE: EXPR: 'FEB'",
+                    "POST: EXPR: 'FEB'",
+                    "PRE: EXPR: 'MAR'",
+                    "POST: EXPR: 'MAR'",
+                    "PRE: EXPR: 'APR'",
+                    "POST: EXPR: 'APR'",
                     "POST: TABLE FACTOR: monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
                     "PRE: EXPR: EMPID",
                     "POST: EXPR: EMPID",

src/parser/mod.rs

@@ -8775,27 +8775,49 @@ impl<'a> Parser<'a> {
         })
     }
 
+    fn parse_aliased_function_call(&mut self) -> Result<ExprWithAlias, ParserError> {
+        let function_name = match self.next_token().token {
+            Token::Word(w) => Ok(w.value),
+            _ => self.expected("a function identifier", self.peek_token()),
+        }?;
+        let expr = self.parse_function(ObjectName(vec![Ident::new(function_name)]))?;
+        let alias = if self.parse_keyword(Keyword::AS) {
+            Some(self.parse_identifier(false)?)
+        } else {
+            None
+        };
+
+        Ok(ExprWithAlias { expr, alias })
+    }
+
+    fn parse_expr_with_alias(&mut self) -> Result<ExprWithAlias, ParserError> {
+        let expr = self.parse_expr()?;
+        let alias = if self.parse_keyword(Keyword::AS) {
+            Some(self.parse_identifier(false)?)
+        } else {
+            None
+        };
+
+        Ok(ExprWithAlias { expr, alias })
+    }
+
     pub fn parse_pivot_table_factor(
         &mut self,
         table: TableFactor,
     ) -> Result<TableFactor, ParserError> {
         self.expect_token(&Token::LParen)?;
-        let function_name = match self.next_token().token {
-            Token::Word(w) => Ok(w.value),
-            _ => self.expected("an aggregate function name", self.peek_token()),
-        }?;
-        let function = self.parse_function(ObjectName(vec![Ident::new(function_name)]))?;
+        let aggregate_functions = self.parse_comma_separated(Self::parse_aliased_function_call)?;
         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(Parser::parse_value)?;
+        let pivot_values = self.parse_comma_separated(Self::parse_expr_with_alias)?;
         self.expect_token(&Token::RParen)?;
         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_function: function,
+            aggregate_functions,
             value_column,
             pivot_values,
             alias,

tests/sqlparser_common.rs

@@ -8453,11 +8453,32 @@ fn parse_escaped_string_without_unescape() {
 #[test]
 fn parse_pivot_table() {
     let sql = concat!(
-        "SELECT * FROM monthly_sales AS a ",
-        "PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
+        "SELECT * FROM monthly_sales AS a PIVOT(",
+        "SUM(a.amount), ",
+        "SUM(b.amount) AS t, ",
+        "SUM(c.amount) AS u ",
+        "FOR a.MONTH IN (1 AS x, 'two', three AS y)) AS p (c, d) ",
         "ORDER BY EMPID"
     );
 
+    fn expected_function(table: &'static str, alias: Option<&'static str>) -> ExprWithAlias {
+        ExprWithAlias {
+            expr: Expr::Function(Function {
+                name: ObjectName(vec![Ident::new("SUM")]),
+                args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                    Expr::CompoundIdentifier(vec![Ident::new(table), Ident::new("amount")]),
+                ))]),
+                null_treatment: None,
+                filter: None,
+                over: None,
+                distinct: false,
+                special: false,
+                order_by: vec![],
+            }),
+            alias: alias.map(Ident::new),
+        }
+    }
+
     assert_eq!(
         verified_only_select(sql).from[0].relation,
         Pivot {
@@ -8472,24 +8493,25 @@ fn parse_pivot_table() {
                 version: None,
                 partitions: vec![],
             }),
-            aggregate_function: Expr::Function(Function {
-                name: ObjectName(vec![Ident::new("SUM")]),
-                args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
-                    Expr::CompoundIdentifier(vec![Ident::new("a"), Ident::new("amount"),])
-                ))]),
-                null_treatment: None,
-                filter: None,
-                over: None,
-                distinct: false,
-                special: false,
-                order_by: vec![],
-            }),
+            aggregate_functions: vec![
+                expected_function("a", None),
+                expected_function("b", Some("t")),
+                expected_function("c", Some("u")),
+            ],
             value_column: vec![Ident::new("a"), Ident::new("MONTH")],
             pivot_values: vec![
-                Value::SingleQuotedString("JAN".to_string()),
-                Value::SingleQuotedString("FEB".to_string()),
-                Value::SingleQuotedString("MAR".to_string()),
-                Value::SingleQuotedString("APR".to_string()),
+                ExprWithAlias {
+                    expr: Expr::Value(number("1")),
+                    alias: Some(Ident::new("x"))
+                },
+                ExprWithAlias {
+                    expr: Expr::Value(Value::SingleQuotedString("two".to_string())),
+                    alias: None
+                },
+                ExprWithAlias {
+                    expr: Expr::Identifier(Ident::new("three")),
+                    alias: Some(Ident::new("y"))
+                },
             ],
             alias: Some(TableAlias {
                 name: Ident {
@@ -8623,22 +8645,31 @@ fn parse_pivot_unpivot_table() {
                     columns: vec![]
                 }),
             }),
-            aggregate_function: Expr::Function(Function {
-                name: ObjectName(vec![Ident::new("sum")]),
-                args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
-                    Expr::Identifier(Ident::new("population"))
-                ))]),
-                null_treatment: None,
-                filter: None,
-                over: None,
-                distinct: false,
-                special: false,
-                order_by: vec![],
-            }),
+            aggregate_functions: vec![ExprWithAlias {
+                expr: Expr::Function(Function {
+                    name: ObjectName(vec![Ident::new("sum")]),
+                    args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
+                        Expr::Identifier(Ident::new("population"))
+                    ))]),
+                    null_treatment: None,
+                    filter: None,
+                    over: None,
+                    distinct: false,
+                    special: false,
+                    order_by: vec![],
+                }),
+                alias: None
+            }],
             value_column: vec![Ident::new("year")],
             pivot_values: vec![
-                Value::SingleQuotedString("population_2000".to_string()),
-                Value::SingleQuotedString("population_2010".to_string())
+                ExprWithAlias {
+                    expr: Expr::Value(Value::SingleQuotedString("population_2000".to_string())),
+                    alias: None
+                },
+                ExprWithAlias {
+                    expr: Expr::Value(Value::SingleQuotedString("population_2010".to_string())),
+                    alias: None
+                },
             ],
             alias: Some(TableAlias {
                 name: Ident::new("p"),