BUG: iterating on a subset of columns in a GroupBy object (pandas-dev#44821)

Fixes issue pandas-dev#44821.

When trying to iterate on a subset of columns in a GroupBy object,
it returned all columns, instead of the selected subset.

GroupBy.__iter__ used self.obj instead of self._selected_obj (see
PR pandas-dev#6570).

Author: maurosilber

pandas/core/groupby/groupby.py

@@ -758,7 +758,7 @@ def __iter__(self) -> Iterator[tuple[Hashable, NDFrameT]]:
         Generator yielding sequence of (name, subsetted object)
         for each group
         """
-        return self.grouper.get_iterator(self.obj, axis=self.axis)
+        return self.grouper.get_iterator(self._selected_obj, axis=self.axis)
 
 
 # To track operations that expand dimensions, like ohlc

pandas/tests/groupby/test_indexing.py

@@ -285,3 +285,14 @@ def test_column_axis(column_group_df):
     expected = column_group_df.iloc[:, [1, 3]]
 
     tm.assert_frame_equal(result, expected)
+
+
+def test_columns_on_iter():
+    # GitHub issue #44821
+    df = pd.DataFrame({k: range(10) for k in "ABC"})
+
+    # Group-by and select columns
+    cols = ["A", "B"]
+    for _, dg in df.groupby(df.A < 4)[cols]:
+        tm.assert_index_equal(dg.columns, pd.Index(cols))
+        assert "C" not in dg.columns