fix(query): keep remaining_predicates when filtering grouping sets (#16971)

* fix(query): keep remaining_predicates when filtering grouping sets

* update

* update

* update

* update

* Update 03_0003_select_group_by.test

* update

---------

Co-authored-by: TCeason <[email protected]>

Author: sundy-li

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/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/tests/it/parser.rs

@@ -591,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"#,