feat: add GroupBy.size() to get number of rows in each group

Author: tswast

bigframes/core/blocks.py

@@ -1129,6 +1129,34 @@ def corr(self):
             index_labels=self.column_labels.names,
         )
 
+    def size(self):
+        """Returns a block object to compute the size(s) of this block."""
+        agg_specs = [
+            (ex.NullaryAggregation(agg_ops.SizeOp()), "size"),
+        ]
+        output_col_ids = [agg_spec[1] for agg_spec in agg_specs]
+        result_expr = self.expr.aggregate(agg_specs)
+
+        # # TODO: where do group by columns come in if as_index=False?
+        # aggregate_labels = self._get_labels_for_columns(
+        #     [agg[0] for agg in aggregations]
+        # )
+        # names: typing.List[Label] = []
+        # for by_col_id in by_column_ids:
+        #     if by_col_id in self.value_columns:
+        #         names.append(self.col_id_to_label[by_col_id])
+        #     else:
+        #         names.append(self.col_id_to_index_name[by_col_id])
+        return (
+            Block(
+                result_expr,
+                # index_columns=by_column_ids,
+                # column_labels=aggregate_labels,
+                # index_labels=names,
+            ),
+            output_col_ids,
+        )
+
     def _standard_stats(self, column_id) -> typing.Sequence[agg_ops.UnaryAggregateOp]:
         """
         Gets a standard set of stats to preemptively fetch for a column if

bigframes/core/compile/aggregate_compiler.py

@@ -35,6 +35,8 @@ def compile_aggregate(
     aggregate: ex.Aggregation,
     bindings: typing.Dict[str, ibis_types.Value],
 ) -> ibis_types.Value:
+    if isinstance(aggregate, ex.NullaryAggregation):
+        return compile_nullary_agg(aggregate.op)
     if isinstance(aggregate, ex.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         return compile_unary_agg(
@@ -54,7 +56,9 @@ def compile_analytic(
     window: window_spec.WindowSpec,
     bindings: typing.Dict[str, ibis_types.Value],
 ) -> ibis_types.Value:
-    if isinstance(aggregate, ex.UnaryAggregation):
+    if isinstance(aggregate, ex.NullaryAggregation):
+        return compile_nullary_agg(aggregate.op, window)
+    elif isinstance(aggregate, ex.UnaryAggregation):
         input = scalar_compiler.compile_expression(aggregate.arg, bindings=bindings)
         return compile_unary_agg(aggregate.op, input, window)
     elif isinstance(aggregate, ex.BinaryAggregation):
@@ -81,6 +85,14 @@ def compile_unary_agg(
     raise ValueError(f"Can't compile unrecognized operation: {op}")
 
 
+@functools.singledispatch
+def compile_nullary_agg(
+    op: agg_ops.WindowOp,
+    window: Optional[window_spec.WindowSpec] = None,
+) -> ibis_types.Value:
+    raise ValueError(f"Can't compile unrecognized operation: {op}")
+
+
 def numeric_op(operation):
     @functools.wraps(operation)
     def constrained_op(op, column: ibis_types.Column, window=None):
@@ -101,6 +113,12 @@ def constrained_op(op, column: ibis_types.Column, window=None):
 ### Specific Op implementations Below
 
 
+@compile_nullary_agg.register
+@numeric_op
+def _(op: agg_ops.SizeOp, window=None) -> ibis_types.NumericValue:
+    return _apply_window_if_present(vendored_ibis_ops.count(1), window)
+
+
 @compile_unary_agg.register
 @numeric_op
 def _(

bigframes/core/expression.py

@@ -40,6 +40,11 @@ class Aggregation(abc.ABC):
     op: agg_ops.WindowOp = dataclasses.field()
 
 
+@dataclasses.dataclass(frozen=True)
+class NullaryAggregation(Aggregation):
+    op: agg_ops.NullaryWindowOp = dataclasses.field()
+
+
 @dataclasses.dataclass(frozen=True)
 class UnaryAggregation(Aggregation):
     op: agg_ops.UnaryWindowOp = dataclasses.field()

bigframes/core/groupby/__init__.py

@@ -102,6 +102,16 @@ def __getitem__(
                 dropna=self._dropna,
             )
 
+    def size(self) -> df.DataFrame:
+        agg_block, _ = self._block.aggregate(
+            by_column_ids=self._by_col_ids,
+            # aggregations=aggregations,
+            dropna=self._dropna,
+        )
+        # agg_block = agg_block.with_column_labels(column_labels)
+        dataframe = df.DataFrame(agg_block)
+        return dataframe if self._as_index else self._convert_index(dataframe)
+
     def sum(self, numeric_only: bool = False, *args) -> df.DataFrame:
         if not numeric_only:
             self._raise_on_non_numeric("sum")

bigframes/operations/aggregations.py

@@ -33,6 +33,13 @@ def handles_ties(self):
         return False
 
 
+@dataclasses.dataclass(frozen=True)
+class NullaryWindowOp(WindowOp):
+    @property
+    def arguments(self) -> int:
+        return 0
+
+
 @dataclasses.dataclass(frozen=True)
 class UnaryWindowOp(WindowOp):
     @property
@@ -55,6 +62,13 @@ def arguments(self) -> int:
         ...
 
 
+@dataclasses.dataclass(frozen=True)
+class NullaryAggregateOp(AggregateOp, NullaryWindowOp):
+    @property
+    def arguments(self) -> int:
+        return 0
+
+
 @dataclasses.dataclass(frozen=True)
 class UnaryAggregateOp(AggregateOp, UnaryWindowOp):
     @property
@@ -69,6 +83,11 @@ def arguments(self) -> int:
         return 2
 
 
+@dataclasses.dataclass(frozen=True)
+class SizeOp(NullaryAggregateOp):
+    name: ClassVar[str] = "size"
+
+
 @dataclasses.dataclass(frozen=True)
 class SumOp(UnaryAggregateOp):
     name: ClassVar[str] = "sum"
@@ -270,6 +289,7 @@ class CovOp(BinaryAggregateOp):
     name: ClassVar[str] = "cov"
 
 
+size_op = SizeOp()
 sum_op = SumOp()
 mean_op = MeanOp()
 median_op = MedianOp()

third_party/bigframes_vendored/ibis/expr/operations/analytic.py

@@ -2,10 +2,17 @@
 
 from __future__ import annotations
 
+import ibis
 import ibis.expr.operations as ops
 import ibis.expr.rules as rlz
 
 
+@ibis.udf.agg.builtin
+def count(value: int) -> int:
+    """Count of a scalar."""
+    return 0  # pragma: NO COVER
+
+
 class FirstNonNullValue(ops.Analytic):
     """Retrieve the first element."""