generate more compact aggregation SQL

Author: chelsea-lin

bigframes/core/compile/sqlglot/aggregate_compiler.py

@@ -30,7 +30,6 @@ def compile_aggregate(
     order_by: tuple[sge.Expression, ...],
 ) -> sge.Expression:
     """Compiles BigFrames aggregation expression into SQLGlot expression."""
-    # TODO: try to remove type: ignore
     if isinstance(aggregate, expression.NullaryAggregation):
         return compile_nullary_agg(aggregate.op)
     if isinstance(aggregate, expression.UnaryAggregation):
@@ -39,9 +38,9 @@ def compile_aggregate(
             aggregate.arg.output_type,
         )
         if not aggregate.op.order_independent:
-            return compile_ordered_unary_agg(aggregate.op, column, order_by=order_by)  # type: ignore
+            return compile_ordered_unary_agg(aggregate.op, column, order_by=order_by)
         else:
-            return compile_unary_agg(aggregate.op, column)  # type: ignore
+            return compile_unary_agg(aggregate.op, column)
     elif isinstance(aggregate, expression.BinaryAggregation):
         left = typed_expr.TypedExpr(
             scalar_compiler.compile_scalar_expression(aggregate.left),
@@ -51,7 +50,7 @@ def compile_aggregate(
             scalar_compiler.compile_scalar_expression(aggregate.right),
             aggregate.right.output_type,
         )
-        return compile_binary_agg(aggregate.op, left, right)  # type: ignore
+        return compile_binary_agg(aggregate.op, left, right)
     else:
         raise ValueError(f"Unexpected aggregation: {aggregate}")
 
@@ -88,6 +87,7 @@ def compile_ordered_unary_agg(
     op: ops.aggregations.WindowOp,
     column: typed_expr.TypedExpr,
     window: typing.Optional[window_spec.WindowSpec] = None,
+    order_by: typing.Sequence[sge.Expression] = [],
 ) -> sge.Expression:
     raise ValueError(f"Can't compile unrecognized operation: {op}")
 

bigframes/core/compile/sqlglot/compiler.py

@@ -18,7 +18,6 @@
 import typing
 
 from google.cloud import bigquery
-import sqlglot as sg
 import sqlglot.expressions as sge
 
 from bigframes.core import expression, guid, identifiers, nodes, pyarrow_utils, rewrite
@@ -279,7 +278,6 @@ def compile_aggregate(
                     ordering.scalar_expression
                 ),
                 desc=ordering.direction.is_ascending is False,
-                # TODO: _convert_row_ordering_to_table_values for overwrite.
                 nulls_first=ordering.na_last is False,
             )
             for ordering in node.order_by
@@ -293,16 +291,13 @@ def compile_aggregate(
             for by_col in node.by_column_ids
         )
 
-        result = child.aggregate(aggregations, by_cols)
+        dropna_cols = []
         if node.dropna:
-            conditions = []
             for key, by_col in zip(node.by_column_ids, by_cols):
                 if node.child.field_by_id[key.id].nullable:
-                    conditions.append(
-                        sg.not_(sge.Is(this=by_col, expression=sge.Null()))
-                    )
-            result = result.filter(tuple(conditions))
-        return result
+                    dropna_cols.append(by_col)
+
+        return child.aggregate(aggregations, by_cols, tuple(dropna_cols))
 
 
 def _replace_unsupported_ops(node: nodes.BigFrameNode):

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -371,15 +371,15 @@ def sample(self, fraction: float) -> SQLGlotIR:
     def aggregate(
         self,
         aggregations: tuple[tuple[str, sge.Expression], ...],
-        by_column_ids: tuple[sge.Expression, ...],
+        by_cols: tuple[sge.Expression, ...],
+        dropna_cols: tuple[sge.Expression, ...],
     ) -> SQLGlotIR:
         """Applies the aggregation expressions.
 
         Args:
             aggregations: output_column_id, aggregation_expr tuples
-            by_column_ids: column ids of the aggregation key, this is preserved through
-              the transform
-            dropna: whether null keys should be dropped
+            by_cols: column expressions for aggregation
+            dropna_cols: columns whether null keys should be dropped
         """
         aggregations_expr = [
             sge.Alias(
@@ -395,9 +395,18 @@ def aggregate(
                 next(self.uid_gen.get_uid_stream("bfcte_")), quoted=self.quoted
             ),
         )
-        new_expr = new_expr.group_by(*by_column_ids).select(
-            *[*by_column_ids, *aggregations_expr], append=False
+        new_expr = new_expr.group_by(*by_cols).select(
+            *[*by_cols, *aggregations_expr], append=False
         )
+
+        condition = _and(
+            tuple(
+                sg.not_(sge.Is(this=drop_col, expression=sge.Null()))
+                for drop_col in dropna_cols
+            )
+        )
+        if condition is not None:
+            new_expr = new_expr.where(condition, append=False)
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
     def insert(

tests/unit/core/compile/sqlglot/snapshots/test_compile_aggregate/test_compile_aggregate/out.sql

@@ -14,18 +14,14 @@ WITH `bfcte_0` AS (
     `bfcol_3`,
     COALESCE(SUM(`bfcol_2`), 0) AS `bfcol_6`
   FROM `bfcte_1`
-  GROUP BY
-    `bfcol_3`
-), `bfcte_3` AS (
-  SELECT
-    *
-  FROM `bfcte_2`
   WHERE
     NOT `bfcol_3` IS NULL
+  GROUP BY
+    `bfcol_3`
 )
 SELECT
   `bfcol_3` AS `bool_col`,
   `bfcol_6` AS `int64_too`
-FROM `bfcte_3`
+FROM `bfcte_2`
 ORDER BY
   `bfcol_3` ASC NULLS LAST