refactor: Generalize Block.aggregate to non-unary aggregates (#1304)

Author: TrevorBergeron

bigframes/core/block_transforms.py

@@ -129,12 +129,16 @@ def quantile(
                 window_spec=window,
             )
             quantile_cols.append(quantile_col)
-    block, results = block.aggregate(
+    block, _ = block.aggregate(
         grouping_column_ids,
-        tuple((col, agg_ops.AnyValueOp()) for col in quantile_cols),
+        tuple(
+            ex.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col))
+            for col in quantile_cols
+        ),
+        column_labels=pd.Index(labels),
         dropna=dropna,
     )
-    return block.select_columns(results).with_column_labels(labels)
+    return block
 
 
 def interpolate(block: blocks.Block, method: str = "linear") -> blocks.Block:
@@ -355,7 +359,7 @@ def value_counts(
     block, dummy = block.create_constant(1)
     block, agg_ids = block.aggregate(
         by_column_ids=columns,
-        aggregations=[(dummy, agg_ops.count_op)],
+        aggregations=[ex.UnaryAggregation(agg_ops.count_op, ex.deref(dummy))],
         dropna=dropna,
     )
     count_id = agg_ids[0]
@@ -589,9 +593,18 @@ def skew(
     # counts, moment3 for each column
     aggregations = []
     for i, col in enumerate(original_columns):
-        count_agg = (col, agg_ops.count_op)
-        moment3_agg = (delta3_ids[i], agg_ops.mean_op)
-        variance_agg = (col, agg_ops.PopVarOp())
+        count_agg = ex.UnaryAggregation(
+            agg_ops.count_op,
+            ex.deref(col),
+        )
+        moment3_agg = ex.UnaryAggregation(
+            agg_ops.mean_op,
+            ex.deref(delta3_ids[i]),
+        )
+        variance_agg = ex.UnaryAggregation(
+            agg_ops.PopVarOp(),
+            ex.deref(col),
+        )
         aggregations.extend([count_agg, moment3_agg, variance_agg])
 
     block, agg_ids = block.aggregate(
@@ -631,9 +644,9 @@ def kurt(
     # counts, moment4 for each column
     aggregations = []
     for i, col in enumerate(original_columns):
-        count_agg = (col, agg_ops.count_op)
-        moment4_agg = (delta4_ids[i], agg_ops.mean_op)
-        variance_agg = (col, agg_ops.PopVarOp())
+        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))
         aggregations.extend([count_agg, moment4_agg, variance_agg])
 
     block, agg_ids = block.aggregate(

bigframes/core/blocks.py

@@ -1174,38 +1174,31 @@ def remap_f(x):
     def aggregate(
         self,
         by_column_ids: typing.Sequence[str] = (),
-        aggregations: typing.Sequence[
-            typing.Tuple[
-                str, typing.Union[agg_ops.UnaryAggregateOp, agg_ops.NullaryAggregateOp]
-            ]
-        ] = (),
+        aggregations: typing.Sequence[ex.Aggregation] = (),
+        column_labels: Optional[pd.Index] = None,
         *,
         dropna: bool = True,
     ) -> typing.Tuple[Block, typing.Sequence[str]]:
         """
-        Apply aggregations to the block. Callers responsible for setting index column(s) after.
+        Apply aggregations to the block.
         Arguments:
             by_column_id: column id of the aggregation key, this is preserved through the transform and used as index.
             aggregations: input_column_id, operation tuples
-            as_index: if True, grouping keys will be index columns in result, otherwise they will be non-index columns.
             dropna: whether null keys should be dropped
         """
+        if column_labels is None:
+            column_labels = pd.Index(range(len(aggregations)))
+
         agg_specs = [
             (
-                ex.UnaryAggregation(operation, ex.deref(input_id))
-                if isinstance(operation, agg_ops.UnaryAggregateOp)
-                else ex.NullaryAggregation(operation),
+                aggregation,
                 guid.generate_guid(),
             )
-            for input_id, operation in aggregations
+            for aggregation in aggregations
         ]
         output_col_ids = [agg_spec[1] for agg_spec in agg_specs]
         result_expr = self.expr.aggregate(agg_specs, by_column_ids, dropna=dropna)
 
-        aggregate_labels = self._get_labels_for_columns(
-            [agg[0] for agg in aggregations]
-        )
-
         names: typing.List[Label] = []
         if len(by_column_ids) == 0:
             result_expr, label_id = result_expr.create_constant(0, pd.Int64Dtype())
@@ -1223,7 +1216,7 @@ def aggregate(
             Block(
                 result_expr,
                 index_columns=index_columns,
-                column_labels=aggregate_labels,
+                column_labels=column_labels,
                 index_labels=names,
             ),
             output_col_ids,
@@ -1561,7 +1554,10 @@ def pivot(
                 column_ids.append(masked_id)
 
         block = block.select_columns(column_ids)
-        aggregations = [(col_id, agg_ops.AnyValueOp()) for col_id in column_ids]
+        aggregations = [
+            ex.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(col_id))
+            for col_id in column_ids
+        ]
         result_block, _ = block.aggregate(
             by_column_ids=self.index_columns,
             aggregations=aggregations,

bigframes/core/groupby/__init__.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Sequence, Union
+from typing import Sequence, Tuple, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.pandas.core.groupby as vendored_pandas_groupby
@@ -26,6 +26,7 @@
 import bigframes.core as core
 import bigframes.core.block_transforms as block_ops
 import bigframes.core.blocks as blocks
+import bigframes.core.expression
 import bigframes.core.ordering as order
 import bigframes.core.utils as utils
 import bigframes.core.validations as validations
@@ -334,24 +335,19 @@ def agg(self, func=None, **kwargs) -> typing.Union[df.DataFrame, series.Series]:
             return self._agg_named(**kwargs)
 
     def _agg_string(self, func: str) -> df.DataFrame:
-        aggregations = [
-            (col_id, agg_ops.lookup_agg_func(func))
-            for col_id in self._aggregated_columns()
-        ]
+        ids, labels = self._aggregated_columns()
+        aggregations = [agg(col_id, agg_ops.lookup_agg_func(func)) for col_id in ids]
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
             aggregations=aggregations,
             dropna=self._dropna,
+            column_labels=labels,
         )
         dataframe = df.DataFrame(agg_block)
         return dataframe if self._as_index else self._convert_index(dataframe)
 
     def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
-        aggregations: typing.List[
-            typing.Tuple[
-                str, typing.Union[agg_ops.UnaryAggregateOp, agg_ops.NullaryAggregateOp]
-            ]
-        ] = []
+        aggregations: typing.List[bigframes.core.expression.Aggregation] = []
         column_labels = []
 
         want_aggfunc_level = any(utils.is_list_like(aggs) for aggs in func.values())
@@ -362,7 +358,7 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
                 funcs_for_id if utils.is_list_like(funcs_for_id) else [funcs_for_id]
             )
             for f in func_list:
-                aggregations.append((col_id, agg_ops.lookup_agg_func(f)))
+                aggregations.append(agg(col_id, agg_ops.lookup_agg_func(f)))
                 column_labels.append(label)
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
@@ -373,7 +369,10 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
             agg_block = agg_block.with_column_labels(
                 utils.combine_indices(
                     pd.Index(column_labels),
-                    pd.Index(agg[1].name for agg in aggregations),
+                    pd.Index(
+                        typing.cast(agg_ops.AggregateOp, agg.op).name
+                        for agg in aggregations
+                    ),
                 )
             )
         else:
@@ -382,34 +381,21 @@ def _agg_dict(self, func: typing.Mapping) -> df.DataFrame:
         return dataframe if self._as_index else self._convert_index(dataframe)
 
     def _agg_list(self, func: typing.Sequence) -> df.DataFrame:
+        ids, labels = self._aggregated_columns()
         aggregations = [
-            (col_id, agg_ops.lookup_agg_func(f))
-            for col_id in self._aggregated_columns()
-            for f in func
+            agg(col_id, agg_ops.lookup_agg_func(f)) for col_id in ids for f in func
         ]
 
         if self._block.column_labels.nlevels > 1:
             # Restructure MultiIndex for proper format: (idx1, idx2, func)
             # rather than ((idx1, idx2), func).
-            aggregated_columns = pd.MultiIndex.from_tuples(
-                [
-                    self._block.col_id_to_label[col_id]
-                    for col_id in self._aggregated_columns()
-                ],
-                names=[*self._block.column_labels.names],
-            ).to_frame(index=False)
-
             column_labels = [
-                tuple(col_id) + (f,)
-                for col_id in aggregated_columns.to_numpy()
-                for f in func
-            ]
-        else:
-            column_labels = [
-                (self._block.col_id_to_label[col_id], f)
-                for col_id in self._aggregated_columns()
+                tuple(label) + (f,)
+                for label in labels.to_frame(index=False).to_numpy()
                 for f in func
             ]
+        else:  # Single-level index
+            column_labels = [(label, f) for label in labels for f in func]
 
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
@@ -435,7 +421,7 @@ def _agg_named(self, **kwargs) -> df.DataFrame:
             if not isinstance(v, tuple) or (len(v) != 2):
                 raise TypeError("kwargs values must be 2-tuples of column, aggfunc")
             col_id = self._resolve_label(v[0])
-            aggregations.append((col_id, agg_ops.lookup_agg_func(v[1])))
+            aggregations.append(agg(col_id, agg_ops.lookup_agg_func(v[1])))
             column_labels.append(k)
         agg_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
@@ -470,15 +456,19 @@ def _raise_on_non_numeric(self, op: str):
             )
         return self
 
-    def _aggregated_columns(self, numeric_only: bool = False) -> typing.Sequence[str]:
+    def _aggregated_columns(
+        self, numeric_only: bool = False
+    ) -> Tuple[typing.Sequence[str], pd.Index]:
         valid_agg_cols: list[str] = []
-        for col_id in self._selected_cols:
+        offsets: list[int] = []
+        for i, col_id in enumerate(self._block.value_columns):
             is_numeric = (
                 self._column_type(col_id) in dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE
             )
-            if is_numeric or not numeric_only:
+            if (col_id in self._selected_cols) and (is_numeric or not numeric_only):
+                offsets.append(i)
                 valid_agg_cols.append(col_id)
-        return valid_agg_cols
+        return valid_agg_cols, self._block.column_labels.take(offsets)
 
     def _column_type(self, col_id: str) -> dtypes.Dtype:
         col_offset = self._block.value_columns.index(col_id)
@@ -488,11 +478,12 @@ def _column_type(self, col_id: str) -> dtypes.Dtype:
     def _aggregate_all(
         self, aggregate_op: agg_ops.UnaryAggregateOp, numeric_only: bool = False
     ) -> df.DataFrame:
-        aggregated_col_ids = self._aggregated_columns(numeric_only=numeric_only)
-        aggregations = [(col_id, aggregate_op) for col_id in aggregated_col_ids]
+        aggregated_col_ids, labels = self._aggregated_columns(numeric_only=numeric_only)
+        aggregations = [agg(col_id, aggregate_op) for col_id in aggregated_col_ids]
         result_block, _ = self._block.aggregate(
             by_column_ids=self._by_col_ids,
             aggregations=aggregations,
+            column_labels=labels,
             dropna=self._dropna,
         )
         dataframe = df.DataFrame(result_block)
@@ -508,7 +499,7 @@ def _apply_window_op(
         window_spec = window or window_specs.cumulative_rows(
             grouping_keys=tuple(self._by_col_ids)
         )
-        columns = self._aggregated_columns(numeric_only=numeric_only)
+        columns, _ = self._aggregated_columns(numeric_only=numeric_only)
         block, result_ids = self._block.multi_apply_window_op(
             columns, op, window_spec=window_spec
         )
@@ -639,11 +630,11 @@ def prod(self, *args) -> series.Series:
     def agg(self, func=None) -> typing.Union[df.DataFrame, series.Series]:
         column_names: list[str] = []
         if isinstance(func, str):
-            aggregations = [(self._value_column, agg_ops.lookup_agg_func(func))]
+            aggregations = [agg(self._value_column, agg_ops.lookup_agg_func(func))]
             column_names = [func]
         elif utils.is_list_like(func):
             aggregations = [
-                (self._value_column, agg_ops.lookup_agg_func(f)) for f in func
+                agg(self._value_column, agg_ops.lookup_agg_func(f)) for f in func
             ]
             column_names = list(func)
         else:
@@ -756,7 +747,7 @@ def expanding(self, min_periods: int = 1) -> windows.Window:
     def _aggregate(self, aggregate_op: agg_ops.UnaryAggregateOp) -> series.Series:
         result_block, _ = self._block.aggregate(
             self._by_col_ids,
-            ((self._value_column, aggregate_op),),
+            (agg(self._value_column, aggregate_op),),
             dropna=self._dropna,
         )
 
@@ -781,3 +772,13 @@ def _apply_window_op(
             window_spec=window_spec,
         )
         return series.Series(block.select_column(result_id))
+
+
+def agg(input: str, op: agg_ops.AggregateOp) -> bigframes.core.expression.Aggregation:
+    if isinstance(op, agg_ops.UnaryAggregateOp):
+        return bigframes.core.expression.UnaryAggregation(
+            op, bigframes.core.expression.deref(input)
+        )
+    else:
+        assert isinstance(op, agg_ops.NullaryAggregateOp)
+        return bigframes.core.expression.NullaryAggregation(op)

bigframes/series.py

@@ -1062,7 +1062,9 @@ def mode(self) -> Series:
         # Approach: Count each value, return each value for which count(x) == max(counts))
         block, agg_ids = block.aggregate(
             by_column_ids=[self._value_column],
-            aggregations=((self._value_column, agg_ops.count_op),),
+            aggregations=(
+                ex.UnaryAggregation(agg_ops.count_op, ex.deref(self._value_column)),
+            ),
         )
         value_count_col_id = agg_ids[0]
         block, max_value_count_col_id = block.apply_window_op(
@@ -1675,7 +1677,8 @@ def unique(self, keep_order=True) -> Series:
             return self.drop_duplicates()
         block, result = self._block.aggregate(
             [self._value_column],
-            [(self._value_column, agg_ops.AnyValueOp())],
+            [ex.UnaryAggregation(agg_ops.AnyValueOp(), ex.deref(self._value_column))],
+            column_labels=self._block.column_labels,
             dropna=False,
         )
         return Series(block.select_columns(result).reset_index())