Skip to content

Commit dd650b8

Browse files
chenkovskychenkovsky
authored
feat: support multi value columns and aliases in unpivot (apache#1969)
1 parent ec0026d commit dd650b8

File tree

4 files changed

+144
-30
lines changed

4 files changed

+144
-30
lines changed

‎src/ast/query.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1349,11 +1349,12 @@ pub enum TableFactor {
13491349
/// ```
13501350
///
13511351
/// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
1352+
/// See <https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-qry-select-unpivot>.
13521353
Unpivot {
13531354
table: Box<TableFactor>,
1354-
value: Ident,
1355+
value: Expr,
13551356
name: Ident,
1356-
columns: Vec<Ident>,
1357+
columns: Vec<ExprWithAlias>,
13571358
null_inclusion: Option<NullInclusion>,
13581359
alias: Option<TableAlias>,
13591360
},

‎src/ast/spans.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2002,9 +2002,9 @@ impl Spanned for TableFactor {
20022002
alias,
20032003
} => union_spans(
20042004
core::iter::once(table.span())
2005-
.chain(core::iter::once(value.span))
2005+
.chain(core::iter::once(value.span()))
20062006
.chain(core::iter::once(name.span))
2007-
.chain(columns.iter().map(|i| i.span))
2007+
.chain(columns.iter().map(|ilist| ilist.span()))
20082008
.chain(alias.as_ref().map(|alias| alias.span())),
20092009
),
20102010
TableFactor::MatchRecognize {

‎src/parser/mod.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14033,11 +14033,13 @@ impl<'a> Parser<'a> {
1403314033
None
1403414034
};
1403514035
self.expect_token(&Token::LParen)?;
14036-
let value = self.parse_identifier()?;
14036+
let value = self.parse_expr()?;
1403714037
self.expect_keyword_is(Keyword::FOR)?;
1403814038
let name = self.parse_identifier()?;
1403914039
self.expect_keyword_is(Keyword::IN)?;
14040-
let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
14040+
let columns = self.parse_parenthesized_column_list_inner(Mandatory, false, |p| {
14041+
p.parse_expr_with_alias()
14042+
})?;
1404114043
self.expect_token(&Token::RParen)?;
1404214044
let alias = self.maybe_parse_table_alias()?;
1404314045
Ok(TableFactor::Unpivot {

‎tests/sqlparser_common.rs

Lines changed: 135 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -11009,20 +11009,14 @@ fn parse_unpivot_table() {
1100911009
index_hints: vec![],
1101011010
}),
1101111011
null_inclusion: None,
11012-
value: Ident {
11013-
value: "quantity".to_string(),
11014-
quote_style: None,
11015-
span: Span::empty(),
11016-
},
11017-
11018-
name: Ident {
11019-
value: "quarter".to_string(),
11020-
quote_style: None,
11021-
span: Span::empty(),
11022-
},
11012+
value: Expr::Identifier(Ident::new("quantity")),
11013+
name: Ident::new("quarter"),
1102311014
columns: ["Q1", "Q2", "Q3", "Q4"]
1102411015
.into_iter()
11025-
.map(Ident::new)
11016+
.map(|col| ExprWithAlias {
11017+
expr: Expr::Identifier(Ident::new(col)),
11018+
alias: None,
11019+
})
1102611020
.collect(),
1102711021
alias: Some(TableAlias {
1102811022
name: Ident::new("u"),
@@ -11084,6 +11078,129 @@ fn parse_unpivot_table() {
1108411078
verified_stmt(sql_unpivot_include_nulls).to_string(),
1108511079
sql_unpivot_include_nulls
1108611080
);
11081+
11082+
let sql_unpivot_with_alias = concat!(
11083+
"SELECT * FROM sales AS s ",
11084+
"UNPIVOT INCLUDE NULLS ",
11085+
"(quantity FOR quarter IN ",
11086+
"(Q1 AS Quater1, Q2 AS Quater2, Q3 AS Quater3, Q4 AS Quater4)) ",
11087+
"AS u (product, quarter, quantity)"
11088+
);
11089+
11090+
if let Unpivot { value, columns, .. } =
11091+
&verified_only_select(sql_unpivot_with_alias).from[0].relation
11092+
{
11093+
assert_eq!(
11094+
*columns,
11095+
vec![
11096+
ExprWithAlias {
11097+
expr: Expr::Identifier(Ident::new("Q1")),
11098+
alias: Some(Ident::new("Quater1")),
11099+
},
11100+
ExprWithAlias {
11101+
expr: Expr::Identifier(Ident::new("Q2")),
11102+
alias: Some(Ident::new("Quater2")),
11103+
},
11104+
ExprWithAlias {
11105+
expr: Expr::Identifier(Ident::new("Q3")),
11106+
alias: Some(Ident::new("Quater3")),
11107+
},
11108+
ExprWithAlias {
11109+
expr: Expr::Identifier(Ident::new("Q4")),
11110+
alias: Some(Ident::new("Quater4")),
11111+
},
11112+
]
11113+
);
11114+
assert_eq!(*value, Expr::Identifier(Ident::new("quantity")));
11115+
}
11116+
11117+
assert_eq!(
11118+
verified_stmt(sql_unpivot_with_alias).to_string(),
11119+
sql_unpivot_with_alias
11120+
);
11121+
11122+
let sql_unpivot_with_alias_and_multi_value = concat!(
11123+
"SELECT * FROM sales AS s ",
11124+
"UNPIVOT INCLUDE NULLS ((first_quarter, second_quarter) ",
11125+
"FOR half_of_the_year IN (",
11126+
"(Q1, Q2) AS H1, ",
11127+
"(Q3, Q4) AS H2",
11128+
"))"
11129+
);
11130+
11131+
if let Unpivot { value, columns, .. } =
11132+
&verified_only_select(sql_unpivot_with_alias_and_multi_value).from[0].relation
11133+
{
11134+
assert_eq!(
11135+
*columns,
11136+
vec![
11137+
ExprWithAlias {
11138+
expr: Expr::Tuple(vec![
11139+
Expr::Identifier(Ident::new("Q1")),
11140+
Expr::Identifier(Ident::new("Q2")),
11141+
]),
11142+
alias: Some(Ident::new("H1")),
11143+
},
11144+
ExprWithAlias {
11145+
expr: Expr::Tuple(vec![
11146+
Expr::Identifier(Ident::new("Q3")),
11147+
Expr::Identifier(Ident::new("Q4")),
11148+
]),
11149+
alias: Some(Ident::new("H2")),
11150+
},
11151+
]
11152+
);
11153+
assert_eq!(
11154+
*value,
11155+
Expr::Tuple(vec![
11156+
Expr::Identifier(Ident::new("first_quarter")),
11157+
Expr::Identifier(Ident::new("second_quarter")),
11158+
])
11159+
);
11160+
}
11161+
11162+
assert_eq!(
11163+
verified_stmt(sql_unpivot_with_alias_and_multi_value).to_string(),
11164+
sql_unpivot_with_alias_and_multi_value
11165+
);
11166+
11167+
let sql_unpivot_with_alias_and_multi_value_and_qualifier = concat!(
11168+
"SELECT * FROM sales AS s ",
11169+
"UNPIVOT INCLUDE NULLS ((first_quarter, second_quarter) ",
11170+
"FOR half_of_the_year IN (",
11171+
"(sales.Q1, sales.Q2) AS H1, ",
11172+
"(sales.Q3, sales.Q4) AS H2",
11173+
"))"
11174+
);
11175+
11176+
if let Unpivot { columns, .. } =
11177+
&verified_only_select(sql_unpivot_with_alias_and_multi_value_and_qualifier).from[0].relation
11178+
{
11179+
assert_eq!(
11180+
*columns,
11181+
vec![
11182+
ExprWithAlias {
11183+
expr: Expr::Tuple(vec![
11184+
Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q1"),]),
11185+
Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q2"),]),
11186+
]),
11187+
alias: Some(Ident::new("H1")),
11188+
},
11189+
ExprWithAlias {
11190+
expr: Expr::Tuple(vec![
11191+
Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q3"),]),
11192+
Expr::CompoundIdentifier(vec![Ident::new("sales"), Ident::new("Q4"),]),
11193+
]),
11194+
alias: Some(Ident::new("H2")),
11195+
},
11196+
]
11197+
);
11198+
}
11199+
11200+
assert_eq!(
11201+
verified_stmt(sql_unpivot_with_alias_and_multi_value_and_qualifier).to_string(),
11202+
sql_unpivot_with_alias_and_multi_value_and_qualifier
11203+
);
1108711204
}
1108811205

1108911206
#[test]
@@ -11181,20 +11298,14 @@ fn parse_pivot_unpivot_table() {
1118111298
index_hints: vec![],
1118211299
}),
1118311300
null_inclusion: None,
11184-
value: Ident {
11185-
value: "population".to_string(),
11186-
quote_style: None,
11187-
span: Span::empty()
11188-
},
11189-
11190-
name: Ident {
11191-
value: "year".to_string(),
11192-
quote_style: None,
11193-
span: Span::empty()
11194-
},
11301+
value: Expr::Identifier(Ident::new("population")),
11302+
name: Ident::new("year"),
1119511303
columns: ["population_2000", "population_2010"]
1119611304
.into_iter()
11197-
.map(Ident::new)
11305+
.map(|col| ExprWithAlias {
11306+
expr: Expr::Identifier(Ident::new(col)),
11307+
alias: None,
11308+
})
1119811309
.collect(),
1119911310
alias: Some(TableAlias {
1120011311
name: Ident::new("u"),

0 commit comments

Comments
 (0)