Skip to content

Commit 46bf2d7

Browse files
iffyioiffyio
committed
Extend pivot operator support
Extends the parser with support for: - Multiple aggregate function calls - Aliased aggregation function calls - Aliased pivot values - also allowing arbitrary expressions as values https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#pivot_operator
1 parent deaa6d8 commit 46bf2d7

File tree

5 files changed

+134
-46
lines changed

5 files changed

+134
-46
lines changed

‎src/ast/mod.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,8 @@ pub use self::dml::{Delete, Insert};
4141
pub use self::operator::{BinaryOperator, UnaryOperator};
4242
pub use self::query::{
4343
AfterMatchSkip, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode, ExceptSelectItem,
44-
ExcludeSelectItem, Fetch, ForClause, ForJson, ForXml, GroupByExpr, IdentWithAlias,
45-
IlikeSelectItem, Join, JoinConstraint, JoinOperator, JsonTableColumn,
44+
ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml, GroupByExpr,
45+
IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator, JsonTableColumn,
4646
JsonTableColumnErrorHandling, LateralView, LockClause, LockType, MatchRecognizePattern,
4747
MatchRecognizeSymbol, Measure, NamedWindowDefinition, NonBlock, Offset, OffsetRows,
4848
OrderByExpr, Query, RenameSelectItem, RepetitionQuantifier, ReplaceSelectElement,

‎src/ast/query.rs

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -731,6 +731,31 @@ impl fmt::Display for TableWithJoins {
731731
}
732732
}
733733

734+
/// An expression optionally followed by an alias.
735+
///
736+
/// Example:
737+
/// ```sql
738+
/// 42 AS myint
739+
/// ```
740+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
741+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
742+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
743+
pub struct ExprWithAlias {
744+
pub expr: Expr,
745+
pub alias: Option<Ident>,
746+
}
747+
748+
impl fmt::Display for ExprWithAlias {
749+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
750+
let ExprWithAlias { expr, alias } = self;
751+
write!(f, "{expr}")?;
752+
if let Some(alias) = alias {
753+
write!(f, " AS {alias}")?;
754+
}
755+
Ok(())
756+
}
757+
}
758+
734759
/// A table name or a parenthesized subquery with an optional alias
735760
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
736761
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -829,12 +854,14 @@ pub enum TableFactor {
829854
},
830855
/// Represents PIVOT operation on a table.
831856
/// For example `FROM monthly_sales PIVOT(sum(amount) FOR MONTH IN ('JAN', 'FEB'))`
832-
/// See <https://docs.snowflake.com/en/sql-reference/constructs/pivot>
857+
///
858+
/// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/query-syntax#pivot_operator)
859+
/// [Snowflake](https://docs.snowflake.com/en/sql-reference/constructs/pivot)
833860
Pivot {
834861
table: Box<TableFactor>,
835-
aggregate_function: Expr, // Function expression
862+
aggregate_functions: Vec<ExprWithAlias>, // Function expression
836863
value_column: Vec<Ident>,
837-
pivot_values: Vec<Value>,
864+
pivot_values: Vec<ExprWithAlias>,
838865
alias: Option<TableAlias>,
839866
},
840867
/// An UNPIVOT operation on a table.
@@ -1199,7 +1226,7 @@ impl fmt::Display for TableFactor {
11991226
}
12001227
TableFactor::Pivot {
12011228
table,
1202-
aggregate_function,
1229+
aggregate_functions,
12031230
value_column,
12041231
pivot_values,
12051232
alias,
@@ -1208,7 +1235,7 @@ impl fmt::Display for TableFactor {
12081235
f,
12091236
"{} PIVOT({} FOR {} IN ({}))",
12101237
table,
1211-
aggregate_function,
1238+
display_comma_separated(aggregate_functions),
12121239
Expr::CompoundIdentifier(value_column.to_vec()),
12131240
display_comma_separated(pivot_values)
12141241
)?;

‎src/ast/visitor.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -850,6 +850,14 @@ mod tests {
850850
"PRE: EXPR: a.amount",
851851
"POST: EXPR: a.amount",
852852
"POST: EXPR: SUM(a.amount)",
853+
"PRE: EXPR: 'JAN'",
854+
"POST: EXPR: 'JAN'",
855+
"PRE: EXPR: 'FEB'",
856+
"POST: EXPR: 'FEB'",
857+
"PRE: EXPR: 'MAR'",
858+
"POST: EXPR: 'MAR'",
859+
"PRE: EXPR: 'APR'",
860+
"POST: EXPR: 'APR'",
853861
"POST: TABLE FACTOR: monthly_sales PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d)",
854862
"PRE: EXPR: EMPID",
855863
"POST: EXPR: EMPID",

‎src/parser/mod.rs

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8674,27 +8674,49 @@ impl<'a> Parser<'a> {
86748674
})
86758675
}
86768676

8677+
fn parse_aliased_function_call(&mut self) -> Result<ExprWithAlias, ParserError> {
8678+
let function_name = match self.next_token().token {
8679+
Token::Word(w) => Ok(w.value),
8680+
_ => self.expected("a function identifier", self.peek_token()),
8681+
}?;
8682+
let expr = self.parse_function(ObjectName(vec![Ident::new(function_name)]))?;
8683+
let alias = if self.parse_keyword(Keyword::AS) {
8684+
Some(self.parse_identifier(false)?)
8685+
} else {
8686+
None
8687+
};
8688+
8689+
Ok(ExprWithAlias { expr, alias })
8690+
}
8691+
8692+
fn parse_expr_with_alias(&mut self) -> Result<ExprWithAlias, ParserError> {
8693+
let expr = self.parse_expr()?;
8694+
let alias = if self.parse_keyword(Keyword::AS) {
8695+
Some(self.parse_identifier(false)?)
8696+
} else {
8697+
None
8698+
};
8699+
8700+
Ok(ExprWithAlias { expr, alias })
8701+
}
8702+
86778703
pub fn parse_pivot_table_factor(
86788704
&mut self,
86798705
table: TableFactor,
86808706
) -> Result<TableFactor, ParserError> {
86818707
self.expect_token(&Token::LParen)?;
8682-
let function_name = match self.next_token().token {
8683-
Token::Word(w) => Ok(w.value),
8684-
_ => self.expected("an aggregate function name", self.peek_token()),
8685-
}?;
8686-
let function = self.parse_function(ObjectName(vec![Ident::new(function_name)]))?;
8708+
let aggregate_functions = self.parse_comma_separated(Self::parse_aliased_function_call)?;
86878709
self.expect_keyword(Keyword::FOR)?;
86888710
let value_column = self.parse_object_name(false)?.0;
86898711
self.expect_keyword(Keyword::IN)?;
86908712
self.expect_token(&Token::LParen)?;
8691-
let pivot_values = self.parse_comma_separated(Parser::parse_value)?;
8713+
let pivot_values = self.parse_comma_separated(Self::parse_expr_with_alias)?;
86928714
self.expect_token(&Token::RParen)?;
86938715
self.expect_token(&Token::RParen)?;
86948716
let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
86958717
Ok(TableFactor::Pivot {
86968718
table: Box::new(table),
8697-
aggregate_function: function,
8719+
aggregate_functions,
86988720
value_column,
86998721
pivot_values,
87008722
alias,

‎tests/sqlparser_common.rs

Lines changed: 63 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8408,11 +8408,32 @@ fn parse_escaped_string_without_unescape() {
84088408
#[test]
84098409
fn parse_pivot_table() {
84108410
let sql = concat!(
8411-
"SELECT * FROM monthly_sales AS a ",
8412-
"PIVOT(SUM(a.amount) FOR a.MONTH IN ('JAN', 'FEB', 'MAR', 'APR')) AS p (c, d) ",
8411+
"SELECT * FROM monthly_sales AS a PIVOT(",
8412+
"SUM(a.amount), ",
8413+
"SUM(b.amount) AS t, ",
8414+
"SUM(c.amount) AS u ",
8415+
"FOR a.MONTH IN (1 AS x, 'two', three AS y)) AS p (c, d) ",
84138416
"ORDER BY EMPID"
84148417
);
84158418

8419+
fn expected_function(table: &'static str, alias: Option<&'static str>) -> ExprWithAlias {
8420+
ExprWithAlias {
8421+
expr: Expr::Function(Function {
8422+
name: ObjectName(vec![Ident::new("SUM")]),
8423+
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
8424+
Expr::CompoundIdentifier(vec![Ident::new(table), Ident::new("amount")]),
8425+
))]),
8426+
null_treatment: None,
8427+
filter: None,
8428+
over: None,
8429+
distinct: false,
8430+
special: false,
8431+
order_by: vec![],
8432+
}),
8433+
alias: alias.map(Ident::new),
8434+
}
8435+
}
8436+
84168437
assert_eq!(
84178438
verified_only_select(sql).from[0].relation,
84188439
Pivot {
@@ -8427,24 +8448,25 @@ fn parse_pivot_table() {
84278448
version: None,
84288449
partitions: vec![],
84298450
}),
8430-
aggregate_function: Expr::Function(Function {
8431-
name: ObjectName(vec![Ident::new("SUM")]),
8432-
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
8433-
Expr::CompoundIdentifier(vec![Ident::new("a"), Ident::new("amount"),])
8434-
))]),
8435-
null_treatment: None,
8436-
filter: None,
8437-
over: None,
8438-
distinct: false,
8439-
special: false,
8440-
order_by: vec![],
8441-
}),
8451+
aggregate_functions: vec![
8452+
expected_function("a", None),
8453+
expected_function("b", Some("t")),
8454+
expected_function("c", Some("u")),
8455+
],
84428456
value_column: vec![Ident::new("a"), Ident::new("MONTH")],
84438457
pivot_values: vec![
8444-
Value::SingleQuotedString("JAN".to_string()),
8445-
Value::SingleQuotedString("FEB".to_string()),
8446-
Value::SingleQuotedString("MAR".to_string()),
8447-
Value::SingleQuotedString("APR".to_string()),
8458+
ExprWithAlias {
8459+
expr: Expr::Value(number("1")),
8460+
alias: Some(Ident::new("x"))
8461+
},
8462+
ExprWithAlias {
8463+
expr: Expr::Value(Value::SingleQuotedString("two".to_string())),
8464+
alias: None
8465+
},
8466+
ExprWithAlias {
8467+
expr: Expr::Identifier(Ident::new("three")),
8468+
alias: Some(Ident::new("y"))
8469+
},
84488470
],
84498471
alias: Some(TableAlias {
84508472
name: Ident {
@@ -8578,22 +8600,31 @@ fn parse_pivot_unpivot_table() {
85788600
columns: vec![]
85798601
}),
85808602
}),
8581-
aggregate_function: Expr::Function(Function {
8582-
name: ObjectName(vec![Ident::new("sum")]),
8583-
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
8584-
Expr::Identifier(Ident::new("population"))
8585-
))]),
8586-
null_treatment: None,
8587-
filter: None,
8588-
over: None,
8589-
distinct: false,
8590-
special: false,
8591-
order_by: vec![],
8592-
}),
8603+
aggregate_functions: vec![ExprWithAlias {
8604+
expr: Expr::Function(Function {
8605+
name: ObjectName(vec![Ident::new("sum")]),
8606+
args: (vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(
8607+
Expr::Identifier(Ident::new("population"))
8608+
))]),
8609+
null_treatment: None,
8610+
filter: None,
8611+
over: None,
8612+
distinct: false,
8613+
special: false,
8614+
order_by: vec![],
8615+
}),
8616+
alias: None
8617+
}],
85938618
value_column: vec![Ident::new("year")],
85948619
pivot_values: vec![
8595-
Value::SingleQuotedString("population_2000".to_string()),
8596-
Value::SingleQuotedString("population_2010".to_string())
8620+
ExprWithAlias {
8621+
expr: Expr::Value(Value::SingleQuotedString("population_2000".to_string())),
8622+
alias: None
8623+
},
8624+
ExprWithAlias {
8625+
expr: Expr::Value(Value::SingleQuotedString("population_2010".to_string())),
8626+
alias: None
8627+
},
85978628
],
85988629
alias: Some(TableAlias {
85998630
name: Ident::new("p"),

0 commit comments

Comments
 (0)