mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-04 14:28:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

snowflake: support for UNPIVOT and a fix for chained PIVOTs (#983)

Browse source
This commit is contained in:
Joey Hain 2023-10-02 10:28:13 -07:00 • committed by GitHub
parent 2786c7eaf1
commit 40e2ecbdf3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 231 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

67
src/ast/query.rs
View file

@ -720,13 +720,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: Ident,
name: Ident,
columns: Vec<Ident>,
alias: Option<TableAlias>,
},
}
@ -810,32 +825,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())
Expr::CompoundIdentifier(value_column.to_vec()),
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(())
}

2
src/keywords.rs
View file

@ -635,6 +635,7 @@ define_keywords!(
UNKNOWN,
UNLOGGED,
UNNEST,
UNPIVOT,
UNSIGNED,
UNTIL,
UPDATE,
@ -693,6 +694,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
Keyword::HAVING,
Keyword::ORDER,
Keyword::PIVOT,
Keyword::UNPIVOT,
Keyword::TOP,
Keyword::LATERAL,
Keyword::VIEW,

54
src/parser/mod.rs
View file

@ -6276,9 +6276,8 @@ impl<'a> Parser<'a> {
| TableFactor::Table { alias, .. }
| TableFactor::UNNEST { alias, .. }
| TableFactor::TableFunction { 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 {
@ -6357,11 +6356,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) {
@ -6373,14 +6367,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)
}
}
@ -6417,8 +6422,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 {
@ -6435,12 +6439,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,
})
}

155
tests/sqlparser_common.rs
View file

@ -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,
@ -7257,10 +7257,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"),
columns: vec![]
table: Box::new(TableFactor::Table {
name: ObjectName(vec![Ident::new("monthly_sales")]),
alias: Some(TableAlias {
name: Ident::new("a"),
columns: vec![]
}),
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}),
aggregate_function: Expr::Function(Function {
name: ObjectName(vec![Ident::new("SUM")]),
@ -7279,7 +7285,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
@ -7290,17 +7296,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(),
@ -7308,6 +7312,135 @@ fn parse_pivot_table() {
);
}
#[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"),
columns: vec![]
}),
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}),
value: Ident {
value: "quantity".to_string(),
quote_style: None
},
name: Ident {
value: "quarter".to_string(),
quote_style: None
},
columns: ["Q1", "Q2", "Q3", "Q4"]
.into_iter()
.map(Ident::new)
.collect(),
alias: Some(TableAlias {
name: Ident::new("u"),
columns: ["product", "quarter", "quantity"]
.into_iter()
.map(Ident::new)
.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"),
columns: vec![]
}),
args: None,
with_hints: vec![],
version: None,
partitions: vec![],
}),
value: Ident {
value: "population".to_string(),
quote_style: None
},
name: Ident {
value: "year".to_string(),
quote_style: None
},
columns: ["population_2000", "population_2010"]
.into_iter()
.map(Ident::new)
.collect(),
alias: Some(TableAlias {
name: Ident::new("u"),
columns: vec![]
}),
}),
aggregate_function: Expr::Function(Function {
name: ObjectName(vec![Ident::new("sum")]),
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
Expr::Identifier(Ident::new("population"))
))]),
over: None,
distinct: false,
special: false,
order_by: vec![],
}),
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"),
columns: vec![]
}),
}
);
assert_eq!(verified_stmt(sql).to_string(), sql);
}
/// Makes a predicate that looks like ((user_id = $id) OR user_id = $2...)
fn make_where_clause(num: usize) -> String {
use std::fmt::Write;