fix: fix generic error message when entering an incorrect column name (#1031)

* fix: fix generic error message when entering incorrect column name

* fix varibale initialization

* Add a test case

* Update code and fix failling tests

* Add test case for multiple columns

---------

Co-authored-by: Chelsea Lin <[email protected]>

Author: arwas11

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
         ]

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,