feat: Support len() on Groupby objects (#2183)

Author: TrevorBergeron

bigframes/core/groupby/dataframe_group_by.py

@@ -177,6 +177,9 @@ def __iter__(self) -> Iterable[Tuple[blocks.Label, df.DataFrame]]:
             filtered_df = df.DataFrame(filtered_block)
             yield group_keys, filtered_df
 
+    def __len__(self) -> int:
+        return len(self.agg([]))
+
     def size(self) -> typing.Union[df.DataFrame, series.Series]:
         agg_block, _ = self._block.aggregate_size(
             by_column_ids=self._by_col_ids,

bigframes/core/groupby/series_group_by.py

@@ -108,6 +108,9 @@ def __iter__(self) -> Iterable[Tuple[blocks.Label, series.Series]]:
             filtered_series.name = self._value_name
             yield group_keys, filtered_series
 
+    def __len__(self) -> int:
+        return len(self.agg([]))
+
     def all(self) -> series.Series:
         return self._aggregate(agg_ops.all_op)
 
@@ -275,9 +278,9 @@ def agg(self, func=None) -> typing.Union[df.DataFrame, series.Series]:
         if column_names:
             agg_block = agg_block.with_column_labels(column_names)
 
-        if len(aggregations) > 1:
-            return df.DataFrame(agg_block)
-        return series.Series(agg_block)
+        if len(aggregations) == 1:
+            return series.Series(agg_block)
+        return df.DataFrame(agg_block)
 
     aggregate = agg
 

tests/system/small/test_groupby.py

@@ -61,6 +61,15 @@ def test_dataframe_groupby_head(scalars_df_index, scalars_pandas_df_index):
     pd.testing.assert_frame_equal(pd_result, bf_result, check_dtype=False)
 
 
+def test_dataframe_groupby_len(scalars_df_index, scalars_pandas_df_index):
+    col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"]
+
+    bf_result = len(scalars_df_index[col_names].groupby("bool_col"))
+    pd_result = len(scalars_pandas_df_index[col_names].groupby("bool_col"))
+
+    assert bf_result == pd_result
+
+
 def test_dataframe_groupby_median(scalars_df_index, scalars_pandas_df_index):
     col_names = ["int64_too", "float64_col", "int64_col", "bool_col", "string_col"]
     bf_result = (
@@ -668,6 +677,13 @@ def test_dataframe_groupby_last(
 # ==============
 
 
+def test_series_groupby_len(scalars_df_index, scalars_pandas_df_index):
+    bf_result = len(scalars_df_index.groupby("bool_col")["int64_col"])
+    pd_result = len(scalars_pandas_df_index.groupby("bool_col")["int64_col"])
+
+    assert bf_result == pd_result
+
+
 @pytest.mark.parametrize(
     ("agg"),
     [