mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-09-26 15:39:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

Extend pivot operator support (#1238)

Browse source
This commit is contained in:
Ifeanyi Ubah 2024-04-30 20:13:26 +02:00 • committed by GitHub
parent 4bfa399919
commit 8626051513
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 140 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

16
src/ast/mod.rs
View file

@ -41,14 +41,14 @@ pub use self::dml::{Delete, Insert};
pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
ExceptSelectItem, ExcludeSelectItem, Fetch, ForClause, ForJson, ForXml, GroupByExpr,
IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator, JsonTableColumn,
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern,
MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
OffsetRows, OrderByExpr, Query, RenameSelectItem, RepetitionQuantifier, ReplaceSelectElement,
ReplaceSelectItem, RowsPerMatch, Select, SelectInto, SelectItem, SetExpr, SetOperator,
SetQuantifier, SymbolDefinition, Table, TableAlias, TableFactor, TableVersion, TableWithJoins,
Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With,
ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
GroupByExpr, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator,
JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem, RepetitionQuantifier,
ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select, SelectInto, SelectItem, SetExpr,
SetOperator, SetQuantifier, SymbolDefinition, Table, TableAlias, TableFactor, TableVersion,
TableWithJoins, Top, TopQuantity, ValueTableMode, Values, WildcardAdditionalOptions, With,
};
pub use self::value::{
escape_double_quote_string, escape_quoted_string, DateTimeField, DollarQuotedString,

37
src/ast/query.rs
View file

@ -802,6 +802,31 @@ impl fmt::Display for ConnectBy {
}
}
/// An expression optionally followed by an alias.
///
/// Example:
/// ```sql
/// 42 AS myint
/// ```
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub struct ExprWithAlias {
pub expr: Expr,
pub alias: Option<Ident>,
}
impl fmt::Display for ExprWithAlias {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let ExprWithAlias { expr, alias } = self;
write!(f, "{expr}")?;
if let Some(alias) = alias {
write!(f, " AS {alias}")?;
}
Ok(())
}
}
/// A table name or a parenthesized subquery with an optional alias
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@ -900,12 +925,14 @@ pub enum TableFactor {
},
/// Represents PIVOT operation on a table.
/// For example `FROM monthly_sales PIVOT(sum(amount) FOR MONTH IN ('JAN', 'FEB'))`
/// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot>
///
/// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#pivot_operator)
/// [Snowflake](https://docs.snowflake.com/en/sql-reference/constructs/pivot)
Pivot {
table: Box<TableFactor>,
aggregate_function: Expr, // Function expression
aggregate_functions: Vec<ExprWithAlias>, // Function expression
value_column: Vec<Ident>,
pivot_values: Vec<Value>,
pivot_values: Vec<ExprWithAlias>,
alias: Option<TableAlias>,
},
/// An UNPIVOT operation on a table.
@ -1270,7 +1297,7 @@ impl fmt::Display for TableFactor {
}
TableFactor::Pivot {
table,
aggregate_function,
aggregate_functions,
value_column,
pivot_values,
alias,
@ -1279,7 +1306,7 @@ impl fmt::Display for TableFactor {
f,
"{} PIVOT({} FOR {} IN ({}))",
table,
aggregate_function,
display_comma_separated(aggregate_functions),
Expr::CompoundIdentifier(value_column.to_vec()),
display_comma_separated(pivot_values)
)?;

8
src/ast/visitor.rs
View file

@ -850,6 +850,14 @@ mod tests {
"PRE: EXPR: a.amount",
"POST: EXPR: a.amount",
"POST: EXPR: SUM(a.amount)",
"PRE: EXPR: 'JAN'",
"POST: EXPR: 'JAN'",
"PRE: EXPR: 'FEB'",
"POST: EXPR: 'FEB'",
"PRE: EXPR: 'MAR'",
"POST: EXPR: 'MAR'",
"PRE: EXPR: 'APR'",
"POST: EXPR: 'APR'",
"POST: TABLE FACTOR: monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
"PRE: EXPR: EMPID",
"POST: EXPR: EMPID",

36
src/parser/mod.rs
View file

@ -8775,27 +8775,49 @@ impl<'a> Parser<'a> {
})
}
fn parse_aliased_function_call(&mut self) -> Result<ExprWithAlias, ParserError> {
let function_name = match self.next_token().token {
Token::Word(w) => Ok(w.value),
_ => self.expected("a function identifier", self.peek_token()),
}?;
let expr = self.parse_function(ObjectName(vec![Ident::new(function_name)]))?;
let alias = if self.parse_keyword(Keyword::AS) {
Some(self.parse_identifier(false)?)
} else {
None
};
Ok(ExprWithAlias { expr, alias })
}
fn parse_expr_with_alias(&mut self) -> Result<ExprWithAlias, ParserError> {
let expr = self.parse_expr()?;
let alias = if self.parse_keyword(Keyword::AS) {
Some(self.parse_identifier(false)?)
} else {
None
};
Ok(ExprWithAlias { expr, alias })
}
pub fn parse_pivot_table_factor(
&mut self,
table: TableFactor,
) -> Result<TableFactor, ParserError> {
self.expect_token(&Token::LParen)?;
let function_name = match self.next_token().token {
Token::Word(w) => Ok(w.value),
_ => self.expected("an aggregate function name", self.peek_token()),
}?;
let function = self.parse_function(ObjectName(vec![Ident::new(function_name)]))?;
let aggregate_functions = self.parse_comma_separated(Self::parse_aliased_function_call)?;
self.expect_keyword(Keyword::FOR)?;
let value_column = self.parse_object_name(false)?.0;
self.expect_keyword(Keyword::IN)?;
self.expect_token(&Token::LParen)?;
let pivot_values = self.parse_comma_separated(Parser::parse_value)?;
let pivot_values = self.parse_comma_separated(Self::parse_expr_with_alias)?;
self.expect_token(&Token::RParen)?;
self.expect_token(&Token::RParen)?;
let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
Ok(TableFactor::Pivot {
table: Box::new(table),
aggregate_function: function,
aggregate_functions,
value_column,
pivot_values,
alias,

73
tests/sqlparser_common.rs
View file

@ -8453,11 +8453,32 @@ fn parse_escaped_string_without_unescape() {
#[test]
fn parse_pivot_table() {
let sql = concat!(
"SELECT * FROM monthly_sales AS a ",
"PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
"SELECT * FROM monthly_sales AS a PIVOT(",
"SUM(a.amount), ",
"SUM(b.amount) AS t, ",
"SUM(c.amount) AS u ",
"FOR a.MONTH IN (1 AS x, 'two', three AS y)) AS p (c, d) ",
"ORDER BY EMPID"
);
fn expected_function(table: &'static str, alias: Option<&'static str>) -> ExprWithAlias {
ExprWithAlias {
expr: Expr::Function(Function {
name: ObjectName(vec![Ident::new("SUM")]),
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
Expr::CompoundIdentifier(vec![Ident::new(table), Ident::new("amount")]),
))]),
null_treatment: None,
filter: None,
over: None,
distinct: false,
special: false,
order_by: vec![],
}),
alias: alias.map(Ident::new),
}
}
assert_eq!(
verified_only_select(sql).from[0].relation,
Pivot {
@ -8472,24 +8493,25 @@ fn parse_pivot_table() {
version: None,
partitions: vec![],
}),
aggregate_function: Expr::Function(Function {
name: ObjectName(vec![Ident::new("SUM")]),
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
Expr::CompoundIdentifier(vec![Ident::new("a"), Ident::new("amount"),])
))]),
null_treatment: None,
filter: None,
over: None,
distinct: false,
special: false,
order_by: vec![],
}),
aggregate_functions: vec![
expected_function("a", None),
expected_function("b", Some("t")),
expected_function("c", Some("u")),
],
value_column: vec![Ident::new("a"), Ident::new("MONTH")],
pivot_values: vec![
Value::SingleQuotedString("JAN".to_string()),
Value::SingleQuotedString("FEB".to_string()),
Value::SingleQuotedString("MAR".to_string()),
Value::SingleQuotedString("APR".to_string()),
ExprWithAlias {
expr: Expr::Value(number("1")),
alias: Some(Ident::new("x"))
},
ExprWithAlias {
expr: Expr::Value(Value::SingleQuotedString("two".to_string())),
alias: None
},
ExprWithAlias {
expr: Expr::Identifier(Ident::new("three")),
alias: Some(Ident::new("y"))
},
],
alias: Some(TableAlias {
name: Ident {
@ -8623,7 +8645,8 @@ fn parse_pivot_unpivot_table() {
columns: vec![]
}),
}),
aggregate_function: Expr::Function(Function {
aggregate_functions: vec![ExprWithAlias {
expr: Expr::Function(Function {
name: ObjectName(vec![Ident::new("sum")]),
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
Expr::Identifier(Ident::new("population"))
@ -8635,10 +8658,18 @@ fn parse_pivot_unpivot_table() {
special: false,
order_by: vec![],
}),
alias: None
}],
value_column: vec![Ident::new("year")],
pivot_values: vec![
Value::SingleQuotedString("population_2000".to_string()),
Value::SingleQuotedString("population_2010".to_string())
ExprWithAlias {
expr: Expr::Value(Value::SingleQuotedString("population_2000".to_string())),
alias: None
},
ExprWithAlias {
expr: Expr::Value(Value::SingleQuotedString("population_2010".to_string())),
alias: None
},
],
alias: Some(TableAlias {
name: Ident::new("p"),