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."""