refactor: Aggregation is now an expression subclass

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -24,6 +24,7 @@
 import pandas
 import pyarrow as pa
 
+from bigframes.core import expression_types
 import bigframes.core.expression as ex
 import bigframes.core.guid
 import bigframes.core.identifiers as ids
@@ -190,7 +191,7 @@ def row_count(self) -> ArrayValue:
                 child=self.node,
                 aggregations=(
                     (
-                        ex.NullaryAggregation(agg_ops.size_op),
+                        expression_types.NullaryAggregation(agg_ops.size_op),
                         ids.ColumnId(bigframes.core.guid.generate_guid()),
                     ),
                 ),
@@ -379,7 +380,7 @@ def drop_columns(self, columns: Iterable[str]) -> ArrayValue:
 
     def aggregate(
         self,
-        aggregations: typing.Sequence[typing.Tuple[ex.Aggregation, str]],
+        aggregations: typing.Sequence[typing.Tuple[expression_types.Aggregation, str]],
         by_column_ids: typing.Sequence[str] = (),
         dropna: bool = True,
     ) -> ArrayValue:
@@ -420,15 +421,15 @@ def project_window_op(
         """
 
         return self.project_window_expr(
-            ex.UnaryAggregation(op, ex.deref(column_name)),
+            expression_types.UnaryAggregation(op, ex.deref(column_name)),
             window_spec,
             never_skip_nulls,
             skip_reproject_unsafe,
         )
 
     def project_window_expr(
         self,
-        expression: ex.Aggregation,
+        expression: expression_types.Aggregation,
         window: WindowSpec,
         never_skip_nulls=False,
         skip_reproject_unsafe: bool = False,

bigframes/core/bigframe_node.py

@@ -20,15 +20,12 @@
 import functools
 import itertools
 import typing
-from typing import Callable, Dict, Generator, Iterable, Mapping, Sequence, Tuple, Union
+from typing import Callable, Dict, Generator, Iterable, Mapping, Sequence, Tuple
 
 from bigframes.core import expression, field, identifiers
 import bigframes.core.schema as schemata
 import bigframes.dtypes
 
-if typing.TYPE_CHECKING:
-    import bigframes.session
-
 COLUMN_SET = frozenset[identifiers.ColumnId]
 
 T = typing.TypeVar("T")
@@ -281,8 +278,8 @@ def field_by_id(self) -> Mapping[identifiers.ColumnId, field.Field]:
     @property
     def _node_expressions(
         self,
-    ) -> Sequence[Union[expression.Expression, expression.Aggregation]]:
-        """List of scalar expressions. Intended for checking engine compatibility with used ops."""
+    ) -> Sequence[expression.Expression]:
+        """List of expressions. Intended for checking engine compatibility with used ops."""
         return ()
 
     # Plan algorithms

bigframes/core/block_transforms.py

@@ -21,12 +21,12 @@
 import pandas as pd
 
 import bigframes.constants
+from bigframes.core import expression_types
 import bigframes.core as core
 import bigframes.core.blocks as blocks
 import bigframes.core.expression as ex
 import bigframes.core.ordering as ordering
 import bigframes.core.window_spec as windows
-import bigframes.dtypes
 import bigframes.dtypes as dtypes
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
@@ -133,7 +133,7 @@ def quantile(
     block, _ = block.aggregate(
         grouping_column_ids,
         tuple(
-            ex.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col))
+            expression_types.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col))
             for col in quantile_cols
         ),
         column_labels=pd.Index(labels),
@@ -363,7 +363,7 @@ def value_counts(
         block = dropna(block, columns, how="any")
     block, agg_ids = block.aggregate(
         by_column_ids=(*grouping_keys, *columns),
-        aggregations=[ex.NullaryAggregation(agg_ops.size_op)],
+        aggregations=[expression_types.NullaryAggregation(agg_ops.size_op)],
         dropna=drop_na and not grouping_keys,
     )
     count_id = agg_ids[0]
@@ -647,15 +647,15 @@ def skew(
     # counts, moment3 for each column
     aggregations = []
     for i, col in enumerate(original_columns):
-        count_agg = ex.UnaryAggregation(
+        count_agg = expression_types.UnaryAggregation(
             agg_ops.count_op,
             ex.deref(col),
         )
-        moment3_agg = ex.UnaryAggregation(
+        moment3_agg = expression_types.UnaryAggregation(
             agg_ops.mean_op,
             ex.deref(delta3_ids[i]),
         )
-        variance_agg = ex.UnaryAggregation(
+        variance_agg = expression_types.UnaryAggregation(
             agg_ops.PopVarOp(),
             ex.deref(col),
         )
@@ -698,9 +698,13 @@ def kurt(
     # counts, moment4 for each column
     aggregations = []
     for i, col in enumerate(original_columns):
-        count_agg = ex.UnaryAggregation(agg_ops.count_op, ex.deref(col))
-        moment4_agg = ex.UnaryAggregation(agg_ops.mean_op, ex.deref(delta4_ids[i]))
-        variance_agg = ex.UnaryAggregation(agg_ops.PopVarOp(), ex.deref(col))
+        count_agg = expression_types.UnaryAggregation(agg_ops.count_op, ex.deref(col))
+        moment4_agg = expression_types.UnaryAggregation(
+            agg_ops.mean_op, ex.deref(delta4_ids[i])
+        )
+        variance_agg = expression_types.UnaryAggregation(
+            agg_ops.PopVarOp(), ex.deref(col)
+        )
         aggregations.extend([count_agg, moment4_agg, variance_agg])
 
     block, agg_ids = block.aggregate(

bigframes/core/blocks.py

@@ -51,11 +51,12 @@
 from bigframes import session
 from bigframes._config import sampling_options
 import bigframes.constants
-from bigframes.core import local_data
+from bigframes.core import expression_types, local_data
 import bigframes.core as core
 import bigframes.core.compile.googlesql as googlesql
 import bigframes.core.expression as ex
 import bigframes.core.expression as scalars
+import bigframes.core.expression_types as ex_types
 import bigframes.core.guid as guid
 import bigframes.core.identifiers
 import bigframes.core.join_def as join_defs
@@ -1143,7 +1144,7 @@ def apply_window_op(
         skip_reproject_unsafe: bool = False,
         never_skip_nulls: bool = False,
     ) -> typing.Tuple[Block, str]:
-        agg_expr = ex.UnaryAggregation(op, ex.deref(column))
+        agg_expr = expression_types.UnaryAggregation(op, ex.deref(column))
         return self.apply_analytic(
             agg_expr,
             window_spec,
@@ -1155,7 +1156,7 @@ def apply_window_op(
 
     def apply_analytic(
         self,
-        agg_expr: ex.Aggregation,
+        agg_expr: expression_types.Aggregation,
         window: windows.WindowSpec,
         result_label: Label,
         *,
@@ -1248,9 +1249,9 @@ def aggregate_all_and_stack(
         if axis_n == 0:
             aggregations = [
                 (
-                    ex.UnaryAggregation(operation, ex.deref(col_id))
+                    expression_types.UnaryAggregation(operation, ex.deref(col_id))
                     if isinstance(operation, agg_ops.UnaryAggregateOp)
-                    else ex.NullaryAggregation(operation),
+                    else expression_types.NullaryAggregation(operation),
                     col_id,
                 )
                 for col_id in self.value_columns
@@ -1279,7 +1280,10 @@ def aggregate_size(
     ):
         """Returns a block object to compute the size(s) of groups."""
         agg_specs = [
-            (ex.NullaryAggregation(agg_ops.SizeOp()), guid.generate_guid()),
+            (
+                expression_types.NullaryAggregation(agg_ops.SizeOp()),
+                guid.generate_guid(),
+            ),
         ]
         output_col_ids = [agg_spec[1] for agg_spec in agg_specs]
         result_expr = self.expr.aggregate(agg_specs, by_column_ids, dropna=dropna)
@@ -1350,7 +1354,7 @@ def remap_f(x):
     def aggregate(
         self,
         by_column_ids: typing.Sequence[str] = (),
-        aggregations: typing.Sequence[ex.Aggregation] = (),
+        aggregations: typing.Sequence[expression_types.Aggregation] = (),
         column_labels: Optional[pd.Index] = None,
         *,
         dropna: bool = True,
@@ -1419,9 +1423,9 @@ def get_stat(
 
         aggregations = [
             (
-                ex.UnaryAggregation(stat, ex.deref(column_id))
+                expression_types.UnaryAggregation(stat, ex.deref(column_id))
                 if isinstance(stat, agg_ops.UnaryAggregateOp)
-                else ex.NullaryAggregation(stat),
+                else expression_types.NullaryAggregation(stat),
                 stat.name,
             )
             for stat in stats_to_fetch
@@ -1447,7 +1451,7 @@ def get_binary_stat(
         # TODO(kemppeterson): Add a cache here.
         aggregations = [
             (
-                ex.BinaryAggregation(
+                expression_types.BinaryAggregation(
                     stat, ex.deref(column_id_left), ex.deref(column_id_right)
                 ),
                 f"{stat.name}_{column_id_left}{column_id_right}",
@@ -1474,9 +1478,9 @@ def summarize(
         labels = pd.Index([stat.name for stat in stats])
         aggregations = [
             (
-                ex.UnaryAggregation(stat, ex.deref(col_id))
+                expression_types.UnaryAggregation(stat, ex.deref(col_id))
                 if isinstance(stat, agg_ops.UnaryAggregateOp)
-                else ex.NullaryAggregation(stat),
+                else expression_types.NullaryAggregation(stat),
                 f"{col_id}-{stat.name}",
             )
             for stat in stats
@@ -1750,7 +1754,7 @@ def pivot(
 
         block = block.select_columns(column_ids)
         aggregations = [
-            ex.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col_id))
+            expression_types.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col_id))
             for col_id in column_ids
         ]
         result_block, _ = block.aggregate(
@@ -2018,7 +2022,7 @@ def _generate_resample_label(
 
         agg_specs = [
             (
-                ex.UnaryAggregation(agg_ops.min_op, ex.deref(col_id)),
+                expression_types.UnaryAggregation(agg_ops.min_op, ex.deref(col_id)),
                 guid.generate_guid(),
             ),
         ]
@@ -2047,13 +2051,13 @@ def _generate_resample_label(
         # Generate integer label sequence.
         min_agg_specs = [
             (
-                ex.UnaryAggregation(agg_ops.min_op, ex.deref(label_col_id)),
+                ex_types.UnaryAggregation(agg_ops.min_op, ex.deref(label_col_id)),
                 guid.generate_guid(),
             ),
         ]
         max_agg_specs = [
             (
-                ex.UnaryAggregation(agg_ops.max_op, ex.deref(label_col_id)),
+                ex_types.UnaryAggregation(agg_ops.max_op, ex.deref(label_col_id)),
                 guid.generate_guid(),
             ),
         ]

bigframes/core/compile/compiled.py

@@ -35,6 +35,7 @@
 import bigframes.core.compile.ibis_compiler.scalar_op_compiler as op_compilers
 import bigframes.core.compile.ibis_types
 import bigframes.core.expression as ex
+import bigframes.core.expression_types as ex_types
 from bigframes.core.ordering import OrderingExpression
 import bigframes.core.sql
 from bigframes.core.window_spec import RangeWindowBounds, RowsWindowBounds, WindowSpec
@@ -215,7 +216,7 @@ def filter(self, predicate: ex.Expression) -> UnorderedIR:
 
     def aggregate(
         self,
-        aggregations: typing.Sequence[tuple[ex.Aggregation, str]],
+        aggregations: typing.Sequence[tuple[ex_types.Aggregation, str]],
         by_column_ids: typing.Sequence[ex.DerefOp] = (),
         order_by: typing.Sequence[OrderingExpression] = (),
     ) -> UnorderedIR:
@@ -401,7 +402,7 @@ def isin_join(
 
     def project_window_op(
         self,
-        expression: ex.Aggregation,
+        expression: ex_types.Aggregation,
         window_spec: WindowSpec,
         output_name: str,
         *,
@@ -467,7 +468,9 @@ def project_window_op(
                     lambda x, y: x & y, per_col_does_count
                 ).cast(int)
                 observation_count = agg_compiler.compile_analytic(
-                    ex.UnaryAggregation(agg_ops.sum_op, ex.deref("_observation_count")),
+                    ex_types.UnaryAggregation(
+                        agg_ops.sum_op, ex.deref("_observation_count")
+                    ),
                     window,
                     bindings={"_observation_count": is_observation},
                 )
@@ -476,7 +479,7 @@ def project_window_op(
                 # notnull is just used to convert null values to non-null (FALSE) values to be counted
                 is_observation = inputs[0].notnull()
                 observation_count = agg_compiler.compile_analytic(
-                    ex.UnaryAggregation(
+                    ex_types.UnaryAggregation(
                         agg_ops.count_op, ex.deref("_observation_count")
                     ),
                     window,

bigframes/core/compile/ibis_compiler/aggregate_compiler.py

@@ -26,10 +26,10 @@
 import bigframes_vendored.ibis.expr.types as ibis_types
 import pandas as pd
 
+from bigframes.core import expression_types
 from bigframes.core.compile import constants as compiler_constants
 import bigframes.core.compile.ibis_compiler.scalar_op_compiler as scalar_compilers
 import bigframes.core.compile.ibis_types as compile_ibis_types
-import bigframes.core.expression as ex
 import bigframes.core.window_spec as window_spec
 import bigframes.operations.aggregations as agg_ops
 
@@ -48,19 +48,19 @@ def approx_quantiles(expression: float, number) -> List[float]:
 
 
 def compile_aggregate(
-    aggregate: ex.Aggregation,
+    aggregate: expression_types.Aggregation,
     bindings: typing.Dict[str, ibis_types.Value],
     order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
-    if isinstance(aggregate, ex.NullaryAggregation):
+    if isinstance(aggregate, expression_types.NullaryAggregation):
         return compile_nullary_agg(aggregate.op)
-    if isinstance(aggregate, ex.UnaryAggregation):
+    if isinstance(aggregate, expression_types.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         if not aggregate.op.order_independent:
             return compile_ordered_unary_agg(aggregate.op, input, order_by=order_by)  # type: ignore
         else:
             return compile_unary_agg(aggregate.op, input)  # type: ignore
-    elif isinstance(aggregate, ex.BinaryAggregation):
+    elif isinstance(aggregate, expression_types.BinaryAggregation):
         left = scalar_compiler.compile_expression(aggregate.left, bindings=bindings)
         right = scalar_compiler.compile_expression(aggregate.right, bindings=bindings)
         return compile_binary_agg(aggregate.op, left, right)  # type: ignore
@@ -69,16 +69,16 @@ def compile_aggregate(
 
 
 def compile_analytic(
-    aggregate: ex.Aggregation,
+    aggregate: expression_types.Aggregation,
     window: window_spec.WindowSpec,
     bindings: typing.Dict[str, ibis_types.Value],
 ) -> ibis_types.Value:
-    if isinstance(aggregate, ex.NullaryAggregation):
+    if isinstance(aggregate, expression_types.NullaryAggregation):
         return compile_nullary_agg(aggregate.op, window)
-    elif isinstance(aggregate, ex.UnaryAggregation):
+    elif isinstance(aggregate, expression_types.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         return compile_unary_agg(aggregate.op, input, window)  # type: ignore
-    elif isinstance(aggregate, ex.BinaryAggregation):
+    elif isinstance(aggregate, expression_types.BinaryAggregation):
         raise NotImplementedError("binary analytic operations not yet supported")
     else:
         raise ValueError(f"Unexpected analytic operation: {aggregate}")