Repository: tobymao/sqlglot

Issue:
More About Presto Group by 
In #1079 , we have solved the order problem about four types. 

However, in Presto, there can be multiple keywords like `group by rollup(field_a), rollup(field_b), rollup(field_c)`, which is equivalent to `rollup(field_a, field_b, field_c)`. So it is still a problem.

```python
sqlglot.parse_one(
"""
select a
from t
group by a, rollup(b), rollup(c), rollup(d)
"""
)
```


Discussion:

Resolving patch:
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -5,7 +5,7 @@
 
 from sqlglot import exp
 from sqlglot.errors import ErrorLevel, UnsupportedError, concat_messages
-from sqlglot.helper import apply_index_offset, csv
+from sqlglot.helper import apply_index_offset, csv, seq_get
 from sqlglot.time import format_time
 from sqlglot.tokens import TokenType
 
@@ -1068,21 +1068,26 @@ def group_sql(self, expression: exp.Group) -> str:
             f"{self.seg('GROUPING SETS')} {self.wrap(grouping_sets)}" if grouping_sets else ""
         )
 
-        cube = expression.args.get("cube")
-        if cube is True:
-            cube = self.seg("WITH CUBE")
+        cube = expression.args.get("cube", [])
+        if seq_get(cube, 0) is True:
+            return f"{group_by}{self.seg('WITH CUBE')}"
         else:
-            cube = self.expressions(expression, key="cube", indent=False)
-            cube = f"{self.seg('CUBE')} {self.wrap(cube)}" if cube else ""
+            cube_sql = self.expressions(expression, key="cube", indent=False)
+            cube_sql = f"{self.seg('CUBE')} {self.wrap(cube_sql)}" if cube_sql else ""
 
-        rollup = expression.args.get("rollup")
-        if rollup is True:
-            rollup = self.seg("WITH ROLLUP")
+        rollup = expression.args.get("rollup", [])
+        if seq_get(rollup, 0) is True:
+            return f"{group_by}{self.seg('WITH ROLLUP')}"
         else:
-            rollup = self.expressions(expression, key="rollup", indent=False)
-            rollup = f"{self.seg('ROLLUP')} {self.wrap(rollup)}" if rollup else ""
+            rollup_sql = self.expressions(expression, key="rollup", indent=False)
+            rollup_sql = f"{self.seg('ROLLUP')} {self.wrap(rollup_sql)}" if rollup_sql else ""
 
-        return f"{group_by}{csv(grouping_sets, cube, rollup, sep=',')}"
+        groupings = csv(grouping_sets, cube_sql, rollup_sql, sep=",")
+
+        if expression.args.get("expressions") and groupings:
+            group_by = f"{group_by},"
+
+        return f"{group_by}{groupings}"
 
     def having_sql(self, expression: exp.Having) -> str:
         this = self.indent(self.sql(expression, "this"))
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -2,6 +2,7 @@
 
 import logging
 import typing as t
+from collections import defaultdict
 
 from sqlglot import exp
 from sqlglot.errors import ErrorLevel, ParseError, concat_messages, merge_errors
@@ -2183,39 +2184,33 @@ def _parse_group(self, skip_group_by_token: bool = False) -> t.Optional[exp.Expr
         if not skip_group_by_token and not self._match(TokenType.GROUP_BY):
             return None
 
-        elements = {}
+        elements = defaultdict(list)
 
-        def parse_element(parser: t.Callable, key: str) -> None:
-            if key in elements:
-                return
-
-            extension = parser()
+        while True:
+            expressions = self._parse_csv(self._parse_conjunction)
+            if expressions:
+                elements["expressions"].extend(expressions)
 
-            if extension:
-                elements[key] = extension
-                self._match(TokenType.COMMA)
+            grouping_sets = self._parse_grouping_sets()
+            if grouping_sets:
+                elements["grouping_sets"].extend(grouping_sets)
 
-        while True:
-            size = len(elements)
+            rollup = None
+            cube = None
 
-            parse_element(lambda: self._parse_csv(self._parse_conjunction), "expressions")
-            parse_element(self._parse_grouping_sets, "grouping_sets")
             with_ = self._match(TokenType.WITH)
-            parse_element(
-                lambda: self._match(TokenType.CUBE)
-                and (with_ or self._parse_wrapped_csv(self._parse_column)),
-                "cube",
-            )
-            parse_element(
-                lambda: self._match(TokenType.ROLLUP)
-                and (with_ or self._parse_wrapped_csv(self._parse_column)),
-                "rollup",
-            )
+            if self._match(TokenType.ROLLUP):
+                rollup = with_ or self._parse_wrapped_csv(self._parse_column)
+                elements["rollup"].extend(ensure_list(rollup))
+
+            if self._match(TokenType.CUBE):
+                cube = with_ or self._parse_wrapped_csv(self._parse_column)
+                elements["cube"].extend(ensure_list(cube))
 
-            if len(elements) == size or len(elements) >= 4:
+            if not (expressions or grouping_sets or rollup or cube):
                 break
 
-        return self.expression(exp.Group, **elements)
+        return self.expression(exp.Group, **elements)  # type: ignore
 
     def _parse_grouping_sets(self) -> t.Optional[t.List[t.Optional[exp.Expression]]]:
         if not self._match(TokenType.GROUPING_SETS):


Test patch:
diff --git a/tests/dialects/test_hive.py b/tests/dialects/test_hive.py
--- a/tests/dialects/test_hive.py
+++ b/tests/dialects/test_hive.py
@@ -345,13 +345,13 @@ def test_hive(self):
             "INSERT OVERWRITE TABLE zipcodes PARTITION(state = 0) VALUES (896, 'US', 'TAMPA', 33607)"
         )
         self.validate_identity(
-            "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, b GROUPING SETS ((a, b), a)"
+            "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, b, GROUPING SETS ((a, b), a)"
         )
         self.validate_identity(
-            "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, b GROUPING SETS ((t.a, b), a)"
+            "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, b, GROUPING SETS ((t.a, b), a)"
         )
         self.validate_identity(
-            "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, FOO(b) GROUPING SETS ((a, FOO(b)), a)"
+            "SELECT a, b, SUM(c) FROM tabl AS t GROUP BY a, FOO(b), GROUPING SETS ((a, FOO(b)), a)"
         )
         self.validate_identity(
             "SELECT key, value, GROUPING__ID, COUNT(*) FROM T1 GROUP BY key, value WITH CUBE"
@@ -648,9 +648,9 @@ def test_hive(self):
             },
         )
         self.validate_all(
-            "SELECT a, SUM(c) FROM t GROUP BY a, DATE_FORMAT(b, 'yyyy') GROUPING SETS ((a, DATE_FORMAT(b, 'yyyy')), a)",
+            "SELECT a, SUM(c) FROM t GROUP BY a, DATE_FORMAT(b, 'yyyy'), GROUPING SETS ((a, DATE_FORMAT(b, 'yyyy')), a)",
             write={
-                "hive": "SELECT a, SUM(c) FROM t GROUP BY a, DATE_FORMAT(CAST(b AS TIMESTAMP), 'yyyy') GROUPING SETS ((a, DATE_FORMAT(CAST(b AS TIMESTAMP), 'yyyy')), a)",
+                "hive": "SELECT a, SUM(c) FROM t GROUP BY a, DATE_FORMAT(CAST(b AS TIMESTAMP), 'yyyy'), GROUPING SETS ((a, DATE_FORMAT(CAST(b AS TIMESTAMP), 'yyyy')), a)",
             },
         )
 
diff --git a/tests/dialects/test_presto.py b/tests/dialects/test_presto.py
--- a/tests/dialects/test_presto.py
+++ b/tests/dialects/test_presto.py
@@ -365,6 +365,12 @@ def test_presto(self):
         self.validate_identity("START TRANSACTION ISOLATION LEVEL REPEATABLE READ")
         self.validate_identity("APPROX_PERCENTILE(a, b, c, d)")
 
+        self.validate_all(
+            "SELECT a FROM t GROUP BY a, ROLLUP(b), ROLLUP(c), ROLLUP(d)",
+            write={
+                "presto": "SELECT a FROM t GROUP BY a, ROLLUP (b, c, d)",
+            },
+        )
         self.validate_all(
             'SELECT a."b" FROM "foo"',
             write={
@@ -510,8 +516,8 @@ def test_presto(self):
         self.validate_all(
             "SELECT a, b, c, d, sum(y) FROM z GROUP BY CUBE(a) ROLLUP(a), GROUPING SETS((b, c)), d",
             write={
-                "presto": "SELECT a, b, c, d, SUM(y) FROM z GROUP BY d GROUPING SETS ((b, c)), CUBE (a), ROLLUP (a)",
-                "hive": "SELECT a, b, c, d, SUM(y) FROM z GROUP BY d GROUPING SETS ((b, c)), CUBE (a), ROLLUP (a)",
+                "presto": "SELECT a, b, c, d, SUM(y) FROM z GROUP BY d, GROUPING SETS ((b, c)), CUBE (a), ROLLUP (a)",
+                "hive": "SELECT a, b, c, d, SUM(y) FROM z GROUP BY d, GROUPING SETS ((b, c)), CUBE (a), ROLLUP (a)",
             },
         )
 
diff --git a/tests/fixtures/pretty.sql b/tests/fixtures/pretty.sql
--- a/tests/fixtures/pretty.sql
+++ b/tests/fixtures/pretty.sql
@@ -40,7 +40,7 @@ WITH cte1 AS (
         FROM (SELECT 1) AS x, y, (SELECT 2) z
         UNION ALL
         SELECT MAX(COALESCE(x AND y, a and b and c, d and e)), FOO(CASE WHEN a and b THEN c and d ELSE 3 END)
-        GROUP BY x, GROUPING SETS (a, (b, c)) CUBE(y, z)
+        GROUP BY x, GROUPING SETS (a, (b, c)), CUBE(y, z)
     ) x
 )
 SELECT a, b c FROM (
@@ -95,7 +95,7 @@ WITH cte1 AS (
       MAX(COALESCE(x AND y, a AND b AND c, d AND e)),
       FOO(CASE WHEN a AND b THEN c AND d ELSE 3 END)
     GROUP BY
-      x
+      x,
     GROUPING SETS (
       a,
       (b, c)