Merge branch 'main' into refactor_append

Author: SkyFan2002

src/common/column/src/offset.rs

@@ -122,7 +122,7 @@ impl<O: Offset> Offsets<O> {
     /// # Error
     /// This function errors iff the new last item is larger than what `O` supports.
     /// # Panic
-    /// This function asserts that `length > 0`.
+    /// This function asserts that `length >= O::zero()`.
     #[inline]
     pub fn try_push(&mut self, length: O) -> Result<(), Error> {
         let old_length = self.last();
@@ -171,7 +171,7 @@ impl<O: Offset> Offsets<O> {
 
     /// Returns a range (start, end) corresponding to the position `index`
     /// # Panic
-    /// This function panics iff `index >= self.len()`
+    /// This function panics iff `index >= self.len_proxy()`
     #[inline]
     pub fn start_end(&self, index: usize) -> (usize, usize) {
         // soundness: the invariant of the function
@@ -456,7 +456,7 @@ impl<O: Offset> OffsetsBuffer<O> {
 
     /// Returns a range (start, end) corresponding to the position `index`
     /// # Panic
-    /// This function panics iff `index >= self.len()`
+    /// This function panics iff `index >= self.len_proxy()`
     #[inline]
     pub fn start_end(&self, index: usize) -> (usize, usize) {
         // soundness: the invariant of the function

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

@@ -272,6 +272,15 @@ fn pretty_group_by(group_by: Option<GroupBy>) -> RcDoc<'static> {
                 )
                 .append(RcDoc::line())
                 .append(RcDoc::text(")")),
+
+            GroupBy::Combined(sets) => RcDoc::line()
+                .append(RcDoc::text("GROUP BY ").append(RcDoc::line().nest(NEST_FACTOR)))
+                .append(
+                    interweave_comma(sets.into_iter().map(|s| RcDoc::text(s.to_string())))
+                        .nest(NEST_FACTOR)
+                        .group(),
+                )
+                .append(RcDoc::line()),
         }
     } else {
         RcDoc::nil()

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

@@ -189,38 +189,8 @@ impl Display for SelectStmt {
         // GROUP BY clause
         if self.group_by.is_some() {
             write!(f, " GROUP BY ")?;
-            match self.group_by.as_ref().unwrap() {
-                GroupBy::Normal(exprs) => {
-                    write_comma_separated_list(f, exprs)?;
-                }
-                GroupBy::All => {
-                    write!(f, "ALL")?;
-                }
-                GroupBy::GroupingSets(sets) => {
-                    write!(f, "GROUPING SETS (")?;
-                    for (i, set) in sets.iter().enumerate() {
-                        if i > 0 {
-                            write!(f, ", ")?;
-                        }
-                        write!(f, "(")?;
-                        write_comma_separated_list(f, set)?;
-                        write!(f, ")")?;
-                    }
-                    write!(f, ")")?;
-                }
-                GroupBy::Cube(exprs) => {
-                    write!(f, "CUBE (")?;
-                    write_comma_separated_list(f, exprs)?;
-                    write!(f, ")")?;
-                }
-                GroupBy::Rollup(exprs) => {
-                    write!(f, "ROLLUP (")?;
-                    write_comma_separated_list(f, exprs)?;
-                    write!(f, ")")?;
-                }
-            }
+            write!(f, "{}", self.group_by.as_ref().unwrap())?;
         }
-
         // HAVING clause
         if let Some(having) = &self.having {
             write!(f, " HAVING {having}")?;
@@ -254,6 +224,60 @@ pub enum GroupBy {
     Cube(Vec<Expr>),
     /// GROUP BY ROLLUP ( expr [, expr]* )
     Rollup(Vec<Expr>),
+    Combined(Vec<GroupBy>),
+}
+
+impl GroupBy {
+    pub fn normal_items(&self) -> Vec<Expr> {
+        match self {
+            GroupBy::Normal(exprs) => exprs.clone(),
+            _ => Vec::new(),
+        }
+    }
+}
+
+impl Display for GroupBy {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            GroupBy::Normal(exprs) => {
+                write_comma_separated_list(f, exprs)?;
+            }
+            GroupBy::All => {
+                write!(f, "ALL")?;
+            }
+            GroupBy::GroupingSets(sets) => {
+                write!(f, "GROUPING SETS (")?;
+                for (i, set) in sets.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "(")?;
+                    write_comma_separated_list(f, set)?;
+                    write!(f, ")")?;
+                }
+                write!(f, ")")?;
+            }
+            GroupBy::Cube(exprs) => {
+                write!(f, "CUBE (")?;
+                write_comma_separated_list(f, exprs)?;
+                write!(f, ")")?;
+            }
+            GroupBy::Rollup(exprs) => {
+                write!(f, "ROLLUP (")?;
+                write_comma_separated_list(f, exprs)?;
+                write!(f, ")")?;
+            }
+            GroupBy::Combined(group_bys) => {
+                for (i, group_by) in group_bys.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "{}", group_by)?;
+                }
+            }
+        }
+        Ok(())
+    }
 }
 
 /// A relational set expression, like `SELECT ... FROM ... {UNION|EXCEPT|INTERSECT} SELECT ... FROM ...`

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

@@ -124,6 +124,9 @@ pub enum Statement {
     ShowCreateCatalog(ShowCreateCatalogStmt),
     CreateCatalog(CreateCatalogStmt),
     DropCatalog(DropCatalogStmt),
+    UseCatalog {
+        catalog: Identifier,
+    },
 
     // Databases
     ShowDatabases(ShowDatabasesStmt),
@@ -410,6 +413,7 @@ impl Statement {
             | Statement::Update(..)
             | Statement::ShowCatalogs(..)
             | Statement::ShowCreateCatalog(..)
+            | Statement::UseCatalog { .. }
             | Statement::ShowDatabases(..)
             | Statement::ShowDropDatabases(..)
             | Statement::ShowCreateDatabase(..)
@@ -718,6 +722,7 @@ impl Display for Statement {
             Statement::ShowCreateCatalog(stmt) => write!(f, "{stmt}")?,
             Statement::CreateCatalog(stmt) => write!(f, "{stmt}")?,
             Statement::DropCatalog(stmt) => write!(f, "{stmt}")?,
+            Statement::UseCatalog { catalog } => write!(f, "USE CATALOG {catalog}")?,
             Statement::ShowDatabases(stmt) => write!(f, "{stmt}")?,
             Statement::ShowDropDatabases(stmt) => write!(f, "{stmt}")?,
             Statement::ShowCreateDatabase(stmt) => write!(f, "{stmt}")?,

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

@@ -1073,10 +1073,6 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement>>> PrattParser<I>
 }
 
 pub fn group_by_items(i: Input) -> IResult<GroupBy> {
-    let normal = map(rule! { ^#comma_separated_list1(expr) }, |groups| {
-        GroupBy::Normal(groups)
-    });
-
     let all = map(rule! { ALL }, |_| GroupBy::All);
 
     let cube = map(
@@ -1096,10 +1092,31 @@ pub fn group_by_items(i: Input) -> IResult<GroupBy> {
         map(rule! { #expr }, |e| vec![e]),
     ));
     let group_sets = map(
-        rule! { GROUPING ~ SETS ~ "(" ~ ^#comma_separated_list1(group_set) ~ ")"  },
+        rule! { GROUPING ~ ^SETS ~ "(" ~ ^#comma_separated_list1(group_set) ~ ")"  },
         |(_, _, _, sets, _)| GroupBy::GroupingSets(sets),
     );
-    rule!(#all | #group_sets | #cube | #rollup | #normal)(i)
+
+    // New rule to handle multiple GroupBy items
+    let single_normal = map(rule! { #expr }, |group| GroupBy::Normal(vec![group]));
+    let group_by_item = alt((all, group_sets, cube, rollup, single_normal));
+    map(rule! { ^#comma_separated_list1(group_by_item) }, |items| {
+        if items.len() > 1 {
+            if items.iter().all(|item| matches!(item, GroupBy::Normal(_))) {
+                let items = items
+                    .into_iter()
+                    .flat_map(|item| match item {
+                        GroupBy::Normal(exprs) => exprs,
+                        _ => unreachable!(),
+                    })
+                    .collect();
+                GroupBy::Normal(items)
+            } else {
+                GroupBy::Combined(items)
+            }
+        } else {
+            items.into_iter().next().unwrap()
+        }
+    })(i)
 }
 
 pub fn window_frame_bound(i: Input) -> IResult<WindowFrameBound> {

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

@@ -516,6 +516,12 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             })
         },
     );
+    let use_catalog = map(
+        rule! {
+            USE ~ CATALOG ~ #ident
+        },
+        |(_, _, catalog)| Statement::UseCatalog { catalog },
+    );
 
     let show_databases = map(
         rule! {
@@ -2281,6 +2287,11 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             | #set_priority: "`SET PRIORITY (HIGH | MEDIUM | LOW) <object_id>`"
             | #system_action: "`SYSTEM (ENABLE | DISABLE) EXCEPTION_BACKTRACE`"
         ),
+        // use
+        rule!(
+                #use_catalog: "`USE CATALOG <catalog>`"
+                | #use_database : "`USE <database>`"
+        ),
         // database
         rule!(
             #show_databases : "`SHOW [FULL] DATABASES [(FROM | IN) <catalog>] [<show_limit>]`"
@@ -2290,7 +2301,6 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
             | #create_database : "`CREATE [OR REPLACE] DATABASE [IF NOT EXISTS] <database> [ENGINE = <engine>]`"
             | #drop_database : "`DROP DATABASE [IF EXISTS] <database>`"
             | #alter_database : "`ALTER DATABASE [IF EXISTS] <action>`"
-            | #use_database : "`USE <database>`"
         ),
         // network policy / password policy
         rule!(

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

@@ -137,6 +137,8 @@ fn test_statement() {
         r#"drop table if exists a."b";"#,
         r#"use "a";"#,
         r#"create catalog ctl type=hive connection=(url='<hive-meta-store>' thrift_protocol='binary');"#,
+        r#"select current_catalog();"#,
+        r#"use catalog ctl;"#,
         r#"create database if not exists a;"#,
         r#"create database ctl.t engine = Default;"#,
         r#"create database t engine = Default;"#,
@@ -589,12 +591,16 @@ fn test_statement() {
         "#,
         r#"SHOW FILE FORMATS"#,
         r#"DROP FILE FORMAT my_csv"#,
+        r#"SELECT * FROM t GROUP BY all"#,
+        r#"SELECT * FROM t GROUP BY a, b, c, d"#,
         r#"SELECT * FROM t GROUP BY GROUPING SETS (a, b, c, d)"#,
         r#"SELECT * FROM t GROUP BY GROUPING SETS (a, b, (c, d))"#,
         r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b), (c), (d, e))"#,
         r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b), (), (d, e))"#,
         r#"SELECT * FROM t GROUP BY CUBE (a, b, c)"#,
         r#"SELECT * FROM t GROUP BY ROLLUP (a, b, c)"#,
+        r#"SELECT * FROM t GROUP BY a, ROLLUP (b, c)"#,
+        r#"SELECT * FROM t GROUP BY GROUPING SETS ((a, b)), a, ROLLUP (b, c)"#,
         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'"#,
         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'"#,
         r#"DESC MASKING POLICY email_mask"#,