Repository: pandas-dev/pandas (Python)

Issue:
BUG: Inconsistent behaviour in DataFrameGroupBy when selecting a subset of columns
### 

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the [latest version](https://pandas.pydata.org/docs/whatsnew/index.html) of pandas.

- [X] I have confirmed this bug exists on the master branch of pandas.


### Reproducible Example

```python
import pandas as pd

df = pd.DataFrame({k: range(10) for k in "ABC"})

# Group by and select columns
dg = df.groupby(df.A < 4)[["A", "B"]]


# Apply on GroupBy
print(dg.apply(lambda x: [x.shape, x.columns]))

# A
# False    [(6, 2), [A, B]]
# True     [(4, 2), [A, B]]
# dtype: object


# Iterate on GroupBy
for ix, dg_ix in dg:
    print(ix, dg_ix.shape, dg_ix.columns)

# False (6, 3) Index(['A', 'B', 'C'], dtype='object')
# True (4, 3) Index(['A', 'B', 'C'], dtype='object')
```


### Issue Description

A subset of columns is not (always) selected with `DataFrameGroupBy.__getitem__`. Selecting a subset of column works for `DataFrameGroupBy.apply` but not for `DataFrameGroupBy.__iter__`.

### Expected Behavior

I expected `dg_ix` to contain only `A` and `B` as columns.

```python
# Iterate on GroupBy
for ix, dg_ix in dg:
    print(ix, dg_ix.shape, dg_ix.columns)

# False (6, 2) Index(['A', 'B'], dtype='object')
# True (4, 2) Index(['A', 'B'], dtype='object')
```

### Installed Versions

<details>

INSTALLED VERSIONS
------------------
commit           : 945c9ed766a61c7d2c0a7cbb251b6edebf9cb7d5
python           : 3.9.7.final.0
python-bits      : 64
OS               : Linux
OS-release       : 5.14.18-300.fc35.x86_64
Version          : #1 SMP Fri Nov 12 16:43:17 UTC 2021
machine          : x86_64
processor        : x86_64
byteorder        : little
LC_ALL           : None
LANG             : en_US.UTF-8
LOCALE           : en_US.UTF-8

pandas           : 1.3.4
numpy            : 1.21.4
pytz             : 2021.3
dateutil         : 2.8.2
pip              : 21.3.1
setuptools       : 59.1.0
Cython           : None
pytest           : 5.4.3
hypothesis       : None
sphinx           : None
blosc            : None
feather          : None
xlsxwriter       : None
lxml.etree       : None
html5lib         : None
pymysql          : None
psycopg2         : None
jinja2           : 3.0.3
IPython          : 7.29.0
pandas_datareader: None
bs4              : None
bottleneck       : None
fsspec           : None
fastparquet      : None
gcsfs            : None
matplotlib       : 3.3.2
numexpr          : None
odfpy            : None
openpyxl         : None
pandas_gbq       : None
pyarrow          : None
pyxlsb           : None
s3fs             : None
scipy            : 1.7.2
sqlalchemy       : None
tables           : None
tabulate         : None
xarray           : 0.20.1
xlrd             : None
xlwt             : None
numba            : 0.53.1

</details>



Discussion:
I found this old issues discussing selection of multiple columns in GroupBy objects:

- https://github.com/pandas-dev/pandas/issues/5264
- https://github.com/pandas-dev/pandas/issues/6524
Thanks @maurosilber for the report and investigation.

> I expected `dg_ix` to contain only `A` and `B` as columns.

makes sense.

> A subset of columns is not (always) selected with `DataFrameGroupBy.__getitem__`.

This looks like a bug to me with just the `__iter__` method? (that was perhaps missed in the previous issue #5264)

PR to fix or further investigation welcome.



I did a pull request, but the following test
https://github.com/pandas-dev/pandas/blob/63e7ef15b3008aa006b563cd2f5ea70cdb0e3ec2/pandas/tests/window/test_base_indexer.py#L411-L454
fails in line 452
https://github.com/pandas-dev/pandas/blob/63e7ef15b3008aa006b563cd2f5ea70cdb0e3ec2/pandas/tests/window/test_base_indexer.py#L452
with
```python
KeyError: "None of ['c'] are in the columns"
```

Which should be expected, as it wasn't selected in line 440
https://github.com/pandas-dev/pandas/blob/63e7ef15b3008aa006b563cd2f5ea70cdb0e3ec2/pandas/tests/window/test_base_indexer.py#L440

Should I correct the test by selecting `"c"`? It might break someone's code if it relies on this behaviour.



Patch:
diff --git a/doc/source/whatsnew/v1.4.0.rst b/doc/source/whatsnew/v1.4.0.rst
index 040c424fb4127..1570fc5e59093 100644
--- a/doc/source/whatsnew/v1.4.0.rst
+++ b/doc/source/whatsnew/v1.4.0.rst
@@ -812,6 +812,7 @@ Groupby/resample/rolling
 - Bug in :meth:`GroupBy.nth` failing on ``axis=1`` (:issue:`43926`)
 - Fixed bug in :meth:`Series.rolling` and :meth:`DataFrame.rolling` not respecting right bound on centered datetime-like windows, if the index contain duplicates (:issue:`3944`)
 - Bug in :meth:`Series.rolling` and :meth:`DataFrame.rolling` when using a :class:`pandas.api.indexers.BaseIndexer` subclass that returned unequal start and end arrays would segfault instead of raising a ``ValueError`` (:issue:`44470`)
+- Fixed bug in :meth:`GroupBy.__iter__` after selecting a subset of columns in a :class:`GroupBy` object, which returned all columns instead of the chosen subset (:issue:`#44821`)
 - Bug in :meth:`Groupby.rolling` when non-monotonic data passed, fails to correctly raise ``ValueError`` (:issue:`43909`)
 - Fixed bug where grouping by a :class:`Series` that has a categorical data type and length unequal to the axis of grouping raised ``ValueError`` (:issue:`44179`)
 
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
index a1866e3bdc9f6..acf65a464a45f 100644
--- a/pandas/core/groupby/groupby.py
+++ b/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


Test patch:
diff --git a/pandas/tests/groupby/test_indexing.py b/pandas/tests/groupby/test_indexing.py
index aea659445801b..0caa17f387a94 100644
--- a/pandas/tests/groupby/test_indexing.py
+++ b/pandas/tests/groupby/test_indexing.py
@@ -288,6 +288,17 @@ def test_column_axis(column_group_df):
     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
+
+
 @pytest.mark.parametrize("func", [list, pd.Index, pd.Series, np.array])
 def test_groupby_duplicated_columns(func):
     # GH#44924
diff --git a/pandas/tests/window/test_base_indexer.py b/pandas/tests/window/test_base_indexer.py
index a7ad409683ec8..5593aa8351c69 100644
--- a/pandas/tests/window/test_base_indexer.py
+++ b/pandas/tests/window/test_base_indexer.py
@@ -437,7 +437,7 @@ def test_rolling_groupby_with_fixed_forward_many(group_keys, window_size):
     result = df.groupby("a")["b"].rolling(window=indexer, min_periods=1).sum()
     result.index.names = ["a", "c"]
 
-    groups = df.groupby("a")[["a", "b"]]
+    groups = df.groupby("a")[["a", "b", "c"]]
     manual = concat(
         [
             g.assign(