diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 837eb28f68..2b447a0190 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -358,8 +358,8 @@ def _convert_index(self, dataframe: df.DataFrame): def _raise_on_non_numeric(self, op: str): if not all( - dtype in dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE - for dtype in self._block.dtypes + self._column_type(col) in dtypes.NUMERIC_BIGFRAMES_TYPES_PERMISSIVE + for col in self._selected_cols ): raise NotImplementedError( f"'{op}' does not support non-numeric columns. " diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index c0f602a598..0f99a3e4db 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -2337,6 +2337,7 @@ def groupby( blocks.Label, bigframes.series.Series, typing.Sequence[typing.Union[blocks.Label, bigframes.series.Series]], + None, ] = None, *, level: typing.Optional[LevelsType] = None, diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 2919c167ef..b38dcaf5d1 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -371,3 +371,20 @@ def test_series_groupby_agg_list(scalars_df_index, scalars_pandas_df_index): pd.testing.assert_frame_equal( pd_result, bf_result_computed, check_dtype=False, check_names=False ) + + +def test_dataframe_groupby_nonnumeric_with_mean(): + df = pd.DataFrame( + { + "key1": ["a", "a", "a", "b"], + "key2": ["a", "a", "c", "c"], + "key3": [1, 2, 3, 4], + "key4": [1.6, 2, 3, 4], + } + ) + pd_result = df.groupby(["key1", "key2"]).mean() + bf_result = bpd.DataFrame(df).groupby(["key1", "key2"]).mean().to_pandas() + + pd.testing.assert_frame_equal( + pd_result, bf_result, check_index_type=False, check_dtype=False + )