Skip to content

Commit ddf96c7

Browse files
jimexistJiayu Liu
jimexist
authored and
Jiayu Liu
committed
parse grouping sets, rollup, and cube
1 parent d7e84be commit ddf96c7

File tree

4 files changed

+211
-2
lines changed

4 files changed

+211
-2
lines changed

‎src/ast/mod.rs‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,12 @@ pub enum Expr {
278278
Subquery(Box<Query>),
279279
/// The `LISTAGG` function `SELECT LISTAGG(...) WITHIN GROUP (ORDER BY ...)`
280280
ListAgg(ListAgg),
281+
/// The `GROUPING SETS` expr.
282+
GroupingSets(Vec<Vec<Expr>>),
283+
/// The `CUBE` expr.
284+
Cube(Vec<Vec<Expr>>),
285+
/// The `ROLLUP` expr.
286+
Rollup(Vec<Vec<Expr>>),
281287
}
282288

283289
impl fmt::Display for Expr {
@@ -376,6 +382,44 @@ impl fmt::Display for Expr {
376382
Expr::Exists(s) => write!(f, "EXISTS ({})", s),
377383
Expr::Subquery(s) => write!(f, "({})", s),
378384
Expr::ListAgg(listagg) => write!(f, "{}", listagg),
385+
Expr::GroupingSets(sets) => {
386+
write!(f, "GROUPING SETS (")?;
387+
let mut sep = "";
388+
for set in sets {
389+
write!(f, "{}", sep)?;
390+
sep = ", ";
391+
write!(f, "({})", display_comma_separated(set))?;
392+
}
393+
write!(f, ")")
394+
}
395+
Expr::Cube(sets) => {
396+
write!(f, "CUBE (")?;
397+
let mut sep = "";
398+
for set in sets {
399+
write!(f, "{}", sep)?;
400+
sep = ", ";
401+
if set.len() == 1 {
402+
write!(f, "{}", set[0])?;
403+
} else {
404+
write!(f, "({})", display_comma_separated(set))?;
405+
}
406+
}
407+
write!(f, ")")
408+
}
409+
Expr::Rollup(sets) => {
410+
write!(f, "ROLLUP (")?;
411+
let mut sep = "";
412+
for set in sets {
413+
write!(f, "{}", sep)?;
414+
sep = ", ";
415+
if set.len() == 1 {
416+
write!(f, "{}", set[0])?;
417+
} else {
418+
write!(f, "({})", display_comma_separated(set))?;
419+
}
420+
}
421+
write!(f, ")")
422+
}
379423
Expr::Substring {
380424
expr,
381425
substring_from,
@@ -1903,4 +1947,93 @@ mod tests {
19031947
let window_frame = WindowFrame::default();
19041948
assert_eq!(WindowFrameBound::Preceding(None), window_frame.start_bound);
19051949
}
1950+
1951+
#[test]
1952+
fn test_grouping_sets_display() {
1953+
// a and b in different group
1954+
let grouping_sets = Expr::GroupingSets(vec![
1955+
vec![Expr::Identifier(Ident::new("a"))],
1956+
vec![Expr::Identifier(Ident::new("b"))],
1957+
]);
1958+
assert_eq!("GROUPING SETS ((a), (b))", format!("{}", grouping_sets));
1959+
1960+
// a and b in the same group
1961+
let grouping_sets = Expr::GroupingSets(vec![vec![
1962+
Expr::Identifier(Ident::new("a")),
1963+
Expr::Identifier(Ident::new("b")),
1964+
]]);
1965+
assert_eq!("GROUPING SETS ((a, b))", format!("{}", grouping_sets));
1966+
1967+
// (a, b) and (c, d) in different group
1968+
let grouping_sets = Expr::GroupingSets(vec![
1969+
vec![
1970+
Expr::Identifier(Ident::new("a")),
1971+
Expr::Identifier(Ident::new("b")),
1972+
],
1973+
vec![
1974+
Expr::Identifier(Ident::new("c")),
1975+
Expr::Identifier(Ident::new("d")),
1976+
],
1977+
]);
1978+
assert_eq!(
1979+
"GROUPING SETS ((a, b), (c, d))",
1980+
format!("{}", grouping_sets)
1981+
);
1982+
}
1983+
1984+
#[test]
1985+
fn test_rollup_display() {
1986+
let rollup = Expr::Rollup(vec![vec![Expr::Identifier(Ident::new("a"))]]);
1987+
assert_eq!("ROLLUP (a)", format!("{}", rollup));
1988+
1989+
let rollup = Expr::Rollup(vec![vec![
1990+
Expr::Identifier(Ident::new("a")),
1991+
Expr::Identifier(Ident::new("b")),
1992+
]]);
1993+
assert_eq!("ROLLUP ((a, b))", format!("{}", rollup));
1994+
1995+
let rollup = Expr::Rollup(vec![
1996+
vec![Expr::Identifier(Ident::new("a"))],
1997+
vec![Expr::Identifier(Ident::new("b"))],
1998+
]);
1999+
assert_eq!("ROLLUP (a, b)", format!("{}", rollup));
2000+
2001+
let rollup = Expr::Rollup(vec![
2002+
vec![Expr::Identifier(Ident::new("a"))],
2003+
vec![
2004+
Expr::Identifier(Ident::new("b")),
2005+
Expr::Identifier(Ident::new("c")),
2006+
],
2007+
vec![Expr::Identifier(Ident::new("d"))],
2008+
]);
2009+
assert_eq!("ROLLUP (a, (b, c), d)", format!("{}", rollup));
2010+
}
2011+
2012+
#[test]
2013+
fn test_cube_display() {
2014+
let cube = Expr::Cube(vec![vec![Expr::Identifier(Ident::new("a"))]]);
2015+
assert_eq!("CUBE (a)", format!("{}", cube));
2016+
2017+
let cube = Expr::Cube(vec![vec![
2018+
Expr::Identifier(Ident::new("a")),
2019+
Expr::Identifier(Ident::new("b")),
2020+
]]);
2021+
assert_eq!("CUBE ((a, b))", format!("{}", cube));
2022+
2023+
let cube = Expr::Cube(vec![
2024+
vec![Expr::Identifier(Ident::new("a"))],
2025+
vec![Expr::Identifier(Ident::new("b"))],
2026+
]);
2027+
assert_eq!("CUBE (a, b)", format!("{}", cube));
2028+
2029+
let cube = Expr::Cube(vec![
2030+
vec![Expr::Identifier(Ident::new("a"))],
2031+
vec![
2032+
Expr::Identifier(Ident::new("b")),
2033+
Expr::Identifier(Ident::new("c")),
2034+
],
2035+
vec![Expr::Identifier(Ident::new("d"))],
2036+
]);
2037+
assert_eq!("CUBE (a, (b, c), d)", format!("{}", cube));
2038+
}
19062039
}

‎src/keywords.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -406,6 +406,7 @@ define_keywords!(
406406
SESSION,
407407
SESSION_USER,
408408
SET,
409+
SETS,
409410
SHOW,
410411
SIMILAR,
411412
SMALLINT,

‎src/parser.rs‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -338,7 +338,7 @@ impl<'a> Parser<'a> {
338338
return_ok_if_some!(self.maybe_parse(|parser| {
339339
match parser.parse_data_type()? {
340340
DataType::Interval => parser.parse_literal_interval(),
341-
// PosgreSQL allows almost any identifier to be used as custom data type name,
341+
// PostgreSQL allows almost any identifier to be used as custom data type name,
342342
// and we support that in `parse_data_type()`. But unlike Postgres we don't
343343
// have a list of globally reserved keywords (since they vary across dialects),
344344
// so given `NOT 'a' LIKE 'b'`, we'd accept `NOT` as a possible custom data type
@@ -559,6 +559,63 @@ impl<'a> Parser<'a> {
559559
}
560560
}
561561

562+
/// parse a group by expr. a group by expr can be one of group sets, roll up, cube, or simple
563+
/// expr.
564+
fn parse_group_by_expr(&mut self) -> Result<Expr, ParserError> {
565+
if self.parse_keywords(&[Keyword::GROUPING, Keyword::SETS]) {
566+
self.expect_token(&Token::LParen)?;
567+
let result = self.parse_comma_separated(|p| p.parse_tuple(false, true))?;
568+
self.expect_token(&Token::RParen)?;
569+
Ok(Expr::GroupingSets(result))
570+
} else if self.parse_keyword(Keyword::CUBE) {
571+
self.expect_token(&Token::LParen)?;
572+
let result = self.parse_comma_separated(|p| p.parse_tuple(true, true))?;
573+
self.expect_token(&Token::RParen)?;
574+
Ok(Expr::Cube(result))
575+
} else if self.parse_keyword(Keyword::ROLLUP) {
576+
self.expect_token(&Token::LParen)?;
577+
let result = self.parse_comma_separated(|p| p.parse_tuple(true, true))?;
578+
self.expect_token(&Token::RParen)?;
579+
Ok(Expr::Rollup(result))
580+
} else {
581+
self.parse_expr()
582+
}
583+
}
584+
585+
/// parse a tuple with `(` and `)`.
586+
/// If `lift_singleton` is true, then a singleton tuple is lifted to a tuple of length 1, otherwise it will fail.
587+
/// If `allow_empty` is true, then an empty tuple is allowed.
588+
fn parse_tuple(
589+
&mut self,
590+
lift_singleton: bool,
591+
allow_empty: bool,
592+
) -> Result<Vec<Expr>, ParserError> {
593+
if lift_singleton {
594+
if self.consume_token(&Token::LParen) {
595+
let result = if allow_empty && self.consume_token(&Token::RParen) {
596+
vec![]
597+
} else {
598+
let result = self.parse_comma_separated(Parser::parse_expr)?;
599+
self.expect_token(&Token::RParen)?;
600+
result
601+
};
602+
Ok(result)
603+
} else {
604+
Ok(vec![self.parse_expr()?])
605+
}
606+
} else {
607+
self.expect_token(&Token::LParen)?;
608+
let result = if allow_empty && self.consume_token(&Token::RParen) {
609+
vec![]
610+
} else {
611+
let result = self.parse_comma_separated(Parser::parse_expr)?;
612+
self.expect_token(&Token::RParen)?;
613+
result
614+
};
615+
Ok(result)
616+
}
617+
}
618+
562619
pub fn parse_case_expr(&mut self) -> Result<Expr, ParserError> {
563620
let mut operand = None;
564621
if !self.parse_keyword(Keyword::WHEN) {
@@ -2494,7 +2551,7 @@ impl<'a> Parser<'a> {
24942551
};
24952552

24962553
let group_by = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
2497-
self.parse_comma_separated(Parser::parse_expr)?
2554+
self.parse_comma_separated(Parser::parse_group_by_expr)?
24982555
} else {
24992556
vec![]
25002557
};

‎tests/sqlparser_common.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1039,6 +1039,24 @@ fn parse_select_group_by() {
10391039
);
10401040
}
10411041

1042+
#[test]
1043+
fn parse_select_group_by_grouping_sets() {
1044+
let sql =
1045+
"SELECT brand, size, sum(sales) FROM items_sold GROUP BY size, GROUPING SETS ((brand), (size), ())";
1046+
let select = verified_only_select(sql);
1047+
assert_eq!(
1048+
vec![
1049+
Expr::Identifier(Ident::new("size")),
1050+
Expr::GroupingSets(vec![
1051+
vec![Expr::Identifier(Ident::new("brand"))],
1052+
vec![Expr::Identifier(Ident::new("size"))],
1053+
vec![],
1054+
])
1055+
],
1056+
select.group_by
1057+
);
1058+
}
1059+
10421060
#[test]
10431061
fn parse_select_having() {
10441062
let sql = "SELECT foo FROM bar GROUP BY foo HAVING COUNT(*) > 1";

0 commit comments

Comments
 (0)