feat: support multiple value for pivot (apache#1970)

Author: chenkovsky

src/ast/query.rs

@@ -1342,7 +1342,7 @@ pub enum TableFactor {
     Pivot {
         table: Box<TableFactor>,
         aggregate_functions: Vec<ExprWithAlias>, // Function expression
-        value_column: Vec<Ident>,
+        value_column: Vec<Expr>,
         value_source: PivotValueSource,
         default_on_null: Option<Expr>,
         alias: Option<TableAlias>,
@@ -2024,10 +2024,15 @@ impl fmt::Display for TableFactor {
             } => {
                 write!(
                     f,
-                    "{table} PIVOT({} FOR {} IN ({value_source})",
+                    "{table} PIVOT({} FOR ",
                     display_comma_separated(aggregate_functions),
-                    Expr::CompoundIdentifier(value_column.to_vec()),
                 )?;
+                if value_column.len() == 1 {
+                    write!(f, "{}", value_column[0])?;
+                } else {
+                    write!(f, "({})", display_comma_separated(value_column))?;
+                }
+                write!(f, " IN ({value_source})")?;
                 if let Some(expr) = default_on_null {
                     write!(f, " DEFAULT ON NULL ({expr})")?;
                 }

src/ast/spans.rs

@@ -2010,7 +2010,7 @@ impl Spanned for TableFactor {
             } => union_spans(
                 core::iter::once(table.span())
                     .chain(aggregate_functions.iter().map(|i| i.span()))
-                    .chain(value_column.iter().map(|i| i.span))
+                    .chain(value_column.iter().map(|i| i.span()))
                     .chain(core::iter::once(value_source.span()))
                     .chain(default_on_null.as_ref().map(|i| i.span()))
                     .chain(alias.as_ref().map(|i| i.span())),

src/ast/visitor.rs

@@ -884,6 +884,8 @@ mod tests {
                     "PRE: EXPR: a.amount",
                     "POST: EXPR: a.amount",
                     "POST: EXPR: SUM(a.amount)",
+                    "PRE: EXPR: a.MONTH",
+                    "POST: EXPR: a.MONTH",
                     "PRE: EXPR: 'JAN'",
                     "POST: EXPR: 'JAN'",
                     "PRE: EXPR: 'FEB'",

src/parser/mod.rs

@@ -11098,6 +11098,18 @@ impl<'a> Parser<'a> {
         self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| p.parse_identifier())
     }
 
+    pub fn parse_parenthesized_compound_identifier_list(
+        &mut self,
+        optional: IsOptional,
+        allow_empty: bool,
+    ) -> Result<Vec<Expr>, ParserError> {
+        self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| {
+            Ok(Expr::CompoundIdentifier(
+                p.parse_period_separated(|p| p.parse_identifier())?,
+            ))
+        })
+    }
+
     /// Parses a parenthesized comma-separated list of index columns, which can be arbitrary
     /// expressions with ordering information (and an opclass in some dialects).
     fn parse_parenthesized_index_column_list(&mut self) -> Result<Vec<IndexColumn>, ParserError> {
@@ -14219,7 +14231,13 @@ impl<'a> Parser<'a> {
         self.expect_token(&Token::LParen)?;
         let aggregate_functions = self.parse_comma_separated(Self::parse_aliased_function_call)?;
         self.expect_keyword_is(Keyword::FOR)?;
-        let value_column = self.parse_period_separated(|p| p.parse_identifier())?;
+        let value_column = if self.peek_token_ref().token == Token::LParen {
+            self.parse_parenthesized_column_list_inner(Mandatory, false, |p| {
+                p.parse_subexpr(self.dialect.prec_value(Precedence::Between))
+            })?
+        } else {
+            vec![self.parse_subexpr(self.dialect.prec_value(Precedence::Between))?]
+        };
         self.expect_keyword_is(Keyword::IN)?;
 
         self.expect_token(&Token::LParen)?;

tests/sqlparser_common.rs

@@ -10964,7 +10964,10 @@ fn parse_pivot_table() {
                 expected_function("b", Some("t")),
                 expected_function("c", Some("u")),
             ],
-            value_column: vec![Ident::new("a"), Ident::new("MONTH")],
+            value_column: vec![Expr::CompoundIdentifier(vec![
+                Ident::new("a"),
+                Ident::new("MONTH")
+            ])],
             value_source: PivotValueSource::List(vec![
                 ExprWithAlias {
                     expr: Expr::value(number("1")),
@@ -11011,6 +11014,75 @@ fn parse_pivot_table() {
         verified_stmt(sql_without_table_alias).to_string(),
         sql_without_table_alias
     );
+
+    let multiple_value_columns_sql = concat!(
+        "SELECT * FROM person ",
+        "PIVOT(",
+        "SUM(age) AS a, AVG(class) AS c ",
+        "FOR (name, age) IN (('John', 30) AS c1, ('Mike', 40) AS c2))",
+    );
+
+    assert_eq!(
+        verified_only_select(multiple_value_columns_sql).from[0].relation,
+        Pivot {
+            table: Box::new(TableFactor::Table {
+                name: ObjectName::from(vec![Ident::new("person")]),
+                alias: None,
+                args: None,
+                with_hints: vec![],
+                version: None,
+                partitions: vec![],
+                with_ordinality: false,
+                json_path: None,
+                sample: None,
+                index_hints: vec![],
+            }),
+            aggregate_functions: vec![
+                ExprWithAlias {
+                    expr: call("SUM", [Expr::Identifier(Ident::new("age"))]),
+                    alias: Some(Ident::new("a"))
+                },
+                ExprWithAlias {
+                    expr: call("AVG", [Expr::Identifier(Ident::new("class"))]),
+                    alias: Some(Ident::new("c"))
+                },
+            ],
+            value_column: vec![
+                Expr::Identifier(Ident::new("name")),
+                Expr::Identifier(Ident::new("age")),
+            ],
+            value_source: PivotValueSource::List(vec![
+                ExprWithAlias {
+                    expr: Expr::Tuple(vec![
+                        Expr::Value(
+                            (Value::SingleQuotedString("John".to_string())).with_empty_span()
+                        ),
+                        Expr::Value(
+                            (Value::Number("30".parse().unwrap(), false)).with_empty_span()
+                        ),
+                    ]),
+                    alias: Some(Ident::new("c1"))
+                },
+                ExprWithAlias {
+                    expr: Expr::Tuple(vec![
+                        Expr::Value(
+                            (Value::SingleQuotedString("Mike".to_string())).with_empty_span()
+                        ),
+                        Expr::Value(
+                            (Value::Number("40".parse().unwrap(), false)).with_empty_span()
+                        ),
+                    ]),
+                    alias: Some(Ident::new("c2"))
+                },
+            ]),
+            default_on_null: None,
+            alias: None,
+        }
+    );
+    assert_eq!(
+        verified_stmt(multiple_value_columns_sql).to_string(),
+        multiple_value_columns_sql
+    );
 }
 
 #[test]
@@ -11343,7 +11415,7 @@ fn parse_pivot_unpivot_table() {
                 expr: call("sum", [Expr::Identifier(Ident::new("population"))]),
                 alias: None
             }],
-            value_column: vec![Ident::new("year")],
+            value_column: vec![Expr::Identifier(Ident::new("year"))],
             value_source: PivotValueSource::List(vec![
                 ExprWithAlias {
                     expr: Expr::Value(