BUG: Inconsistent behaviour in DataFrameGroupBy when selecting a subset of columns · Issue #44821 · pandas-dev/pandas (original) (raw)

• I have checked that this issue has not already been reported.
• I have confirmed this bug exists on the latest version of pandas.
• I have confirmed this bug exists on the master branch of pandas.

Reproducible Example

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.

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

INSTALLED VERSIONS

commit : 945c9ed
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