feat: support multiple value for pivot

This commit is contained in:
Chongchen Chen 2025-07-23 21:52:16 +08:00
parent 492184643a
commit 3fade60c96
4 changed files with 32 additions and 7 deletions

View file

@ -1336,7 +1336,7 @@ pub enum TableFactor {
Pivot { Pivot {
table: Box<TableFactor>, table: Box<TableFactor>,
aggregate_functions: Vec<ExprWithAlias>, // Function expression aggregate_functions: Vec<ExprWithAlias>, // Function expression
value_column: Vec<Ident>, value_column: Vec<Expr>, // Expr is a identifier or a compound identifier
value_source: PivotValueSource, value_source: PivotValueSource,
default_on_null: Option<Expr>, default_on_null: Option<Expr>,
alias: Option<TableAlias>, alias: Option<TableAlias>,
@ -2010,10 +2010,15 @@ impl fmt::Display for TableFactor {
} => { } => {
write!( write!(
f, f,
"{table} PIVOT({} FOR {} IN ({value_source})", "{table} PIVOT({} FOR ",
display_comma_separated(aggregate_functions), 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 { if let Some(expr) = default_on_null {
write!(f, " DEFAULT ON NULL ({expr})")?; write!(f, " DEFAULT ON NULL ({expr})")?;
} }

View file

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

View file

@ -10820,6 +10820,16 @@ impl<'a> Parser<'a> {
self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| p.parse_identifier()) 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 /// Parses a parenthesized comma-separated list of index columns, which can be arbitrary
/// expressions with ordering information (and an opclass in some dialects). /// expressions with ordering information (and an opclass in some dialects).
fn parse_parenthesized_index_column_list(&mut self) -> Result<Vec<IndexColumn>, ParserError> { fn parse_parenthesized_index_column_list(&mut self) -> Result<Vec<IndexColumn>, ParserError> {
@ -13828,7 +13838,11 @@ impl<'a> Parser<'a> {
self.expect_token(&Token::LParen)?; self.expect_token(&Token::LParen)?;
let aggregate_functions = self.parse_comma_separated(Self::parse_aliased_function_call)?; let aggregate_functions = self.parse_comma_separated(Self::parse_aliased_function_call)?;
self.expect_keyword_is(Keyword::FOR)?; self.expect_keyword_is(Keyword::FOR)?;
let value_column = self.parse_period_separated(|p| p.parse_identifier())?; let value_column = if self.peek_token().token == Token::LParen {
self.parse_parenthesized_compound_identifier_list(Mandatory, false)?
} else {
vec![Expr::CompoundIdentifier(self.parse_period_separated(|p| p.parse_identifier())?)]
};
self.expect_keyword_is(Keyword::IN)?; self.expect_keyword_is(Keyword::IN)?;
self.expect_token(&Token::LParen)?; self.expect_token(&Token::LParen)?;

View file

@ -10875,7 +10875,7 @@ fn parse_pivot_table() {
expected_function("b", Some("t")), expected_function("b", Some("t")),
expected_function("c", Some("u")), 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![ value_source: PivotValueSource::List(vec![
ExprWithAlias { ExprWithAlias {
expr: Expr::value(number("1")), expr: Expr::value(number("1")),
@ -10922,6 +10922,12 @@ fn parse_pivot_table() {
verified_stmt(sql_without_table_alias).to_string(), verified_stmt(sql_without_table_alias).to_string(),
sql_without_table_alias sql_without_table_alias
); );
let sql_with_multiple_value_column = 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_stmt(sql_with_multiple_value_column).to_string(), sql_with_multiple_value_column);
} }
#[test] #[test]
@ -11143,7 +11149,7 @@ fn parse_pivot_unpivot_table() {
expr: call("sum", [Expr::Identifier(Ident::new("population"))]), expr: call("sum", [Expr::Identifier(Ident::new("population"))]),
alias: None alias: None
}], }],
value_column: vec![Ident::new("year")], value_column: vec![Expr::CompoundIdentifier(vec![Ident::new("year")])],
value_source: PivotValueSource::List(vec![ value_source: PivotValueSource::List(vec![
ExprWithAlias { ExprWithAlias {
expr: Expr::Value( expr: Expr::Value(