diff --git a/bigframes/core/groupby/__init__.py b/bigframes/core/groupby/__init__.py index 5cb0e65729..2d351cf82d 100644 --- a/bigframes/core/groupby/__init__.py +++ b/bigframes/core/groupby/__init__.py @@ -88,6 +88,12 @@ def __getitem__( keys = list(key) else: keys = [key] + + bad_keys = [key for key in keys if key not in self._block.column_labels] + + if len(bad_keys) > 0: + raise KeyError(f"Columns not found: {str(bad_keys)[1:-1]}") + columns = [ col_id for col_id, label in self._col_id_labels.items() if label in keys ] diff --git a/tests/system/small/test_groupby.py b/tests/system/small/test_groupby.py index 8574860daa..2d5ae21bb4 100644 --- a/tests/system/small/test_groupby.py +++ b/tests/system/small/test_groupby.py @@ -421,6 +421,34 @@ def test_dataframe_groupby_getitem( pd.testing.assert_series_equal(pd_result, bf_result, check_dtype=False) +def test_dataframe_groupby_getitem_error( + scalars_df_index, + scalars_pandas_df_index, +): + col_names = ["float64_col", "int64_col", "bool_col", "string_col"] + with pytest.raises(KeyError, match="\"Columns not found: 'not_in_group'\""): + ( + scalars_df_index[col_names] + .groupby("string_col")["not_in_group"] + .min() + .to_pandas() + ) + + +def test_dataframe_groupby_getitem_multiple_columns_error( + scalars_df_index, + scalars_pandas_df_index, +): + col_names = ["float64_col", "int64_col", "bool_col", "string_col"] + with pytest.raises(KeyError, match="\"Columns not found: 'col1', 'col2'\""): + ( + scalars_df_index[col_names] + .groupby("string_col")["col1", "col2"] + .min() + .to_pandas() + ) + + def test_dataframe_groupby_getitem_list( scalars_df_index, scalars_pandas_df_index,