Skip to content

Commit bc898b3

Browse files
BohuTANGBohuTANG
authored
Merge pull request #10601 from RinChanNOWWW/rollup-cube
feat: support `CUBE` and `ROLLUP`.
2 parents 70e1cf9 + 38674f6 commit bc898b3

File tree

10 files changed

+463
-62
lines changed

10 files changed

+463
-62
lines changed

‎src/query/ast/src/ast/format/ast_format.rs‎

Lines changed: 66 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2074,45 +2074,76 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
20742074
FormatTreeNode::with_children(selection_format_ctx, vec![selection_child]);
20752075
children.push(selection_node);
20762076
}
2077-
match &stmt.group_by {
2078-
Some(GroupBy::Normal(exprs)) => {
2079-
let mut group_by_list_children = Vec::with_capacity(exprs.len());
2080-
for group_by in exprs.iter() {
2081-
self.visit_expr(group_by);
2082-
group_by_list_children.push(self.children.pop().unwrap());
2077+
if let Some(group_by) = &stmt.group_by {
2078+
match group_by {
2079+
GroupBy::Normal(exprs) => {
2080+
let mut group_by_list_children = Vec::with_capacity(exprs.len());
2081+
for group_by in exprs.iter() {
2082+
self.visit_expr(group_by);
2083+
group_by_list_children.push(self.children.pop().unwrap());
2084+
}
2085+
let group_by_list_name = "GroupByList".to_string();
2086+
let group_by_list_format_ctx = AstFormatContext::with_children(
2087+
group_by_list_name,
2088+
group_by_list_children.len(),
2089+
);
2090+
let group_by_list_node = FormatTreeNode::with_children(
2091+
group_by_list_format_ctx,
2092+
group_by_list_children,
2093+
);
2094+
children.push(group_by_list_node);
20832095
}
2084-
let group_by_list_name = "GroupByList".to_string();
2085-
let group_by_list_format_ctx = AstFormatContext::with_children(
2086-
group_by_list_name,
2087-
group_by_list_children.len(),
2088-
);
2089-
let group_by_list_node =
2090-
FormatTreeNode::with_children(group_by_list_format_ctx, group_by_list_children);
2091-
children.push(group_by_list_node);
2092-
}
2093-
Some(GroupBy::GroupingSets(sets)) => {
2094-
let mut grouping_sets = Vec::with_capacity(sets.len());
2095-
for set in sets.iter() {
2096-
let mut grouping_set = Vec::with_capacity(set.len());
2097-
for expr in set.iter() {
2098-
self.visit_expr(expr);
2099-
grouping_set.push(self.children.pop().unwrap());
2096+
GroupBy::GroupingSets(sets) => {
2097+
let mut grouping_sets = Vec::with_capacity(sets.len());
2098+
for set in sets.iter() {
2099+
let mut grouping_set = Vec::with_capacity(set.len());
2100+
for expr in set.iter() {
2101+
self.visit_expr(expr);
2102+
grouping_set.push(self.children.pop().unwrap());
2103+
}
2104+
let name = "GroupingSet".to_string();
2105+
let grouping_set_format_ctx =
2106+
AstFormatContext::with_children(name, grouping_set.len());
2107+
let grouping_set_node =
2108+
FormatTreeNode::with_children(grouping_set_format_ctx, grouping_set);
2109+
grouping_sets.push(grouping_set_node);
2110+
}
2111+
let group_by_list_name = "GroupByList".to_string();
2112+
let group_by_list_format_ctx =
2113+
AstFormatContext::with_children(group_by_list_name, grouping_sets.len());
2114+
let group_by_list_node =
2115+
FormatTreeNode::with_children(group_by_list_format_ctx, grouping_sets);
2116+
children.push(group_by_list_node);
2117+
}
2118+
GroupBy::Rollup(exprs) => {
2119+
let mut rollup_list_children = Vec::with_capacity(exprs.len());
2120+
for group_by in exprs.iter() {
2121+
self.visit_expr(group_by);
2122+
rollup_list_children.push(self.children.pop().unwrap());
2123+
}
2124+
let rollup_list_name = "GroupByRollUpList".to_string();
2125+
let rollup_list_format_ctx = AstFormatContext::with_children(
2126+
rollup_list_name,
2127+
rollup_list_children.len(),
2128+
);
2129+
let rollup_list_node =
2130+
FormatTreeNode::with_children(rollup_list_format_ctx, rollup_list_children);
2131+
children.push(rollup_list_node);
2132+
}
2133+
GroupBy::Cube(exprs) => {
2134+
let mut cube_list_children = Vec::with_capacity(exprs.len());
2135+
for group_by in exprs.iter() {
2136+
self.visit_expr(group_by);
2137+
cube_list_children.push(self.children.pop().unwrap());
21002138
}
2101-
let name = "GroupingSet".to_string();
2102-
let grouping_set_format_ctx =
2103-
AstFormatContext::with_children(name, grouping_set.len());
2104-
let grouping_set_node =
2105-
FormatTreeNode::with_children(grouping_set_format_ctx, grouping_set);
2106-
grouping_sets.push(grouping_set_node);
2139+
let cube_list_name = "GroupByCubeList".to_string();
2140+
let cube_list_format_ctx =
2141+
AstFormatContext::with_children(cube_list_name, cube_list_children.len());
2142+
let cube_list_node =
2143+
FormatTreeNode::with_children(cube_list_format_ctx, cube_list_children);
2144+
children.push(cube_list_node);
21072145
}
2108-
let group_by_list_name = "GroupByList".to_string();
2109-
let group_by_list_format_ctx =
2110-
AstFormatContext::with_children(group_by_list_name, grouping_sets.len());
2111-
let group_by_list_node =
2112-
FormatTreeNode::with_children(group_by_list_format_ctx, grouping_sets);
2113-
children.push(group_by_list_node);
21142146
}
2115-
_ => {}
21162147
}
21172148

21182149
if let Some(having) = &stmt.having {

‎src/query/ast/src/ast/format/syntax/query.rs‎

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -203,33 +203,56 @@ fn pretty_group_set(set: Vec<Expr>) -> RcDoc<'static> {
203203
}
204204

205205
fn pretty_group_by(group_by: Option<GroupBy>) -> RcDoc<'static> {
206-
match group_by {
207-
Some(GroupBy::Normal(exprs)) => RcDoc::line()
208-
.append(
209-
RcDoc::text("GROUP BY").append(
210-
if exprs.len() > 1 {
211-
RcDoc::line()
212-
} else {
213-
RcDoc::space()
214-
}
215-
.nest(NEST_FACTOR),
206+
if let Some(group_by) = group_by {
207+
match group_by {
208+
GroupBy::Normal(exprs) => RcDoc::line()
209+
.append(
210+
RcDoc::text("GROUP BY").append(
211+
if exprs.len() > 1 {
212+
RcDoc::line()
213+
} else {
214+
RcDoc::space()
215+
}
216+
.nest(NEST_FACTOR),
217+
),
218+
)
219+
.append(
220+
interweave_comma(exprs.into_iter().map(pretty_expr))
221+
.nest(NEST_FACTOR)
222+
.group(),
216223
),
217-
)
218-
.append(
219-
interweave_comma(exprs.into_iter().map(pretty_expr))
220-
.nest(NEST_FACTOR)
221-
.group(),
222-
),
223-
Some(GroupBy::GroupingSets(sets)) => RcDoc::line()
224-
.append(RcDoc::text("GROUP BY GROUPING SETS (").append(RcDoc::line().nest(NEST_FACTOR)))
225-
.append(
226-
interweave_comma(sets.into_iter().map(pretty_group_set))
227-
.nest(NEST_FACTOR)
228-
.group(),
229-
)
230-
.append(RcDoc::line())
231-
.append(RcDoc::text(")")),
232-
_ => RcDoc::nil(),
224+
GroupBy::GroupingSets(sets) => RcDoc::line()
225+
.append(
226+
RcDoc::text("GROUP BY GROUPING SETS (").append(RcDoc::line().nest(NEST_FACTOR)),
227+
)
228+
.append(
229+
interweave_comma(sets.into_iter().map(pretty_group_set))
230+
.nest(NEST_FACTOR)
231+
.group(),
232+
)
233+
.append(RcDoc::line())
234+
.append(RcDoc::text(")")),
235+
GroupBy::Rollup(exprs) => RcDoc::line()
236+
.append(RcDoc::text("GROUP BY ROLLUP (").append(RcDoc::line().nest(NEST_FACTOR)))
237+
.append(
238+
interweave_comma(exprs.into_iter().map(pretty_expr))
239+
.nest(NEST_FACTOR)
240+
.group(),
241+
)
242+
.append(RcDoc::line())
243+
.append(RcDoc::text(")")),
244+
GroupBy::Cube(exprs) => RcDoc::line()
245+
.append(RcDoc::text("GROUP BY CUBE (").append(RcDoc::line().nest(NEST_FACTOR)))
246+
.append(
247+
interweave_comma(exprs.into_iter().map(pretty_expr))
248+
.nest(NEST_FACTOR)
249+
.group(),
250+
)
251+
.append(RcDoc::line())
252+
.append(RcDoc::text(")")),
253+
}
254+
} else {
255+
RcDoc::nil()
233256
}
234257
}
235258

‎src/query/ast/src/ast/query.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,10 @@ pub enum GroupBy {
9696
///
9797
/// GroupSet := (expr [, expr]*) | expr
9898
GroupingSets(Vec<Vec<Expr>>),
99+
/// GROUP BY CUBE ( expr [, expr]* )
100+
Cube(Vec<Expr>),
101+
/// GROUP BY ROLLUP ( expr [, expr]* )
102+
Rollup(Vec<Expr>),
99103
}
100104

101105
/// A relational set expression, like `SELECT ... FROM ... {UNION|EXCEPT|INTERSECT} SELECT ... FROM ...`
@@ -471,6 +475,16 @@ impl Display for SelectStmt {
471475
}
472476
write!(f, ")")?;
473477
}
478+
GroupBy::Cube(exprs) => {
479+
write!(f, "CUBE (")?;
480+
write_comma_separated_list(f, exprs)?;
481+
write!(f, ")")?;
482+
}
483+
GroupBy::Rollup(exprs) => {
484+
write!(f, "ROLLUP (")?;
485+
write_comma_separated_list(f, exprs)?;
486+
write!(f, ")")?;
487+
}
474488
}
475489
}
476490

‎src/query/ast/src/parser/query.rs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -569,6 +569,14 @@ pub fn group_by_items(i: Input) -> IResult<GroupBy> {
569569
let normal = map(rule! { ^#comma_separated_list1(expr) }, |groups| {
570570
GroupBy::Normal(groups)
571571
});
572+
let cube = map(
573+
rule! { CUBE ~ "(" ~ ^#comma_separated_list1(expr) ~ ")" },
574+
|(_, _, groups, _)| GroupBy::Cube(groups),
575+
);
576+
let rollup = map(
577+
rule! { ROLLUP ~ "(" ~ ^#comma_separated_list1(expr) ~ ")" },
578+
|(_, _, groups, _)| GroupBy::Rollup(groups),
579+
);
572580
let group_set = alt((
573581
map(rule! {"(" ~ ")"}, |(_, _)| vec![]), // empty grouping set
574582
map(
@@ -581,7 +589,7 @@ pub fn group_by_items(i: Input) -> IResult<GroupBy> {
581589
rule! { GROUPING ~ SETS ~ "(" ~ ^#comma_separated_list1(group_set) ~ ")" },
582590
|(_, _, _, sets, _)| GroupBy::GroupingSets(sets),
583591
);
584-
rule!(#group_sets | #normal)(i)
592+
rule!(#group_sets | #cube | #rollup | #normal)(i)
585593
}
586594
587595
pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>> {

‎src/query/ast/src/parser/token.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -879,6 +879,10 @@ pub enum TokenKind {
879879
GROUPING,
880880
#[token("SETS", ignore(ascii_case))]
881881
SETS,
882+
#[token("CUBE", ignore(ascii_case))]
883+
CUBE,
884+
#[token("ROLLUP", ignore(ascii_case))]
885+
ROLLUP,
882886
}
883887

884888
// Reference: https://www.postgresql.org/docs/current/sql-keywords-appendix.html
@@ -982,6 +986,8 @@ impl TokenKind {
982986
// | TokenKind::FOREIGN
983987
// | TokenKind::GREATEST
984988
// | TokenKind::GROUPING
989+
| TokenKind::CUBE
990+
| TokenKind::ROLLUP
985991
// | TokenKind::IFNULL
986992
| TokenKind::IN
987993
// | TokenKind::INITIALLY

‎src/query/ast/tests/it/parser.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,6 +372,8 @@ fn test_statement() {
372372
r#"SELECT * FROM t GROUP BY GROUPING SETS (a, b, (c, d))"#,
373373
r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b), (c), (d, e))"#,
374374
r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b), (), (d, e))"#,
375+
r#"SELECT * FROM t GROUP BY CUBE (a, b, c)"#,
376+
r#"SELECT * FROM t GROUP BY ROLLUP (a, b, c)"#,
375377
];
376378

377379
for case in cases {

0 commit comments

Comments
 (0)