googleapis · chelsea-lin · Oct 6, 2024 · Sep 30, 2024 · Sep 30, 2024 · Oct 1, 2024
@@ -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
        ]

@@ -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,