rename expression_types.py to agg_expressions.py

Author: TrevorBergeron

bigframes/core/expression_types.py bigframes/core/agg_expressions.pybigframes/core/expression_types.py renamed to bigframes/core/agg_expressions.py

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

bigframes/core/array_value.py

@@ -21,7 +21,7 @@
 import pandas as pd
 
 import bigframes.constants
-from bigframes.core import expression_types
+from bigframes.core import agg_expressions
 import bigframes.core as core
 import bigframes.core.blocks as blocks
 import bigframes.core.expression as ex
@@ -133,7 +133,7 @@ def quantile(
     block, _ = block.aggregate(
         grouping_column_ids,
         tuple(
-            expression_types.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col))
+            agg_expressions.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=[expression_types.NullaryAggregation(agg_ops.size_op)],
+        aggregations=[agg_expressions.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 = expression_types.UnaryAggregation(
+        count_agg = agg_expressions.UnaryAggregation(
             agg_ops.count_op,
             ex.deref(col),
         )
-        moment3_agg = expression_types.UnaryAggregation(
+        moment3_agg = agg_expressions.UnaryAggregation(
             agg_ops.mean_op,
             ex.deref(delta3_ids[i]),
         )
-        variance_agg = expression_types.UnaryAggregation(
+        variance_agg = agg_expressions.UnaryAggregation(
             agg_ops.PopVarOp(),
             ex.deref(col),
         )
@@ -698,11 +698,11 @@ def kurt(
     # counts, moment4 for each column
     aggregations = []
     for i, col in enumerate(original_columns):
-        count_agg = expression_types.UnaryAggregation(agg_ops.count_op, ex.deref(col))
-        moment4_agg = expression_types.UnaryAggregation(
+        count_agg = agg_expressions.UnaryAggregation(agg_ops.count_op, ex.deref(col))
+        moment4_agg = agg_expressions.UnaryAggregation(
             agg_ops.mean_op, ex.deref(delta4_ids[i])
         )
-        variance_agg = expression_types.UnaryAggregation(
+        variance_agg = agg_expressions.UnaryAggregation(
             agg_ops.PopVarOp(), ex.deref(col)
         )
         aggregations.extend([count_agg, moment4_agg, variance_agg])

bigframes/core/block_transforms.py

@@ -51,12 +51,12 @@
 from bigframes import session
 from bigframes._config import sampling_options
 import bigframes.constants
-from bigframes.core import expression_types, local_data
+from bigframes.core import agg_expressions, local_data
 import bigframes.core as core
+import bigframes.core.agg_expressions as ex_types
 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
@@ -1144,7 +1144,7 @@ def apply_window_op(
         skip_reproject_unsafe: bool = False,
         never_skip_nulls: bool = False,
     ) -> typing.Tuple[Block, str]:
-        agg_expr = expression_types.UnaryAggregation(op, ex.deref(column))
+        agg_expr = agg_expressions.UnaryAggregation(op, ex.deref(column))
         return self.apply_analytic(
             agg_expr,
             window_spec,
@@ -1156,7 +1156,7 @@ def apply_window_op(
 
     def apply_analytic(
         self,
-        agg_expr: expression_types.Aggregation,
+        agg_expr: agg_expressions.Aggregation,
         window: windows.WindowSpec,
         result_label: Label,
         *,
@@ -1249,9 +1249,9 @@ def aggregate_all_and_stack(
         if axis_n == 0:
             aggregations = [
                 (
-                    expression_types.UnaryAggregation(operation, ex.deref(col_id))
+                    agg_expressions.UnaryAggregation(operation, ex.deref(col_id))
                     if isinstance(operation, agg_ops.UnaryAggregateOp)
-                    else expression_types.NullaryAggregation(operation),
+                    else agg_expressions.NullaryAggregation(operation),
                     col_id,
                 )
                 for col_id in self.value_columns
@@ -1281,7 +1281,7 @@ def aggregate_size(
         """Returns a block object to compute the size(s) of groups."""
         agg_specs = [
             (
-                expression_types.NullaryAggregation(agg_ops.SizeOp()),
+                agg_expressions.NullaryAggregation(agg_ops.SizeOp()),
                 guid.generate_guid(),
             ),
         ]
@@ -1354,7 +1354,7 @@ def remap_f(x):
     def aggregate(
         self,
         by_column_ids: typing.Sequence[str] = (),
-        aggregations: typing.Sequence[expression_types.Aggregation] = (),
+        aggregations: typing.Sequence[agg_expressions.Aggregation] = (),
         column_labels: Optional[pd.Index] = None,
         *,
         dropna: bool = True,
@@ -1423,9 +1423,9 @@ def get_stat(
 
         aggregations = [
             (
-                expression_types.UnaryAggregation(stat, ex.deref(column_id))
+                agg_expressions.UnaryAggregation(stat, ex.deref(column_id))
                 if isinstance(stat, agg_ops.UnaryAggregateOp)
-                else expression_types.NullaryAggregation(stat),
+                else agg_expressions.NullaryAggregation(stat),
                 stat.name,
             )
             for stat in stats_to_fetch
@@ -1451,7 +1451,7 @@ def get_binary_stat(
         # TODO(kemppeterson): Add a cache here.
         aggregations = [
             (
-                expression_types.BinaryAggregation(
+                agg_expressions.BinaryAggregation(
                     stat, ex.deref(column_id_left), ex.deref(column_id_right)
                 ),
                 f"{stat.name}_{column_id_left}{column_id_right}",
@@ -1478,9 +1478,9 @@ def summarize(
         labels = pd.Index([stat.name for stat in stats])
         aggregations = [
             (
-                expression_types.UnaryAggregation(stat, ex.deref(col_id))
+                agg_expressions.UnaryAggregation(stat, ex.deref(col_id))
                 if isinstance(stat, agg_ops.UnaryAggregateOp)
-                else expression_types.NullaryAggregation(stat),
+                else agg_expressions.NullaryAggregation(stat),
                 f"{col_id}-{stat.name}",
             )
             for stat in stats
@@ -1754,7 +1754,7 @@ def pivot(
 
         block = block.select_columns(column_ids)
         aggregations = [
-            expression_types.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col_id))
+            agg_expressions.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col_id))
             for col_id in column_ids
         ]
         result_block, _ = block.aggregate(
@@ -2022,7 +2022,7 @@ def _generate_resample_label(
 
         agg_specs = [
             (
-                expression_types.UnaryAggregation(agg_ops.min_op, ex.deref(col_id)),
+                agg_expressions.UnaryAggregation(agg_ops.min_op, ex.deref(col_id)),
                 guid.generate_guid(),
             ),
         ]

bigframes/core/blocks.py

@@ -30,12 +30,12 @@
 import pyarrow as pa
 
 from bigframes.core import utils
+import bigframes.core.agg_expressions as ex_types
 import bigframes.core.compile.googlesql
 import bigframes.core.compile.ibis_compiler.aggregate_compiler as agg_compiler
 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

bigframes/core/compile/compiled.py

@@ -26,7 +26,7 @@
 import bigframes_vendored.ibis.expr.types as ibis_types
 import pandas as pd
 
-from bigframes.core import expression_types
+from bigframes.core import agg_expressions
 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
@@ -48,19 +48,19 @@ def approx_quantiles(expression: float, number) -> List[float]:
 
 
 def compile_aggregate(
-    aggregate: expression_types.Aggregation,
+    aggregate: agg_expressions.Aggregation,
     bindings: typing.Dict[str, ibis_types.Value],
     order_by: typing.Sequence[ibis_types.Value] = [],
 ) -> ibis_types.Value:
-    if isinstance(aggregate, expression_types.NullaryAggregation):
+    if isinstance(aggregate, agg_expressions.NullaryAggregation):
         return compile_nullary_agg(aggregate.op)
-    if isinstance(aggregate, expression_types.UnaryAggregation):
+    if isinstance(aggregate, agg_expressions.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, expression_types.BinaryAggregation):
+    elif isinstance(aggregate, agg_expressions.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: expression_types.Aggregation,
+    aggregate: agg_expressions.Aggregation,
     window: window_spec.WindowSpec,
     bindings: typing.Dict[str, ibis_types.Value],
 ) -> ibis_types.Value:
-    if isinstance(aggregate, expression_types.NullaryAggregation):
+    if isinstance(aggregate, agg_expressions.NullaryAggregation):
         return compile_nullary_agg(aggregate.op, window)
-    elif isinstance(aggregate, expression_types.UnaryAggregation):
+    elif isinstance(aggregate, agg_expressions.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         return compile_unary_agg(aggregate.op, input, window)  # type: ignore
-    elif isinstance(aggregate, expression_types.BinaryAggregation):
+    elif isinstance(aggregate, agg_expressions.BinaryAggregation):
         raise NotImplementedError("binary analytic operations not yet supported")
     else:
         raise ValueError(f"Unexpected analytic operation: {aggregate}")

bigframes/core/compile/ibis_compiler/aggregate_compiler.py

@@ -22,7 +22,7 @@
 import pandas as pd
 
 import bigframes.core
-from bigframes.core import expression_types, identifiers, nodes, ordering, window_spec
+from bigframes.core import agg_expressions, identifiers, nodes, ordering, window_spec
 from bigframes.core.compile.polars import lowering
 import bigframes.core.expression as ex
 import bigframes.core.guid as guid
@@ -443,15 +443,15 @@ class PolarsAggregateCompiler:
 
         def get_args(
             self,
-            agg: expression_types.Aggregation,
+            agg: agg_expressions.Aggregation,
         ) -> Sequence[pl.Expr]:
             """Prepares arguments for aggregation by compiling them."""
-            if isinstance(agg, expression_types.NullaryAggregation):
+            if isinstance(agg, agg_expressions.NullaryAggregation):
                 return []
-            elif isinstance(agg, expression_types.UnaryAggregation):
+            elif isinstance(agg, agg_expressions.UnaryAggregation):
                 arg = self.scalar_compiler.compile_expression(agg.arg)
                 return [arg]
-            elif isinstance(agg, expression_types.BinaryAggregation):
+            elif isinstance(agg, agg_expressions.BinaryAggregation):
                 larg = self.scalar_compiler.compile_expression(agg.left)
                 rarg = self.scalar_compiler.compile_expression(agg.right)
                 return [larg, rarg]
@@ -460,13 +460,13 @@ def get_args(
                 f"Aggregation {agg} not yet supported in polars engine."
             )
 
-        def compile_agg_expr(self, expr: expression_types.Aggregation):
-            if isinstance(expr, expression_types.NullaryAggregation):
+        def compile_agg_expr(self, expr: agg_expressions.Aggregation):
+            if isinstance(expr, agg_expressions.NullaryAggregation):
                 inputs: Tuple = ()
-            elif isinstance(expr, expression_types.UnaryAggregation):
+            elif isinstance(expr, agg_expressions.UnaryAggregation):
                 assert isinstance(expr.arg, ex.DerefOp)
                 inputs = (expr.arg.id.sql,)
-            elif isinstance(expr, expression_types.BinaryAggregation):
+            elif isinstance(expr, agg_expressions.BinaryAggregation):
                 assert isinstance(expr.left, ex.DerefOp)
                 assert isinstance(expr.right, ex.DerefOp)
                 inputs = (
@@ -770,7 +770,7 @@ def _aggregate(
         self,
         df: pl.LazyFrame,
         aggregations: Sequence[
-            Tuple[expression_types.Aggregation, identifiers.ColumnId]
+            Tuple[agg_expressions.Aggregation, identifiers.ColumnId]
         ],
         grouping_keys: Tuple[ex.DerefOp, ...],
     ) -> pl.LazyFrame:
@@ -860,7 +860,7 @@ def compile_window(self, node: nodes.WindowOpNode):
     def _calc_row_analytic_func(
         self,
         frame: pl.LazyFrame,
-        agg_expr: expression_types.Aggregation,
+        agg_expr: agg_expressions.Aggregation,
         window: window_spec.WindowSpec,
         name: str,
     ) -> pl.LazyFrame: