Skip to content

Commit

Permalink
snowflake: support for UNPIVOT and a fix for chained PIVOTs (apache#983…
Browse files Browse the repository at this point in the history 
…) (#16)

Co-authored-by: Joey Hain <[email protected]>
  • Loading branch information
@lustefaniak @jmhain
lustefaniak and jmhain authored Oct 26, 2023
1 parent 019d5f2 commit 954cff4
Show file tree
Hide file tree
Showing 4 changed files with 221 additions and 47 deletions.
67 changes: 46 additions & 21 deletions src/ast/query.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -736,13 +736,28 @@ pub enum TableFactor {
/// For example `FROM monthly_sales PIVOT(sum(amount) FOR MONTH IN ('JAN', 'FEB'))`
/// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot>
Pivot {
#[cfg_attr(feature = "visitor", visit(with = "visit_relation"))]
name: ObjectName,
table_alias: Option<TableAlias>,
#[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
table: Box<TableFactor>,
aggregate_function: Expr, // Function expression
value_column: Vec<Ident>,
pivot_values: Vec<Value>,
pivot_alias: Option<TableAlias>,
alias: Option<TableAlias>,
},
/// An UNPIVOT operation on a table.
///
/// Syntax:
/// ```sql
/// table UNPIVOT(value FOR name IN (column1, [ column2, ... ])) [ alias ]
/// ```
///
/// See <https://docs.snowflake.com/en/sql-reference/constructs/unpivot>.
Unpivot {
#[cfg_attr(feature = "visitor", visit(with = "visit_table_factor"))]
table: Box<TableFactor>,
value: WithSpan<Ident>,
name: WithSpan<Ident>,
columns: Vec<WithSpan<Ident>>,
alias: Option<TableAlias>,
},
}

Expand Down Expand Up @@ -849,32 +864,42 @@ impl fmt::Display for TableFactor {
Ok(())
}
TableFactor::Pivot {
name,
table_alias,
table,
aggregate_function,
value_column,
pivot_values,
pivot_alias,
alias,
} => {
write!(f, "{}", name)?;
if table_alias.is_some() {
write!(f, " AS {}", table_alias.as_ref().unwrap())?;
}
write!(
f,
" PIVOT({} FOR {} IN (",
"{} PIVOT({} FOR {} IN ({}))",
table,
aggregate_function,
Expr::CompoundIdentifier(value_column.to_vec().empty_span())
Expr::CompoundIdentifier(value_column.to_vec().empty_span()),
display_comma_separated(pivot_values)
)?;
for value in pivot_values {
write!(f, "{}", value)?;
if !value.eq(pivot_values.last().unwrap()) {
write!(f, ", ")?;
}
if alias.is_some() {
write!(f, " AS {}", alias.as_ref().unwrap())?;
}
write!(f, "))")?;
if pivot_alias.is_some() {
write!(f, " AS {}", pivot_alias.as_ref().unwrap())?;
Ok(())
}
TableFactor::Unpivot {
table,
value,
name,
columns,
alias,
} => {
write!(
f,
"{} UNPIVOT({} FOR {} IN ({}))",
table,
value,
name,
display_comma_separated(columns)
)?;
if alias.is_some() {
write!(f, " AS {}", alias.as_ref().unwrap())?;
}
Ok(())
}
Expand Down
2 changes: 2 additions & 0 deletions src/keywords.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -641,6 +641,7 @@ define_keywords!(
UNKNOWN,
UNLOGGED,
UNNEST,
UNPIVOT,
UNSIGNED,
UNTIL,
UPDATE,
Expand Down Expand Up @@ -699,6 +700,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
Keyword::HAVING,
Keyword::ORDER,
Keyword::PIVOT,
Keyword::UNPIVOT,
Keyword::TOP,
Keyword::LATERAL,
Keyword::VIEW,
Expand Down
54 changes: 39 additions & 15 deletions src/parser/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6615,9 +6615,8 @@ impl<'a> Parser<'a> {
| TableFactor::UNNEST { alias, .. }
| TableFactor::TableFunction { alias, .. }
| TableFactor::FieldAccessor { alias, .. }
| TableFactor::Pivot {
pivot_alias: alias, ..
}
| TableFactor::Pivot { alias, .. }
| TableFactor::Unpivot { alias, .. }
| TableFactor::NestedJoin { alias, .. } => {
// but not `FROM (mytable AS alias1) AS alias2`.
if let Some(inner_alias) = alias {
Expand Down Expand Up @@ -6696,11 +6695,6 @@ impl<'a> Parser<'a> {

let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;

// Pivot
if self.parse_keyword(Keyword::PIVOT) {
return self.parse_pivot_table_factor(name, alias);
}

// MSSQL-specific table hints:
let mut with_hints = vec![];
if self.parse_keyword(Keyword::WITH) {
Expand All @@ -6712,14 +6706,25 @@ impl<'a> Parser<'a> {
self.prev_token();
}
};
Ok(TableFactor::Table {

let mut table = TableFactor::Table {
name,
alias,
args,
with_hints,
version,
partitions,
})
};

while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {
table = match kw {
Keyword::PIVOT => self.parse_pivot_table_factor(table)?,
Keyword::UNPIVOT => self.parse_unpivot_table_factor(table)?,
_ => unreachable!(),
}
}

Ok(table)
}
}

Expand Down Expand Up @@ -6765,8 +6770,7 @@ impl<'a> Parser<'a> {

pub fn parse_pivot_table_factor(
&mut self,
name: ObjectName,
table_alias: Option<TableAlias>,
table: TableFactor,
) -> Result<TableFactor, ParserError> {
self.expect_token(&Token::LParen)?;
let function_name = match self.next_token().token {
Expand All @@ -6783,12 +6787,32 @@ impl<'a> Parser<'a> {
self.expect_token(&Token::RParen)?;
let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
Ok(TableFactor::Pivot {
name,
table_alias,
table: Box::new(table),
aggregate_function: function,
value_column,
pivot_values,
pivot_alias: alias,
alias,
})
}

pub fn parse_unpivot_table_factor(
&mut self,
table: TableFactor,
) -> Result<TableFactor, ParserError> {
self.expect_token(&Token::LParen)?;
let value = self.parse_identifier()?;
self.expect_keyword(Keyword::FOR)?;
let name = self.parse_identifier()?;
self.expect_keyword(Keyword::IN)?;
let columns = self.parse_parenthesized_column_list(Mandatory, false)?;
self.expect_token(&Token::RParen)?;
let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
Ok(TableFactor::Unpivot {
table: Box::new(table),
value,
name,
columns,
alias,
})
}

Expand Down
145 changes: 134 additions & 11 deletions tests/sqlparser_common.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
use matches::assert_matches;
use sqlparser::ast::SelectItem::UnnamedExpr;
use sqlparser::ast::TableFactor::Pivot;
use sqlparser::ast::TableFactor::{Pivot, Unpivot};
use sqlparser::ast::*;
use sqlparser::dialect::{
AnsiDialect, BigQueryDialect, ClickHouseDialect, DuckDbDialect, GenericDialect, HiveDialect,
Expand Down Expand Up @@ -7736,10 +7736,16 @@ fn parse_pivot_table() {
assert_eq!(
verified_only_select(sql).from[0].relation,
Pivot {
name: ObjectName(vec![Ident::new("monthly_sales")]),
table_alias: Some(TableAlias {
name: Ident::new("a").empty_span(),
columns: vec![]
table: Box::new(TableFactor::Table {
name: ObjectName(vec![Ident::new("monthly_sales")]),
alias: Some(TableAlias {
name: Ident::new("a").empty_span(),
columns: vec![]
}),
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}),
aggregate_function: Expr::Function(Function {
name: ObjectName(vec![Ident::new("SUM")]),
Expand All @@ -7764,7 +7770,7 @@ fn parse_pivot_table() {
Value::SingleQuotedString("MAR".to_string()),
Value::SingleQuotedString("APR".to_string()),
],
pivot_alias: Some(TableAlias {
alias: Some(TableAlias {
name: Ident {
value: "p".to_string(),
quote_style: None
Expand All @@ -7776,17 +7782,15 @@ fn parse_pivot_table() {
);
assert_eq!(verified_stmt(sql).to_string(), sql);

// parsing should succeed with empty alias
let sql_without_table_alias = concat!(
"SELECT * FROM monthly_sales ",
"PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
"ORDER BY EMPID"
);
assert_matches!(
verified_only_select(sql_without_table_alias).from[0].relation,
Pivot {
table_alias: None, // parsing should succeed with empty alias
..
}
&verified_only_select(sql_without_table_alias).from[0].relation,
Pivot { table, .. } if matches!(&**table, TableFactor::Table { alias: None, .. })
);
assert_eq!(
verified_stmt(sql_without_table_alias).to_string(),
Expand Down Expand Up @@ -7820,6 +7824,125 @@ fn parse_within_group() {
}]),
})
);
}

#[test]
fn parse_unpivot_table() {
let sql = concat!(
"SELECT * FROM sales AS s ",
"UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4)) AS u (product, quarter, quantity)"
);

pretty_assertions::assert_eq!(
verified_only_select(sql).from[0].relation,
Unpivot {
table: Box::new(TableFactor::Table {
name: ObjectName(vec![Ident::new("sales")]),
alias: Some(TableAlias {
name: Ident::new("s").empty_span(),
columns: vec![]
}),
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}),
value: Ident::new("quantity").empty_span(),

name: Ident::new("quarter").empty_span(),
columns: ["Q1", "Q2", "Q3", "Q4"]
.into_iter()
.map(|f| Ident::new(f).empty_span())
.collect(),
alias: Some(TableAlias {
name: Ident::new("u").empty_span(),
columns: ["product", "quarter", "quantity"]
.into_iter()
.map(|f| Ident::new(f).empty_span())
.collect()
}),
}
);
assert_eq!(verified_stmt(sql).to_string(), sql);

let sql_without_aliases = concat!(
"SELECT * FROM sales ",
"UNPIVOT(quantity FOR quarter IN (Q1, Q2, Q3, Q4))"
);

assert_matches!(
&verified_only_select(sql_without_aliases).from[0].relation,
Unpivot {
table,
alias: None,
..
} if matches!(&**table, TableFactor::Table { alias: None, .. })
);
assert_eq!(
verified_stmt(sql_without_aliases).to_string(),
sql_without_aliases
);
}

#[test]
fn parse_pivot_unpivot_table() {
let sql = concat!(
"SELECT * FROM census AS c ",
"UNPIVOT(population FOR year IN (population_2000, population_2010)) AS u ",
"PIVOT(sum(population) FOR year IN ('population_2000', 'population_2010')) AS p"
);

pretty_assertions::assert_eq!(
verified_only_select(sql).from[0].relation,
Pivot {
table: Box::new(Unpivot {
table: Box::new(TableFactor::Table {
name: ObjectName(vec![Ident::new("census")]),
alias: Some(TableAlias {
name: Ident::new("c").empty_span(),
columns: vec![]
}),
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}),
value: Ident::new("population").empty_span(),
name: Ident::new("year").empty_span(),
columns: ["population_2000", "population_2010"]
.into_iter()
.map(|f| Ident::new(f).empty_span())
.collect(),
alias: Some(TableAlias {
name: Ident::new("u").empty_span(),
columns: vec![]
}),
}),
aggregate_function: Expr::Function(Function {
name: ObjectName(vec![Ident::new("sum")]),
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
Expr::Identifier(Ident::new("population").empty_span())
))]),
within_group: None,
over: None,
distinct: false,
special: false,
order_by: vec![],
limit: None,
on_overflow: None,
null_treatment: None,
}),
value_column: vec![Ident::new("year")],
pivot_values: vec![
Value::SingleQuotedString("population_2000".to_string()),
Value::SingleQuotedString("population_2010".to_string())
],
alias: Some(TableAlias {
name: Ident::new("p").empty_span(),
columns: vec![]
}),
}
);
assert_eq!(verified_stmt(sql).to_string(), sql);
}

Expand Down

0 comments on commit 954cff4

Please sign in to comment.