GROUP BY CUBE / ROLLUP / GROUPING SETS

beanquery/parser/ast.py

@@ -148,9 +148,19 @@ class From(Node):
 # A GROUP BY clause.
 #
 # Attributes:
-#   columns: A list of group-by expressions, simple Column() or otherwise.
+#   elements: A list of grouping elements. See GroupByElement.
 #   having: An expression tree for the optional HAVING clause, or None.
-GroupBy = node('GroupBy', 'columns having')
+GroupBy = node('GroupBy', 'elements having')
+
+# A GROUP BY grouping element. 
+# Attributes:
+#   columns: If type == '': ast.Column() or an integer column index. If 
+#      type != '': A list of ast.Column() or integer column indexes.
+#   type: Distinguishes the grouping modes relating to the keywords 
+#     ROLLUP, CUBE, GROUPING SETS. 'type' has the keyword in lower case. For
+#     simple grouping (no keyword), 'type' is the empty string.
+#    
+GroupByElement = node('GroupByElement', 'columns type')
 
 # An ORDER BY clause.
 #

beanquery/parser/bql.ebnf

@@ -1,9 +1,9 @@
 @@grammar :: BQL
 @@parseinfo :: True
 @@ignorecase :: True
-@@keyword :: 'AND' 'AS' 'ASC' 'BY' 'DESC' 'DISTINCT' 'FALSE' 'FROM'
-    'GROUP' 'HAVING' 'IN' 'IS' 'LIMIT' 'NOT' 'OR' 'ORDER' 'PIVOT'
-    'SELECT' 'TRUE' 'WHERE'
+@@keyword :: 'AND' 'AS' 'ASC' 'BY' 'CUBE' 'DESC' 'DISTINCT' 'FALSE' 'FROM'
+    'GROUP' 'GROUPING' 'HAVING' 'IN' 'IS' 'LIMIT' 'NOT' 'OR' 'ORDER' 'PIVOT'
+    'ROLLUP' 'SELECT' 'SETS' 'TRUE' 'WHERE'
 @@keyword :: 'CREATE' 'TABLE' 'USING' 'INSERT' 'INTO'
 @@keyword :: 'BALANCES' 'JOURNAL' 'PRINT'
 @@comments :: /(\/\*([^*]|[\r\n]|(\*+([^*\/]|[\r\n])))*\*+\/)/
@@ -55,8 +55,36 @@ table::Table
     = name:identifier
     ;
 
+
+# A GROUP BY clause supports multiple syntaxes:
+# 1. Regular GROUP BY: GROUP BY col1, col2 [HAVING condition]
+# 2. Full ROLLUP: GROUP BY ROLLUP (col1, col2) [HAVING condition]
+# 3. Full CUBE: GROUP BY CUBE (col1, col2) [HAVING condition]
+# 4. GROUPING SETS: GROUP BY GROUPING SETS ((col1, col2), (col1), ()) [HAVING condition]
+# 5. Mixed grouping: GROUP BY col1, ROLLUP (col2, col3) [HAVING condition]
+#
+# Examples:
+# - GROUP BY account, year
+# - GROUP BY account, year HAVING SUM(position) > 100
+# - GROUP BY ROLLUP (account, year)
+# - GROUP BY ROLLUP (account, year) HAVING SUM(position) > 100
+# - GROUP BY CUBE (account, year)
+# - GROUP BY GROUPING SETS ((account, year), (account), ())
+# - GROUP BY region, ROLLUP (year, month)
+# - GROUP BY region, CUBE (year, month)
 groupby::GroupBy
-    = columns:','.{ (integer | expression) }+ ['HAVING' having:expression]
+    = elements:','.{ grouping_element }+ ['HAVING' having:expression]
+    ;
+
+grouping_element::GroupByElement
+    = 'ROLLUP' '(' columns:','.{ (integer | expression) }+ ')' type:`rollup`
+    | 'CUBE' '(' columns:','.{ (integer | expression) }+ ')' type:`cube`
+    | 'GROUPING' 'SETS' '(' columns:','.{ grouping_set }+ ')' type:`grouping sets`
+    |  columns: (integer | expression) type:`` 
+    ;
+
+grouping_set
+    = '(' @:','.{ (integer | expression) }* ')'
     ;
 
 order::OrderBy

beanquery/parser/parser.py

@@ -29,11 +29,13 @@
     'AS',
     'ASC',
     'BY',
+    'CUBE',
     'DESC',
     'DISTINCT',
     'FALSE',
     'FROM',
     'GROUP',
+    'GROUPING',
     'HAVING',
     'IN',
     'IS',
@@ -42,7 +44,9 @@
     'OR',
     'ORDER',
     'PIVOT',
+    'ROLLUP',
     'SELECT',
+    'SETS',
     'TRUE',
     'WHERE',
     'CREATE',
@@ -340,6 +344,113 @@ def _table_(self):
     @tatsumasu('GroupBy')
     def _groupby_(self):
 
+        def sep0():
+            self._token(',')
+
+        def block1():
+            self._grouping_element_()
+        self._positive_gather(block1, sep0)
+        self.name_last_node('elements')
+        with self._optional():
+            self._token('HAVING')
+            self._expression_()
+            self.name_last_node('having')
+            self._define(['having'], [])
+        self._define(['elements', 'having'], [])
+
+    @tatsumasu('GroupByElement')
+    def _grouping_element_(self):
+        with self._choice():
+            with self._option():
+                self._token('ROLLUP')
+                self._token('(')
+
+                def sep0():
+                    self._token(',')
+
+                def block1():
+                    with self._group():
+                        with self._choice():
+                            with self._option():
+                                self._integer_()
+                            with self._option():
+                                self._expression_()
+                            self._error(
+                                'expecting one of: '
+                                '<expression> <integer>'
+                            )
+                self._positive_gather(block1, sep0)
+                self.name_last_node('columns')
+                self._token(')')
+                self._constant('rollup')
+                self.name_last_node('type')
+                self._define(['columns', 'type'], [])
+            with self._option():
+                self._token('CUBE')
+                self._token('(')
+
+                def sep2():
+                    self._token(',')
+
+                def block3():
+                    with self._group():
+                        with self._choice():
+                            with self._option():
+                                self._integer_()
+                            with self._option():
+                                self._expression_()
+                            self._error(
+                                'expecting one of: '
+                                '<expression> <integer>'
+                            )
+                self._positive_gather(block3, sep2)
+                self.name_last_node('columns')
+                self._token(')')
+                self._constant('cube')
+                self.name_last_node('type')
+                self._define(['columns', 'type'], [])
+            with self._option():
+                self._token('GROUPING')
+                self._token('SETS')
+                self._token('(')
+
+                def sep4():
+                    self._token(',')
+
+                def block5():
+                    self._grouping_set_()
+                self._positive_gather(block5, sep4)
+                self.name_last_node('columns')
+                self._token(')')
+                self._constant('grouping sets')
+                self.name_last_node('type')
+                self._define(['columns', 'type'], [])
+            with self._option():
+                with self._group():
+                    with self._choice():
+                        with self._option():
+                            self._integer_()
+                        with self._option():
+                            self._expression_()
+                        self._error(
+                            'expecting one of: '
+                            '<expression> <integer>'
+                        )
+                self.name_last_node('columns')
+                self._constant('')
+                self.name_last_node('type')
+                self._define(['columns', 'type'], [])
+            self._error(
+                'expecting one of: '
+                "'CUBE' 'GROUPING' 'ROLLUP' <conjunction>"
+                '<disjunction> <expression> <integer>'
+                '[0-9]+'
+            )
+
+    @tatsumasu()
+    def _grouping_set_(self):
+        self._token('(')
+
         def sep0():
             self._token(',')
 
@@ -354,14 +465,9 @@ def block1():
                         'expecting one of: '
                         '<expression> <integer>'
                     )
-        self._positive_gather(block1, sep0)
-        self.name_last_node('columns')
-        with self._optional():
-            self._token('HAVING')
-            self._expression_()
-            self.name_last_node('having')
-            self._define(['having'], [])
-        self._define(['columns', 'having'], [])
+        self._gather(block1, sep0)
+        self.name_last_node('@')
+        self._token(')')
 
     @tatsumasu('OrderBy')
     def _order_(self):

beanquery/parser_test.py

@@ -309,52 +309,112 @@ def test_groupby_one(self):
         self.assertParse(
             "SELECT * GROUP BY a;",
             Select(ast.Asterisk(),
-                   group_by=ast.GroupBy([ast.Column('a')], None)))
+                group_by=ast.GroupBy(
+                    [ast.GroupByElement(ast.Column('a'), '')], None)))
 
     def test_groupby_many(self):
+        ge = ast.GroupByElement
         self.assertParse(
             "SELECT * GROUP BY a, b, c;",
             Select(ast.Asterisk(),
-                   group_by=ast.GroupBy([
-                       ast.Column('a'),
-                       ast.Column('b'),
-                       ast.Column('c')], None)))
+                group_by=ast.GroupBy([
+                    ge(ast.Column('a'), ''),
+                    ge(ast.Column('b'), ''),
+                    ge(ast.Column('c'), '')],
+                None)))
 
     def test_groupby_expr(self):
+        ge = ast.GroupByElement
         self.assertParse(
             "SELECT * GROUP BY length(a) > 0, b;",
             Select(ast.Asterisk(),
-                   group_by=ast.GroupBy([
-                       ast.Greater(
-                           ast.Function('length', [
-                               ast.Column('a')]),
-                           ast.Constant(0)),
-                       ast.Column('b')], None)))
+                group_by=ast.GroupBy([
+                    ge( ast.Greater(
+                            ast.Function('length', [
+                                ast.Column('a')]),
+                            ast.Constant(0)), ''),
+                    ge(ast.Column('b'), '')],
+                    None)))
 
     def test_groupby_having(self):
         self.assertParse(
             "SELECT * GROUP BY a HAVING sum(x) = 0;",
             Select(ast.Asterisk(),
-                   group_by=ast.GroupBy([ast.Column('a')],
-                        ast.Equal(
-                            ast.Function('sum', [
-                                ast.Column('x')]),
-                            ast.Constant(0)))))
+                group_by=ast.GroupBy(
+                    [ast.GroupByElement(ast.Column('a'), '')],
+                    ast.Equal(
+                        ast.Function('sum', [
+                            ast.Column('x')]),
+                        ast.Constant(0)))))
 
     def test_groupby_numbers(self):
         self.assertParse(
             "SELECT * GROUP BY 1;",
             Select(ast.Asterisk(),
-                   group_by=ast.GroupBy([1], None)))
+                group_by=ast.GroupBy([ast.GroupByElement(1, '')], None)))
 
+        ge = ast.GroupByElement
         self.assertParse(
             "SELECT * GROUP BY 2, 4, 5;",
             Select(ast.Asterisk(),
-                   group_by=ast.GroupBy([2, 4, 5], None)))
+                group_by=ast.GroupBy(
+                    [ge(2, ''), ge(4, ''), ge(5, '')], None)))
 
     def test_groupby_empty(self):
         with self.assertRaises(parser.ParseError):
             parser.parse("SELECT * GROUP BY;")
+
+    def test_groupby_rollup(self):
+        """Test ROLLUP syntax in GROUP BY clause."""
+        self.assertParse(
+            "SELECT * GROUP BY ROLLUP (account, year);",
+            Select(ast.Asterisk(),
+                group_by=ast.GroupBy([
+                    ast.GroupByElement(
+                        [ast.Column('account'), ast.Column('year')], 'rollup')
+                ], None)
+            )
+        )
+
+    def test_groupby_cube(self):
+        """Test CUBE syntax in GROUP BY clause."""
+        self.assertParse(
+            "SELECT * GROUP BY CUBE (account, year);",
+            Select(ast.Asterisk(),
+                group_by=ast.GroupBy([
+                    ast.GroupByElement(
+                          [ast.Column('account'), ast.Column('year')], 'cube')
+                    ], None)
+            )
+        )
+
+    def test_groupby_grouping_sets(self):
+        """Test GROUPING SETS syntax in GROUP BY clause."""
+        self.assertParse(
+            "SELECT * GROUP BY GROUPING SETS ((account, year), (account), ());",
+            Select(ast.Asterisk(),
+                group_by=ast.GroupBy([
+                    ast.GroupByElement([
+                        [ast.Column('account'), ast.Column('year')],
+                        [ast.Column('account')],
+                        []
+                    ], 'grouping sets')
+                ], None)
+            )
+        )
+
+    def test_groupby_mixed(self):
+        """Test mixed grouping elements in GROUP BY clause."""
+        ge = ast.GroupByElement
+        self.assertParse(
+            "SELECT * GROUP BY region, ROLLUP (year, month);",
+            Select(ast.Asterisk(),
+                group_by=ast.GroupBy([
+                    ge(ast.Column('region'),''),
+                    ge([ast.Column('year'), ast.Column('month')], 'rollup')
+                ], None)
+            )
+        )
 
 
 class TestSelectOrderBy(QueryParserTestBase):