Skip to content

Commit fad46e3

Browse files
chenkovskychenkovsky
committed
feat: support multi value column unpivot
1 parent 4921846 commit fad46e3

5 files changed

Lines changed: 167 additions & 20 deletions

File tree

‎src/ast/mod.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -76,12 +76,12 @@ pub use self::query::{
7676
AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
7777
ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, ExprWithAliasAndOrderBy, Fetch, ForClause,
7878
ForJson, ForXml, FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias,
79-
IlikeSelectItem, InputFormatClause, Interpolate, InterpolateExpr, Join, JoinConstraint,
80-
JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, JsonTableNamedColumn,
81-
JsonTableNestedColumn, LateralView, LimitClause, LockClause, LockType, MatchRecognizePattern,
82-
MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
83-
OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind, OrderByOptions,
84-
PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
79+
IdentsWithAlias, IlikeSelectItem, InputFormatClause, Interpolate, InterpolateExpr, Join,
80+
JoinConstraint, JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling,
81+
JsonTableNamedColumn, JsonTableNestedColumn, LateralView, LimitClause, LockClause, LockType,
82+
MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
83+
NonBlock, Offset, OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind,
84+
OrderByOptions, PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
8585
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
8686
SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SetExpr, SetOperator,
8787
SetQuantifier, Setting, SymbolDefinition, Table, TableAlias, TableAliasColumnDef, TableFactor,

‎src/ast/query.rs‎

Lines changed: 52 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -745,6 +745,47 @@ impl fmt::Display for IdentWithAlias {
745745
}
746746
}
747747

748+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
749+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
750+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
751+
pub struct IdentsWithAlias {
752+
pub idents: Vec<Ident>,
753+
pub alias: Option<Ident>,
754+
}
755+
756+
impl IdentsWithAlias {
757+
pub fn new(idents: Vec<Ident>, alias: Option<Ident>) -> Self {
758+
Self { idents, alias }
759+
}
760+
}
761+
762+
impl fmt::Display for IdentsWithAlias {
763+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
764+
match self.idents.len() {
765+
0 => Ok(()),
766+
1 => {
767+
if let Some(alias) = &self.alias {
768+
write!(f, "{} AS {}", self.idents[0], alias)
769+
} else {
770+
write!(f, "{}", self.idents[0])
771+
}
772+
}
773+
_ => {
774+
if let Some(alias) = &self.alias {
775+
write!(
776+
f,
777+
"({}) AS {}",
778+
display_comma_separated(&self.idents),
779+
alias
780+
)
781+
} else {
782+
write!(f, "({})", display_comma_separated(&self.idents))
783+
}
784+
}
785+
}
786+
}
787+
}
788+
748789
/// Additional options for wildcards, e.g. Snowflake `EXCLUDE`/`RENAME` and Bigquery `EXCEPT`.
749790
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
750791
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1351,9 +1392,9 @@ pub enum TableFactor {
13511392
/// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
13521393
Unpivot {
13531394
table: Box<TableFactor>,
1354-
value: Ident,
1395+
value: Vec<Ident>,
13551396
name: Ident,
1356-
columns: Vec<Ident>,
1397+
columns: Vec<IdentsWithAlias>,
13571398
null_inclusion: Option<NullInclusion>,
13581399
alias: Option<TableAlias>,
13591400
},
@@ -2035,10 +2076,17 @@ impl fmt::Display for TableFactor {
20352076
if let Some(null_inclusion) = null_inclusion {
20362077
write!(f, " {null_inclusion} ")?;
20372078
}
2079+
write!(f, "(")?;
2080+
if value.len() == 1 {
2081+
// single value column unpivot
2082+
write!(f, "{}", value[0])?;
2083+
} else {
2084+
// multi value column unpivot
2085+
write!(f, "({})", display_comma_separated(value))?;
2086+
}
20382087
write!(
20392088
f,
2040-
"({} FOR {} IN ({}))",
2041-
value,
2089+
" FOR {} IN ({}))",
20422090
name,
20432091
display_comma_separated(columns)
20442092
)?;

‎src/ast/spans.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1985,9 +1985,15 @@ impl Spanned for TableFactor {
19851985
alias,
19861986
} => union_spans(
19871987
core::iter::once(table.span())
1988-
.chain(core::iter::once(value.span))
1988+
.chain(value.iter().map(|i| i.span))
19891989
.chain(core::iter::once(name.span))
1990-
.chain(columns.iter().map(|i| i.span))
1990+
.chain(columns.iter().flat_map(|ilist| {
1991+
ilist
1992+
.idents
1993+
.iter()
1994+
.map(|i| i.span)
1995+
.chain(ilist.alias.as_ref().map(|a| a.span))
1996+
}))
19911997
.chain(alias.as_ref().map(|alias| alias.span())),
19921998
),
19931999
TableFactor::MatchRecognize {

‎src/parser/mod.rs‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10810,6 +10810,29 @@ impl<'a> Parser<'a> {
1081010810
}
1081110811
}
1081210812

10813+
pub fn parse_identifiers_with_alias(&mut self) -> Result<IdentsWithAlias, ParserError> {
10814+
let idents = match self.peek_token_ref().token {
10815+
Token::LParen => self.parse_parenthesized_column_list(Mandatory, false)?,
10816+
_ => vec![self.parse_identifier()?],
10817+
};
10818+
let alias = if self.parse_keyword(Keyword::AS) {
10819+
Some(self.parse_identifier()?)
10820+
} else {
10821+
None
10822+
};
10823+
Ok(IdentsWithAlias { idents, alias })
10824+
}
10825+
10826+
pub fn parse_parenthesized_columns_with_alias_list(
10827+
&mut self,
10828+
optional: IsOptional,
10829+
allow_empty: bool,
10830+
) -> Result<Vec<IdentsWithAlias>, ParserError> {
10831+
self.parse_parenthesized_column_list_inner(optional, allow_empty, |p| {
10832+
p.parse_identifiers_with_alias()
10833+
})
10834+
}
10835+
1081310836
/// Parses a parenthesized comma-separated list of unqualified, possibly quoted identifiers.
1081410837
/// For example: `(col1, "col 2", ...)`
1081510838
pub fn parse_parenthesized_column_list(
@@ -13882,11 +13905,20 @@ impl<'a> Parser<'a> {
1388213905
None
1388313906
};
1388413907
self.expect_token(&Token::LParen)?;
13885-
let value = self.parse_identifier()?;
13908+
let value = match self.peek_token_ref().token {
13909+
Token::LParen => {
13910+
// multi value column unpivot
13911+
self.parse_parenthesized_column_list(Mandatory, false)?
13912+
}
13913+
_ => {
13914+
// single value column unpivot
13915+
vec![self.parse_identifier()?]
13916+
}
13917+
};
1388613918
self.expect_keyword_is(Keyword::FOR)?;
1388713919
let name = self.parse_identifier()?;
1388813920
self.expect_keyword_is(Keyword::IN)?;
13889-
let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
13921+
let columns = self.parse_parenthesized_columns_with_alias_list(Mandatory, false)?;
1389013922
self.expect_token(&Token::RParen)?;
1389113923
let alias = self.maybe_parse_table_alias()?;
1389213924
Ok(TableFactor::Unpivot {

‎tests/sqlparser_common.rs‎

Lines changed: 67 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10947,11 +10947,11 @@ fn parse_unpivot_table() {
1094710947
index_hints: vec![],
1094810948
}),
1094910949
null_inclusion: None,
10950-
value: Ident {
10950+
value: vec![Ident {
1095110951
value: "quantity".to_string(),
1095210952
quote_style: None,
1095310953
span: Span::empty(),
10954-
},
10954+
}],
1095510955

1095610956
name: Ident {
1095710957
value: "quarter".to_string(),
@@ -10960,7 +10960,7 @@ fn parse_unpivot_table() {
1096010960
},
1096110961
columns: ["Q1", "Q2", "Q3", "Q4"]
1096210962
.into_iter()
10963-
.map(Ident::new)
10963+
.map(|col| IdentsWithAlias::new(vec![Ident::new(col)], None))
1096410964
.collect(),
1096510965
alias: Some(TableAlias {
1096610966
name: Ident::new("u"),
@@ -11022,6 +11022,67 @@ fn parse_unpivot_table() {
1102211022
verified_stmt(sql_unpivot_include_nulls).to_string(),
1102311023
sql_unpivot_include_nulls
1102411024
);
11025+
11026+
let sql_unpivot_with_alias = concat!(
11027+
"SELECT * FROM sales AS s ",
11028+
"UNPIVOT INCLUDE NULLS (quantity FOR quarter IN (Q1 AS Quater1, Q2 AS Quater2, Q3 AS Quater3, Q4 AS Quater4)) AS u (product, quarter, quantity)"
11029+
);
11030+
11031+
if let Unpivot { value, columns, .. } =
11032+
&verified_only_select(sql_unpivot_with_alias).from[0].relation
11033+
{
11034+
assert_eq!(
11035+
*columns,
11036+
vec![
11037+
IdentsWithAlias::new(vec![Ident::new("Q1")], Some(Ident::new("Quater1"))),
11038+
IdentsWithAlias::new(vec![Ident::new("Q2")], Some(Ident::new("Quater2"))),
11039+
IdentsWithAlias::new(vec![Ident::new("Q3")], Some(Ident::new("Quater3"))),
11040+
IdentsWithAlias::new(vec![Ident::new("Q4")], Some(Ident::new("Quater4"))),
11041+
]
11042+
);
11043+
assert_eq!(*value, vec![Ident::new("quantity")]);
11044+
}
11045+
11046+
assert_eq!(
11047+
verified_stmt(sql_unpivot_with_alias).to_string(),
11048+
sql_unpivot_with_alias
11049+
);
11050+
11051+
let sql_unpivot_with_alias = concat!(
11052+
"SELECT * FROM sales AS s ",
11053+
"UNPIVOT INCLUDE NULLS ((first_quarter, second_quarter) ",
11054+
"FOR half_of_the_year IN (",
11055+
"(Q1, Q2) AS H1, ",
11056+
"(Q3, Q4) AS H2",
11057+
"))"
11058+
);
11059+
11060+
if let Unpivot { value, columns, .. } =
11061+
&verified_only_select(sql_unpivot_with_alias).from[0].relation
11062+
{
11063+
assert_eq!(
11064+
*columns,
11065+
vec![
11066+
IdentsWithAlias::new(
11067+
vec![Ident::new("Q1"), Ident::new("Q2")],
11068+
Some(Ident::new("H1"))
11069+
),
11070+
IdentsWithAlias::new(
11071+
vec![Ident::new("Q3"), Ident::new("Q4")],
11072+
Some(Ident::new("H2"))
11073+
),
11074+
]
11075+
);
11076+
assert_eq!(
11077+
*value,
11078+
vec![Ident::new("first_quarter"), Ident::new("second_quarter")]
11079+
);
11080+
}
11081+
11082+
assert_eq!(
11083+
verified_stmt(sql_unpivot_with_alias).to_string(),
11084+
sql_unpivot_with_alias
11085+
);
1102511086
}
1102611087

1102711088
#[test]
@@ -11119,11 +11180,11 @@ fn parse_pivot_unpivot_table() {
1111911180
index_hints: vec![],
1112011181
}),
1112111182
null_inclusion: None,
11122-
value: Ident {
11183+
value: vec![Ident {
1112311184
value: "population".to_string(),
1112411185
quote_style: None,
1112511186
span: Span::empty()
11126-
},
11187+
}],
1112711188

1112811189
name: Ident {
1112911190
value: "year".to_string(),
@@ -11132,7 +11193,7 @@ fn parse_pivot_unpivot_table() {
1113211193
},
1113311194
columns: ["population_2000", "population_2010"]
1113411195
.into_iter()
11135-
.map(Ident::new)
11196+
.map(|col| IdentsWithAlias::new(vec![Ident::new(col)], None))
1113611197
.collect(),
1113711198
alias: Some(TableAlias {
1113811199
name: Ident::new("u"),

0 commit comments

Comments
 (0)