fixed render label as label assignment

Author: kiraksi

sqlalchemy_bigquery/base.py

@@ -332,11 +332,25 @@ def visit_column(
 
             return self.preparer.quote(tablename) + "." + name
 
-    def visit_label(self, *args, **kwargs):
+    def visit_label(self, *args, within_group_by=False, **kwargs):
+        # Use labels in GROUP BY clause.
+        #
+        # Flag set in the group_by_clause method. Works around missing
+        # equivalent to supports_simple_order_by_label for group by.
+        if within_group_by:
+            if all(
+                keyword not in str(args[0])
+                for keyword in ("GROUPING SETS", "ROLLUP", "CUBE")
+            ):
+                kwargs["render_label_as_label"] = args[0]
         return super(BigQueryCompiler, self).visit_label(*args, **kwargs)
 
-    def group_by_clause(self, select, **kwargs):
-        return super(BigQueryCompiler, self).group_by_clause(select, **kwargs)
+    def group_by_clause(self, select, **kw):
+        return super(BigQueryCompiler, self).group_by_clause(
+            select,
+            **kw,
+            within_group_by=True,
+        )
 
     ############################################################################
     # Handle parameters in in

tests/unit/test_compiler.py

@@ -288,7 +288,7 @@ def test_grouping_sets(faux_conn, metadata):
         "table1",
         metadata,
         sqlalchemy.Column("foo", sqlalchemy.Integer),
-        sqlalchemy.Column("bar", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.ARRAY(sqlalchemy.Integer)),
     )
 
     q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(
@@ -309,7 +309,7 @@ def test_rollup(faux_conn, metadata):
         "table1",
         metadata,
         sqlalchemy.Column("foo", sqlalchemy.Integer),
-        sqlalchemy.Column("bar", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.ARRAY(sqlalchemy.Integer)),
     )
 
     q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(
@@ -330,7 +330,7 @@ def test_cube(faux_conn, metadata):
         "table1",
         metadata,
         sqlalchemy.Column("foo", sqlalchemy.Integer),
-        sqlalchemy.Column("bar", sqlalchemy.Integer),
+        sqlalchemy.Column("bar", sqlalchemy.ARRAY(sqlalchemy.Integer)),
     )
 
     q = sqlalchemy.select(table.c.foo, table.c.bar).group_by(