Skip to content

Commit 44d7a20

Browse files
git-hulkiffyio
andauthored
Support GROUP BY WITH MODIFIER for ClickHouse (#1323)
Co-authored-by: Ifeanyi Ubah <[email protected]>
1 parent 0b1a413 commit 44d7a20

File tree

10 files changed

+215
-76
lines changed

10 files changed

+215
-76
lines changed

src/ast/mod.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,8 @@ pub use self::operator::{BinaryOperator, UnaryOperator};
4343
pub use self::query::{
4444
AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
4545
ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
46-
GroupByExpr, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator,
47-
JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
46+
GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint,
47+
JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
4848
MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
4949
NonBlock, Offset, OffsetRows, OrderByExpr, PivotValueSource, Query, RenameSelectItem,
5050
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,

src/ast/query.rs

Lines changed: 47 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -299,10 +299,10 @@ impl fmt::Display for Select {
299299
write!(f, " WHERE {selection}")?;
300300
}
301301
match &self.group_by {
302-
GroupByExpr::All => write!(f, " GROUP BY ALL")?,
303-
GroupByExpr::Expressions(exprs) => {
302+
GroupByExpr::All(_) => write!(f, " {}", self.group_by)?,
303+
GroupByExpr::Expressions(exprs, _) => {
304304
if !exprs.is_empty() {
305-
write!(f, " GROUP BY {}", display_comma_separated(exprs))?;
305+
write!(f, " {}", self.group_by)?
306306
}
307307
}
308308
}
@@ -1866,27 +1866,65 @@ impl fmt::Display for SelectInto {
18661866
}
18671867
}
18681868

1869+
/// ClickHouse supports GROUP BY WITH modifiers(includes ROLLUP|CUBE|TOTALS).
1870+
/// e.g. GROUP BY year WITH ROLLUP WITH TOTALS
1871+
///
1872+
/// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#rollup-modifier>
1873+
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
1874+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
1875+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
1876+
pub enum GroupByWithModifier {
1877+
Rollup,
1878+
Cube,
1879+
Totals,
1880+
}
1881+
1882+
impl fmt::Display for GroupByWithModifier {
1883+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
1884+
match self {
1885+
GroupByWithModifier::Rollup => write!(f, "WITH ROLLUP"),
1886+
GroupByWithModifier::Cube => write!(f, "WITH CUBE"),
1887+
GroupByWithModifier::Totals => write!(f, "WITH TOTALS"),
1888+
}
1889+
}
1890+
}
1891+
18691892
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
18701893
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
18711894
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
18721895
pub enum GroupByExpr {
1873-
/// ALL syntax of [Snowflake], and [DuckDB]
1896+
/// ALL syntax of [Snowflake], [DuckDB] and [ClickHouse].
18741897
///
18751898
/// [Snowflake]: <https://docs.snowflake.com/en/sql-reference/constructs/group-by#label-group-by-all-columns>
18761899
/// [DuckDB]: <https://duckdb.org/docs/sql/query_syntax/groupby.html>
1877-
All,
1900+
/// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#group-by-all>
1901+
///
1902+
/// ClickHouse also supports WITH modifiers after GROUP BY ALL and expressions.
1903+
///
1904+
/// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#rollup-modifier>
1905+
All(Vec<GroupByWithModifier>),
18781906

18791907
/// Expressions
1880-
Expressions(Vec<Expr>),
1908+
Expressions(Vec<Expr>, Vec<GroupByWithModifier>),
18811909
}
18821910

18831911
impl fmt::Display for GroupByExpr {
18841912
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
18851913
match self {
1886-
GroupByExpr::All => write!(f, "GROUP BY ALL"),
1887-
GroupByExpr::Expressions(col_names) => {
1914+
GroupByExpr::All(modifiers) => {
1915+
write!(f, "GROUP BY ALL")?;
1916+
if !modifiers.is_empty() {
1917+
write!(f, " {}", display_separated(modifiers, " "))?;
1918+
}
1919+
Ok(())
1920+
}
1921+
GroupByExpr::Expressions(col_names, modifiers) => {
18881922
let col_names = display_comma_separated(col_names);
1889-
write!(f, "GROUP BY ({col_names})")
1923+
write!(f, "GROUP BY {col_names}")?;
1924+
if !modifiers.is_empty() {
1925+
write!(f, " {}", display_separated(modifiers, " "))?;
1926+
}
1927+
Ok(())
18901928
}
18911929
}
18921930
}

src/keywords.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -721,6 +721,7 @@ define_keywords!(
721721
TINYINT,
722722
TO,
723723
TOP,
724+
TOTALS,
724725
TRAILING,
725726
TRANSACTION,
726727
TRANSIENT,

src/parser/mod.rs

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8319,13 +8319,42 @@ impl<'a> Parser<'a> {
83198319
};
83208320

83218321
let group_by = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
8322-
if self.parse_keyword(Keyword::ALL) {
8323-
GroupByExpr::All
8322+
let expressions = if self.parse_keyword(Keyword::ALL) {
8323+
None
83248324
} else {
8325-
GroupByExpr::Expressions(self.parse_comma_separated(Parser::parse_group_by_expr)?)
8325+
Some(self.parse_comma_separated(Parser::parse_group_by_expr)?)
8326+
};
8327+
8328+
let mut modifiers = vec![];
8329+
if dialect_of!(self is ClickHouseDialect | GenericDialect) {
8330+
loop {
8331+
if !self.parse_keyword(Keyword::WITH) {
8332+
break;
8333+
}
8334+
let keyword = self.expect_one_of_keywords(&[
8335+
Keyword::ROLLUP,
8336+
Keyword::CUBE,
8337+
Keyword::TOTALS,
8338+
])?;
8339+
modifiers.push(match keyword {
8340+
Keyword::ROLLUP => GroupByWithModifier::Rollup,
8341+
Keyword::CUBE => GroupByWithModifier::Cube,
8342+
Keyword::TOTALS => GroupByWithModifier::Totals,
8343+
_ => {
8344+
return parser_err!(
8345+
"BUG: expected to match GroupBy modifier keyword",
8346+
self.peek_token().location
8347+
)
8348+
}
8349+
});
8350+
}
8351+
}
8352+
match expressions {
8353+
None => GroupByExpr::All(modifiers),
8354+
Some(exprs) => GroupByExpr::Expressions(exprs, modifiers),
83268355
}
83278356
} else {
8328-
GroupByExpr::Expressions(vec![])
8357+
GroupByExpr::Expressions(vec![], vec![])
83298358
};
83308359

83318360
let cluster_by = if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {

tests/sqlparser_clickhouse.rs

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ fn parse_map_access_expr() {
8888
right: Box::new(Expr::Value(Value::SingleQuotedString("foo".to_string()))),
8989
}),
9090
}),
91-
group_by: GroupByExpr::Expressions(vec![]),
91+
group_by: GroupByExpr::Expressions(vec![], vec![]),
9292
cluster_by: vec![],
9393
distribute_by: vec![],
9494
sort_by: vec![],
@@ -626,6 +626,61 @@ fn parse_create_materialized_view() {
626626
clickhouse_and_generic().verified_stmt(sql);
627627
}
628628

629+
#[test]
630+
fn parse_group_by_with_modifier() {
631+
let clauses = ["x", "a, b", "ALL"];
632+
let modifiers = [
633+
"WITH ROLLUP",
634+
"WITH CUBE",
635+
"WITH TOTALS",
636+
"WITH ROLLUP WITH CUBE",
637+
];
638+
let expected_modifiers = [
639+
vec![GroupByWithModifier::Rollup],
640+
vec![GroupByWithModifier::Cube],
641+
vec![GroupByWithModifier::Totals],
642+
vec![GroupByWithModifier::Rollup, GroupByWithModifier::Cube],
643+
];
644+
for clause in &clauses {
645+
for (modifier, expected_modifier) in modifiers.iter().zip(expected_modifiers.iter()) {
646+
let sql = format!("SELECT * FROM t GROUP BY {clause} {modifier}");
647+
match clickhouse_and_generic().verified_stmt(&sql) {
648+
Statement::Query(query) => {
649+
let group_by = &query.body.as_select().unwrap().group_by;
650+
if clause == &"ALL" {
651+
assert_eq!(group_by, &GroupByExpr::All(expected_modifier.to_vec()));
652+
} else {
653+
assert_eq!(
654+
group_by,
655+
&GroupByExpr::Expressions(
656+
clause
657+
.split(", ")
658+
.map(|c| Identifier(Ident::new(c)))
659+
.collect(),
660+
expected_modifier.to_vec()
661+
)
662+
);
663+
}
664+
}
665+
_ => unreachable!(),
666+
}
667+
}
668+
}
669+
670+
// invalid cases
671+
let invalid_cases = [
672+
"SELECT * FROM t GROUP BY x WITH",
673+
"SELECT * FROM t GROUP BY x WITH ROLLUP CUBE",
674+
"SELECT * FROM t GROUP BY x WITH WITH ROLLUP",
675+
"SELECT * FROM t GROUP BY WITH ROLLUP",
676+
];
677+
for sql in invalid_cases {
678+
clickhouse_and_generic()
679+
.parse_sql_statements(sql)
680+
.expect_err("Expected: one of ROLLUP or CUBE or TOTALS, found: WITH");
681+
}
682+
}
683+
629684
fn clickhouse() -> TestedDialects {
630685
TestedDialects {
631686
dialects: vec![Box::new(ClickHouseDialect {})],

tests/sqlparser_common.rs

Lines changed: 30 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -392,9 +392,10 @@ fn parse_update_set_from() {
392392
}],
393393
lateral_views: vec![],
394394
selection: None,
395-
group_by: GroupByExpr::Expressions(vec![Expr::Identifier(Ident::new(
396-
"id"
397-
))]),
395+
group_by: GroupByExpr::Expressions(
396+
vec![Expr::Identifier(Ident::new("id"))],
397+
vec![]
398+
),
398399
cluster_by: vec![],
399400
distribute_by: vec![],
400401
sort_by: vec![],
@@ -2119,10 +2120,13 @@ fn parse_select_group_by() {
21192120
let sql = "SELECT id, fname, lname FROM customer GROUP BY lname, fname";
21202121
let select = verified_only_select(sql);
21212122
assert_eq!(
2122-
GroupByExpr::Expressions(vec![
2123-
Expr::Identifier(Ident::new("lname")),
2124-
Expr::Identifier(Ident::new("fname")),
2125-
]),
2123+
GroupByExpr::Expressions(
2124+
vec![
2125+
Expr::Identifier(Ident::new("lname")),
2126+
Expr::Identifier(Ident::new("fname")),
2127+
],
2128+
vec![]
2129+
),
21262130
select.group_by
21272131
);
21282132

@@ -2137,7 +2141,7 @@ fn parse_select_group_by() {
21372141
fn parse_select_group_by_all() {
21382142
let sql = "SELECT id, fname, lname, SUM(order) FROM customer GROUP BY ALL";
21392143
let select = verified_only_select(sql);
2140-
assert_eq!(GroupByExpr::All, select.group_by);
2144+
assert_eq!(GroupByExpr::All(vec![]), select.group_by);
21412145

21422146
one_statement_parses_to(
21432147
"SELECT id, fname, lname, SUM(order) FROM customer GROUP BY ALL",
@@ -4545,7 +4549,7 @@ fn test_parse_named_window() {
45454549
}],
45464550
lateral_views: vec![],
45474551
selection: None,
4548-
group_by: GroupByExpr::Expressions(vec![]),
4552+
group_by: GroupByExpr::Expressions(vec![], vec![]),
45494553
cluster_by: vec![],
45504554
distribute_by: vec![],
45514555
sort_by: vec![],
@@ -4974,7 +4978,7 @@ fn parse_interval_and_or_xor() {
49744978
}),
49754979
}),
49764980
}),
4977-
group_by: GroupByExpr::Expressions(vec![]),
4981+
group_by: GroupByExpr::Expressions(vec![], vec![]),
49784982
cluster_by: vec![],
49794983
distribute_by: vec![],
49804984
sort_by: vec![],
@@ -6908,7 +6912,7 @@ fn lateral_function() {
69086912
}],
69096913
lateral_views: vec![],
69106914
selection: None,
6911-
group_by: GroupByExpr::Expressions(vec![]),
6915+
group_by: GroupByExpr::Expressions(vec![], vec![]),
69126916
cluster_by: vec![],
69136917
distribute_by: vec![],
69146918
sort_by: vec![],
@@ -7627,7 +7631,7 @@ fn parse_merge() {
76277631
}],
76287632
lateral_views: vec![],
76297633
selection: None,
7630-
group_by: GroupByExpr::Expressions(vec![]),
7634+
group_by: GroupByExpr::Expressions(vec![], vec![]),
76317635
cluster_by: vec![],
76327636
distribute_by: vec![],
76337637
sort_by: vec![],
@@ -9133,7 +9137,7 @@ fn parse_unload() {
91339137
}],
91349138
lateral_views: vec![],
91359139
selection: None,
9136-
group_by: GroupByExpr::Expressions(vec![]),
9140+
group_by: GroupByExpr::Expressions(vec![], vec![]),
91379141
cluster_by: vec![],
91389142
distribute_by: vec![],
91399143
sort_by: vec![],
@@ -9276,7 +9280,7 @@ fn parse_connect_by() {
92769280
into: None,
92779281
lateral_views: vec![],
92789282
selection: None,
9279-
group_by: GroupByExpr::Expressions(vec![]),
9283+
group_by: GroupByExpr::Expressions(vec![], vec![]),
92809284
cluster_by: vec![],
92819285
distribute_by: vec![],
92829286
sort_by: vec![],
@@ -9364,7 +9368,7 @@ fn parse_connect_by() {
93649368
op: BinaryOperator::NotEq,
93659369
right: Box::new(Expr::Value(number("42"))),
93669370
}),
9367-
group_by: GroupByExpr::Expressions(vec![]),
9371+
group_by: GroupByExpr::Expressions(vec![], vec![]),
93689372
cluster_by: vec![],
93699373
distribute_by: vec![],
93709374
sort_by: vec![],
@@ -9484,15 +9488,18 @@ fn test_group_by_grouping_sets() {
94849488
all_dialects_where(|d| d.supports_group_by_expr())
94859489
.verified_only_select(sql)
94869490
.group_by,
9487-
GroupByExpr::Expressions(vec![Expr::GroupingSets(vec![
9488-
vec![
9489-
Expr::Identifier(Ident::new("city")),
9490-
Expr::Identifier(Ident::new("car_model"))
9491-
],
9492-
vec![Expr::Identifier(Ident::new("city")),],
9493-
vec![Expr::Identifier(Ident::new("car_model"))],
9491+
GroupByExpr::Expressions(
9492+
vec![Expr::GroupingSets(vec![
9493+
vec![
9494+
Expr::Identifier(Ident::new("city")),
9495+
Expr::Identifier(Ident::new("car_model"))
9496+
],
9497+
vec![Expr::Identifier(Ident::new("city")),],
9498+
vec![Expr::Identifier(Ident::new("car_model"))],
9499+
vec![]
9500+
])],
94949501
vec![]
9495-
])])
9502+
)
94969503
);
94979504
}
94989505

tests/sqlparser_duckdb.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ fn test_select_union_by_name() {
171171
}],
172172
lateral_views: vec![],
173173
selection: None,
174-
group_by: GroupByExpr::Expressions(vec![]),
174+
group_by: GroupByExpr::Expressions(vec![], vec![]),
175175
cluster_by: vec![],
176176
distribute_by: vec![],
177177
sort_by: vec![],
@@ -209,7 +209,7 @@ fn test_select_union_by_name() {
209209
}],
210210
lateral_views: vec![],
211211
selection: None,
212-
group_by: GroupByExpr::Expressions(vec![]),
212+
group_by: GroupByExpr::Expressions(vec![], vec![]),
213213
cluster_by: vec![],
214214
distribute_by: vec![],
215215
sort_by: vec![],

tests/sqlparser_mssql.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ fn parse_create_procedure() {
111111
from: vec![],
112112
lateral_views: vec![],
113113
selection: None,
114-
group_by: GroupByExpr::Expressions(vec![]),
114+
group_by: GroupByExpr::Expressions(vec![], vec![]),
115115
cluster_by: vec![],
116116
distribute_by: vec![],
117117
sort_by: vec![],
@@ -528,7 +528,7 @@ fn parse_substring_in_select() {
528528
}],
529529
lateral_views: vec![],
530530
selection: None,
531-
group_by: GroupByExpr::Expressions(vec![]),
531+
group_by: GroupByExpr::Expressions(vec![], vec![]),
532532
cluster_by: vec![],
533533
distribute_by: vec![],
534534
sort_by: vec![],

0 commit comments

Comments
 (0)