Skip to content

Commit b01fe51

Browse files
SkyFan2002SkyFan2002
committed
fix ut
1 parent 286fb36 commit b01fe51

File tree

42 files changed

+1536
-358
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+1536
-358
lines changed

‎src/common/column/src/offset.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ impl<O: Offset> Offsets<O> {
122122
/// # Error
123123
/// This function errors iff the new last item is larger than what `O` supports.
124124
/// # Panic
125-
/// This function asserts that `length > 0`.
125+
/// This function asserts that `length >= O::zero()`.
126126
#[inline]
127127
pub fn try_push(&mut self, length: O) -> Result<(), Error> {
128128
let old_length = self.last();
@@ -171,7 +171,7 @@ impl<O: Offset> Offsets<O> {
171171

172172
/// Returns a range (start, end) corresponding to the position `index`
173173
/// # Panic
174-
/// This function panics iff `index >= self.len()`
174+
/// This function panics iff `index >= self.len_proxy()`
175175
#[inline]
176176
pub fn start_end(&self, index: usize) -> (usize, usize) {
177177
// soundness: the invariant of the function
@@ -456,7 +456,7 @@ impl<O: Offset> OffsetsBuffer<O> {
456456

457457
/// Returns a range (start, end) corresponding to the position `index`
458458
/// # Panic
459-
/// This function panics iff `index >= self.len()`
459+
/// This function panics iff `index >= self.len_proxy()`
460460
#[inline]
461461
pub fn start_end(&self, index: usize) -> (usize, usize) {
462462
// soundness: the invariant of the function

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,15 @@ fn pretty_group_by(group_by: Option<GroupBy>) -> RcDoc<'static> {
272272
)
273273
.append(RcDoc::line())
274274
.append(RcDoc::text(")")),
275+
276+
GroupBy::Combined(sets) => RcDoc::line()
277+
.append(RcDoc::text("GROUP BY ").append(RcDoc::line().nest(NEST_FACTOR)))
278+
.append(
279+
interweave_comma(sets.into_iter().map(|s| RcDoc::text(s.to_string())))
280+
.nest(NEST_FACTOR)
281+
.group(),
282+
)
283+
.append(RcDoc::line()),
275284
}
276285
} else {
277286
RcDoc::nil()

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

Lines changed: 55 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -189,38 +189,8 @@ impl Display for SelectStmt {
189189
// GROUP BY clause
190190
if self.group_by.is_some() {
191191
write!(f, " GROUP BY ")?;
192-
match self.group_by.as_ref().unwrap() {
193-
GroupBy::Normal(exprs) => {
194-
write_comma_separated_list(f, exprs)?;
195-
}
196-
GroupBy::All => {
197-
write!(f, "ALL")?;
198-
}
199-
GroupBy::GroupingSets(sets) => {
200-
write!(f, "GROUPING SETS (")?;
201-
for (i, set) in sets.iter().enumerate() {
202-
if i > 0 {
203-
write!(f, ", ")?;
204-
}
205-
write!(f, "(")?;
206-
write_comma_separated_list(f, set)?;
207-
write!(f, ")")?;
208-
}
209-
write!(f, ")")?;
210-
}
211-
GroupBy::Cube(exprs) => {
212-
write!(f, "CUBE (")?;
213-
write_comma_separated_list(f, exprs)?;
214-
write!(f, ")")?;
215-
}
216-
GroupBy::Rollup(exprs) => {
217-
write!(f, "ROLLUP (")?;
218-
write_comma_separated_list(f, exprs)?;
219-
write!(f, ")")?;
220-
}
221-
}
192+
write!(f, "{}", self.group_by.as_ref().unwrap())?;
222193
}
223-
224194
// HAVING clause
225195
if let Some(having) = &self.having {
226196
write!(f, " HAVING {having}")?;
@@ -254,6 +224,60 @@ pub enum GroupBy {
254224
Cube(Vec<Expr>),
255225
/// GROUP BY ROLLUP ( expr [, expr]* )
256226
Rollup(Vec<Expr>),
227+
Combined(Vec<GroupBy>),
228+
}
229+
230+
impl GroupBy {
231+
pub fn normal_items(&self) -> Vec<Expr> {
232+
match self {
233+
GroupBy::Normal(exprs) => exprs.clone(),
234+
_ => Vec::new(),
235+
}
236+
}
237+
}
238+
239+
impl Display for GroupBy {
240+
fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
241+
match self {
242+
GroupBy::Normal(exprs) => {
243+
write_comma_separated_list(f, exprs)?;
244+
}
245+
GroupBy::All => {
246+
write!(f, "ALL")?;
247+
}
248+
GroupBy::GroupingSets(sets) => {
249+
write!(f, "GROUPING SETS (")?;
250+
for (i, set) in sets.iter().enumerate() {
251+
if i > 0 {
252+
write!(f, ", ")?;
253+
}
254+
write!(f, "(")?;
255+
write_comma_separated_list(f, set)?;
256+
write!(f, ")")?;
257+
}
258+
write!(f, ")")?;
259+
}
260+
GroupBy::Cube(exprs) => {
261+
write!(f, "CUBE (")?;
262+
write_comma_separated_list(f, exprs)?;
263+
write!(f, ")")?;
264+
}
265+
GroupBy::Rollup(exprs) => {
266+
write!(f, "ROLLUP (")?;
267+
write_comma_separated_list(f, exprs)?;
268+
write!(f, ")")?;
269+
}
270+
GroupBy::Combined(group_bys) => {
271+
for (i, group_by) in group_bys.iter().enumerate() {
272+
if i > 0 {
273+
write!(f, ", ")?;
274+
}
275+
write!(f, "{}", group_by)?;
276+
}
277+
}
278+
}
279+
Ok(())
280+
}
257281
}
258282

259283
/// A relational set expression, like `SELECT ... FROM ... {UNION|EXCEPT|INTERSECT} SELECT ... FROM ...`

‎src/query/ast/src/ast/statements/statement.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,9 @@ pub enum Statement {
124124
ShowCreateCatalog(ShowCreateCatalogStmt),
125125
CreateCatalog(CreateCatalogStmt),
126126
DropCatalog(DropCatalogStmt),
127+
UseCatalog {
128+
catalog: Identifier,
129+
},
127130

128131
// Databases
129132
ShowDatabases(ShowDatabasesStmt),
@@ -410,6 +413,7 @@ impl Statement {
410413
| Statement::Update(..)
411414
| Statement::ShowCatalogs(..)
412415
| Statement::ShowCreateCatalog(..)
416+
| Statement::UseCatalog { .. }
413417
| Statement::ShowDatabases(..)
414418
| Statement::ShowDropDatabases(..)
415419
| Statement::ShowCreateDatabase(..)
@@ -718,6 +722,7 @@ impl Display for Statement {
718722
Statement::ShowCreateCatalog(stmt) => write!(f, "{stmt}")?,
719723
Statement::CreateCatalog(stmt) => write!(f, "{stmt}")?,
720724
Statement::DropCatalog(stmt) => write!(f, "{stmt}")?,
725+
Statement::UseCatalog { catalog } => write!(f, "USE CATALOG {catalog}")?,
721726
Statement::ShowDatabases(stmt) => write!(f, "{stmt}")?,
722727
Statement::ShowDropDatabases(stmt) => write!(f, "{stmt}")?,
723728
Statement::ShowCreateDatabase(stmt) => write!(f, "{stmt}")?,

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

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1073,10 +1073,6 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement>>> PrattParser<I>
10731073
}
10741074

10751075
pub fn group_by_items(i: Input) -> IResult<GroupBy> {
1076-
let normal = map(rule! { ^#comma_separated_list1(expr) }, |groups| {
1077-
GroupBy::Normal(groups)
1078-
});
1079-
10801076
let all = map(rule! { ALL }, |_| GroupBy::All);
10811077

10821078
let cube = map(
@@ -1096,10 +1092,31 @@ pub fn group_by_items(i: Input) -> IResult<GroupBy> {
10961092
map(rule! { #expr }, |e| vec![e]),
10971093
));
10981094
let group_sets = map(
1099-
rule! { GROUPING ~ SETS ~ "(" ~ ^#comma_separated_list1(group_set) ~ ")" },
1095+
rule! { GROUPING ~ ^SETS ~ "(" ~ ^#comma_separated_list1(group_set) ~ ")" },
11001096
|(_, _, _, sets, _)| GroupBy::GroupingSets(sets),
11011097
);
1102-
rule!(#all | #group_sets | #cube | #rollup | #normal)(i)
1098+
1099+
// New rule to handle multiple GroupBy items
1100+
let single_normal = map(rule! { #expr }, |group| GroupBy::Normal(vec![group]));
1101+
let group_by_item = alt((all, group_sets, cube, rollup, single_normal));
1102+
map(rule! { ^#comma_separated_list1(group_by_item) }, |items| {
1103+
if items.len() > 1 {
1104+
if items.iter().all(|item| matches!(item, GroupBy::Normal(_))) {
1105+
let items = items
1106+
.into_iter()
1107+
.flat_map(|item| match item {
1108+
GroupBy::Normal(exprs) => exprs,
1109+
_ => unreachable!(),
1110+
})
1111+
.collect();
1112+
GroupBy::Normal(items)
1113+
} else {
1114+
GroupBy::Combined(items)
1115+
}
1116+
} else {
1117+
items.into_iter().next().unwrap()
1118+
}
1119+
})(i)
11031120
}
11041121

11051122
pub fn window_frame_bound(i: Input) -> IResult<WindowFrameBound> {

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -516,6 +516,12 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
516516
})
517517
},
518518
);
519+
let use_catalog = map(
520+
rule! {
521+
USE ~ CATALOG ~ #ident
522+
},
523+
|(_, _, catalog)| Statement::UseCatalog { catalog },
524+
);
519525

520526
let show_databases = map(
521527
rule! {
@@ -2281,6 +2287,11 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
22812287
| #set_priority: "`SET PRIORITY (HIGH | MEDIUM | LOW) <object_id>`"
22822288
| #system_action: "`SYSTEM (ENABLE | DISABLE) EXCEPTION_BACKTRACE`"
22832289
),
2290+
// use
2291+
rule!(
2292+
#use_catalog: "`USE CATALOG <catalog>`"
2293+
| #use_database : "`USE <database>`"
2294+
),
22842295
// database
22852296
rule!(
22862297
#show_databases : "`SHOW [FULL] DATABASES [(FROM | IN) <catalog>] [<show_limit>]`"
@@ -2290,7 +2301,6 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
22902301
| #create_database : "`CREATE [OR REPLACE] DATABASE [IF NOT EXISTS] <database> [ENGINE = <engine>]`"
22912302
| #drop_database : "`DROP DATABASE [IF EXISTS] <database>`"
22922303
| #alter_database : "`ALTER DATABASE [IF EXISTS] <action>`"
2293-
| #use_database : "`USE <database>`"
22942304
),
22952305
// network policy / password policy
22962306
rule!(

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,8 @@ fn test_statement() {
137137
r#"drop table if exists a."b";"#,
138138
r#"use "a";"#,
139139
r#"create catalog ctl type=hive connection=(url='<hive-meta-store>' thrift_protocol='binary');"#,
140+
r#"select current_catalog();"#,
141+
r#"use catalog ctl;"#,
140142
r#"create database if not exists a;"#,
141143
r#"create database ctl.t engine = Default;"#,
142144
r#"create database t engine = Default;"#,
@@ -589,12 +591,16 @@ fn test_statement() {
589591
"#,
590592
r#"SHOW FILE FORMATS"#,
591593
r#"DROP FILE FORMAT my_csv"#,
594+
r#"SELECT * FROM t GROUP BY all"#,
595+
r#"SELECT * FROM t GROUP BY a, b, c, d"#,
592596
r#"SELECT * FROM t GROUP BY GROUPING SETS (a, b, c, d)"#,
593597
r#"SELECT * FROM t GROUP BY GROUPING SETS (a, b, (c, d))"#,
594598
r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b), (c), (d, e))"#,
595599
r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b), (), (d, e))"#,
596600
r#"SELECT * FROM t GROUP BY CUBE (a, b, c)"#,
597601
r#"SELECT * FROM t GROUP BY ROLLUP (a, b, c)"#,
602+
r#"SELECT * FROM t GROUP BY a, ROLLUP (b, c)"#,
603+
r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b)), a, ROLLUP (b, c)"#,
598604
r#"CREATE MASKING POLICY email_mask AS (val STRING) RETURNS STRING -> CASE WHEN current_role() IN ('ANALYST') THEN VAL ELSE '*********'END comment = 'this is a masking policy'"#,
599605
r#"CREATE OR REPLACE MASKING POLICY email_mask AS (val STRING) RETURNS STRING -> CASE WHEN current_role() IN ('ANALYST') THEN VAL ELSE '*********'END comment = 'this is a masking policy'"#,
600606
r#"DESC MASKING POLICY email_mask"#,

0 commit comments

Comments
 (0)