Skip to content

Support GROUP BY WITH MODIFIER for ClickHouse #1323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jun 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,8 @@ pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
GroupByExpr, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint, JoinOperator,
JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem, Join, JoinConstraint,
JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
MatchRecognizePattern, MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr,
NonBlock, Offset, OffsetRows, OrderByExpr, PivotValueSource, Query, RenameSelectItem,
RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
Expand Down
56 changes: 47 additions & 9 deletions src/ast/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -299,10 +299,10 @@ impl fmt::Display for Select {
write!(f, " WHERE {selection}")?;
}
match &self.group_by {
GroupByExpr::All => write!(f, " GROUP BY ALL")?,
GroupByExpr::Expressions(exprs) => {
GroupByExpr::All(_) => write!(f, " {}", self.group_by)?,
GroupByExpr::Expressions(exprs, _) => {
if !exprs.is_empty() {
write!(f, " GROUP BY {}", display_comma_separated(exprs))?;
write!(f, " {}", self.group_by)?
}
}
}
Expand Down Expand Up @@ -1865,27 +1865,65 @@ impl fmt::Display for SelectInto {
}
}

/// ClickHouse supports GROUP BY WITH modifiers(includes ROLLUP|CUBE|TOTALS).
/// e.g. GROUP BY year WITH ROLLUP WITH TOTALS
///
/// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#rollup-modifier>
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum GroupByWithModifier {
Rollup,
Cube,
Totals,
}

impl fmt::Display for GroupByWithModifier {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
GroupByWithModifier::Rollup => write!(f, "WITH ROLLUP"),
GroupByWithModifier::Cube => write!(f, "WITH CUBE"),
GroupByWithModifier::Totals => write!(f, "WITH TOTALS"),
}
}
}

#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum GroupByExpr {
/// ALL syntax of [Snowflake], and [DuckDB]
/// ALL syntax of [Snowflake], [DuckDB] and [ClickHouse].
///
/// [Snowflake]: <https://docs.snowflake.com/en/sql-reference/constructs/group-by#label-group-by-all-columns>
/// [DuckDB]: <https://duckdb.org/docs/sql/query_syntax/groupby.html>
All,
/// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#group-by-all>
///
/// ClickHouse also supports WITH modifiers after GROUP BY ALL and expressions.
///
/// [ClickHouse]: <https://clickhouse.com/docs/en/sql-reference/statements/select/group-by#rollup-modifier>
All(Vec<GroupByWithModifier>),

/// Expressions
Expressions(Vec<Expr>),
Expressions(Vec<Expr>, Vec<GroupByWithModifier>),
}

impl fmt::Display for GroupByExpr {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
GroupByExpr::All => write!(f, "GROUP BY ALL"),
GroupByExpr::Expressions(col_names) => {
GroupByExpr::All(modifiers) => {
write!(f, "GROUP BY ALL")?;
if !modifiers.is_empty() {
write!(f, " {}", display_separated(modifiers, " "))?;
}
Ok(())
}
GroupByExpr::Expressions(col_names, modifiers) => {
let col_names = display_comma_separated(col_names);
write!(f, "GROUP BY ({col_names})")
write!(f, "GROUP BY {col_names}")?;
if !modifiers.is_empty() {
write!(f, " {}", display_separated(modifiers, " "))?;
}
Ok(())
}
}
}
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -721,6 +721,7 @@ define_keywords!(
TINYINT,
TO,
TOP,
TOTALS,
TRAILING,
TRANSACTION,
TRANSIENT,
Expand Down
37 changes: 33 additions & 4 deletions src/parser/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8287,13 +8287,42 @@ impl<'a> Parser<'a> {
};

let group_by = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY]) {
if self.parse_keyword(Keyword::ALL) {
GroupByExpr::All
let expressions = if self.parse_keyword(Keyword::ALL) {
None
} else {
GroupByExpr::Expressions(self.parse_comma_separated(Parser::parse_group_by_expr)?)
Some(self.parse_comma_separated(Parser::parse_group_by_expr)?)
};

let mut modifiers = vec![];
if dialect_of!(self is ClickHouseDialect | GenericDialect) {
loop {
if !self.parse_keyword(Keyword::WITH) {
break;
}
let keyword = self.expect_one_of_keywords(&[
Keyword::ROLLUP,
Keyword::CUBE,
Keyword::TOTALS,
])?;
modifiers.push(match keyword {
Keyword::ROLLUP => GroupByWithModifier::Rollup,
Keyword::CUBE => GroupByWithModifier::Cube,
Keyword::TOTALS => GroupByWithModifier::Totals,
_ => {
return parser_err!(
"BUG: expected to match GroupBy modifier keyword",
self.peek_token().location
)
}
});
}
}
match expressions {
None => GroupByExpr::All(modifiers),
Some(exprs) => GroupByExpr::Expressions(exprs, modifiers),
}
} else {
GroupByExpr::Expressions(vec![])
GroupByExpr::Expressions(vec![], vec![])
};

let cluster_by = if self.parse_keywords(&[Keyword::CLUSTER, Keyword::BY]) {
Expand Down
57 changes: 56 additions & 1 deletion tests/sqlparser_clickhouse.rs
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ fn parse_map_access_expr() {
right: Box::new(Expr::Value(Value::SingleQuotedString("foo".to_string()))),
}),
}),
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -626,6 +626,61 @@ fn parse_create_materialized_view() {
clickhouse_and_generic().verified_stmt(sql);
}

#[test]
fn parse_group_by_with_modifier() {
let clauses = ["x", "a, b", "ALL"];
let modifiers = [
"WITH ROLLUP",
"WITH CUBE",
"WITH TOTALS",
"WITH ROLLUP WITH CUBE",
];
let expected_modifiers = [
vec![GroupByWithModifier::Rollup],
vec![GroupByWithModifier::Cube],
vec![GroupByWithModifier::Totals],
vec![GroupByWithModifier::Rollup, GroupByWithModifier::Cube],
];
for clause in &clauses {
for (modifier, expected_modifier) in modifiers.iter().zip(expected_modifiers.iter()) {
let sql = format!("SELECT * FROM t GROUP BY {clause} {modifier}");
match clickhouse_and_generic().verified_stmt(&sql) {
Statement::Query(query) => {
let group_by = &query.body.as_select().unwrap().group_by;
if clause == &"ALL" {
assert_eq!(group_by, &GroupByExpr::All(expected_modifier.to_vec()));
} else {
assert_eq!(
group_by,
&GroupByExpr::Expressions(
clause
.split(", ")
.map(|c| Identifier(Ident::new(c)))
.collect(),
expected_modifier.to_vec()
)
);
}
}
_ => unreachable!(),
}
}
}

// invalid cases
let invalid_cases = [
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💯 for negative cases

Copy link
Member Author

@git-hulk git-hulk Jun 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be credited to the guidance of @iffyio, thank you!

"SELECT * FROM t GROUP BY x WITH",
"SELECT * FROM t GROUP BY x WITH ROLLUP CUBE",
"SELECT * FROM t GROUP BY x WITH WITH ROLLUP",
"SELECT * FROM t GROUP BY WITH ROLLUP",
];
for sql in invalid_cases {
clickhouse_and_generic()
.parse_sql_statements(sql)
.expect_err("Expected: one of ROLLUP or CUBE or TOTALS, found: WITH");
}
}

fn clickhouse() -> TestedDialects {
TestedDialects {
dialects: vec![Box::new(ClickHouseDialect {})],
Expand Down
53 changes: 30 additions & 23 deletions tests/sqlparser_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -392,9 +392,10 @@ fn parse_update_set_from() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![Expr::Identifier(Ident::new(
"id"
))]),
group_by: GroupByExpr::Expressions(
vec![Expr::Identifier(Ident::new("id"))],
vec![]
),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -2119,10 +2120,13 @@ fn parse_select_group_by() {
let sql = "SELECT id, fname, lname FROM customer GROUP BY lname, fname";
let select = verified_only_select(sql);
assert_eq!(
GroupByExpr::Expressions(vec![
Expr::Identifier(Ident::new("lname")),
Expr::Identifier(Ident::new("fname")),
]),
GroupByExpr::Expressions(
vec![
Expr::Identifier(Ident::new("lname")),
Expr::Identifier(Ident::new("fname")),
],
vec![]
),
select.group_by
);

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

one_statement_parses_to(
"SELECT id, fname, lname, SUM(order) FROM customer GROUP BY ALL",
Expand Down Expand Up @@ -4545,7 +4549,7 @@ fn test_parse_named_window() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -4974,7 +4978,7 @@ fn parse_interval_and_or_xor() {
}),
}),
}),
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -6908,7 +6912,7 @@ fn lateral_function() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -7627,7 +7631,7 @@ fn parse_merge() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -9133,7 +9137,7 @@ fn parse_unload() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -9276,7 +9280,7 @@ fn parse_connect_by() {
into: None,
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -9364,7 +9368,7 @@ fn parse_connect_by() {
op: BinaryOperator::NotEq,
right: Box::new(Expr::Value(number("42"))),
}),
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -9484,15 +9488,18 @@ fn test_group_by_grouping_sets() {
all_dialects_where(|d| d.supports_group_by_expr())
.verified_only_select(sql)
.group_by,
GroupByExpr::Expressions(vec![Expr::GroupingSets(vec![
vec![
Expr::Identifier(Ident::new("city")),
Expr::Identifier(Ident::new("car_model"))
],
vec![Expr::Identifier(Ident::new("city")),],
vec![Expr::Identifier(Ident::new("car_model"))],
GroupByExpr::Expressions(
vec![Expr::GroupingSets(vec![
vec![
Expr::Identifier(Ident::new("city")),
Expr::Identifier(Ident::new("car_model"))
],
vec![Expr::Identifier(Ident::new("city")),],
vec![Expr::Identifier(Ident::new("car_model"))],
vec![]
])],
vec![]
])])
)
);
}

Expand Down
4 changes: 2 additions & 2 deletions tests/sqlparser_duckdb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,7 +171,7 @@ fn test_select_union_by_name() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -209,7 +209,7 @@ fn test_select_union_by_name() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down
4 changes: 2 additions & 2 deletions tests/sqlparser_mssql.rs
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ fn parse_create_procedure() {
from: vec![],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down Expand Up @@ -528,7 +528,7 @@ fn parse_substring_in_select() {
}],
lateral_views: vec![],
selection: None,
group_by: GroupByExpr::Expressions(vec![]),
group_by: GroupByExpr::Expressions(vec![], vec![]),
cluster_by: vec![],
distribute_by: vec![],
sort_by: vec![],
Expand Down
Loading
Loading