Repository: pandas-dev/pandas
Pull request: BUG: Attributes are lost when subsetting columns in groupby
- [x] closes #9959
- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry

Avoiding the behavior of Series here, e.g. `df.groupby('a')['b']` as I think that will involve some API changes. Will follow up: ref #35443

Problem statement:
columns selection after groupby reset group_keys to True
```
df = pd.DataFrame({"g":[1, 2, 1, 2, 1], "a":range(5), "b":range(1, 6), "c":range(2, 7)})
g = df.groupby("g", group_keys=False)
print g.group_keys  #False
print g[["a", "b", "c"]].group_keys  #True
```

So the `apply()` results are different for these two case:

```
print g.apply(lambda x:x[:2])
```

ouput:

```
   a  b  c  g
0  0  1  2  1
2  2  3  4  1
1  1  2  3  2
3  3  4  5  2
```

but 

```
print g[["a", "b", "c"]].apply(lambda x:x[:2])
```

output:

```
     a  b  c
g           
1 0  0  1  2
  2  2  3  4
2 1  1  2  3
  3  3  4  5
```

Hints from issue discussion:
Came across the same problem today. I am surprised that this problem dates back to as far as 2 years ago. Same problem still occurs in the current pandas version.

After poking into the code, I believe the problem is in `DataFrameGroupBy._gotitem`. It constructs a new GroupBy object and returns it. However, it does not pass the original `group_keys` to the constructor, so the default `group_keys=True` is used.

```Python
def _gotitem(self, key, ndim, subset=None):
        if ndim == 2:
            if subset is None:
                subset = self.obj
            return DataFrameGroupBy(subset, self.grouper, selection=key,
                                    grouper=self.grouper,
                                    exclusions=self.exclusions,
                                    as_index=self.as_index)
        elif ndim == 1:
            if subset is None:
                subset = self.obj[key]
            return SeriesGroupBy(subset, selection=key,
                                 grouper=self.grouper)

        raise AssertionError("invalid ndim for _gotitem")
```
Any plan to fix this? I think it would be kind of annoying if the behavior of `group_keys=False` is not consistent in column selection.
@jbrockmendel  @mroeschke  can you please label this issue? I think it was missed because it wasn't tag.
This problem is indeed annoying.
Possibly xref #14927

Commit message: Changed some tests from skip to xfail

Patch to review:
diff --git a/doc/source/whatsnew/v1.2.0.rst b/doc/source/whatsnew/v1.2.0.rst
index 8ec75b4846ae2..57ef72e90a27a 100644
--- a/doc/source/whatsnew/v1.2.0.rst
+++ b/doc/source/whatsnew/v1.2.0.rst
@@ -251,6 +251,8 @@ Groupby/resample/rolling
 - Bug in :meth:`DataFrameGroupBy.apply` that would some times throw an erroneous ``ValueError`` if the grouping axis had duplicate entries (:issue:`16646`)
 - Bug when combining methods :meth:`DataFrame.groupby` with :meth:`DataFrame.resample` and :meth:`DataFrame.interpolate` raising an ``TypeError`` (:issue:`35325`)
 - Bug in :meth:`DataFrameGroupBy.apply` where a non-nuisance grouping column would be dropped from the output columns if another groupby method was called before ``.apply()`` (:issue:`34656`)
+- Bug when subsetting columns on a :class:`~pandas.core.groupby.DataFrameGroupBy` (e.g. ``df.groupby('a')[['b']])``) would reset the attributes ``axis``, ``dropna``, ``group_keys``, ``level``, ``mutated``, ``sort``, and ``squeeze`` to their default values. (:issue:`9959`)
+-
 
 Reshaping
 ^^^^^^^^^
diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py
index 166631e69f523..3809df8b36488 100644
--- a/pandas/core/groupby/generic.py
+++ b/pandas/core/groupby/generic.py
@@ -1650,17 +1650,32 @@ def _gotitem(self, key, ndim: int, subset=None):
             return DataFrameGroupBy(
                 subset,
                 self.grouper,
-                selection=key,
+                axis=self.axis,
+                level=self.level,
                 grouper=self.grouper,
                 exclusions=self.exclusions,
+                selection=key,
                 as_index=self.as_index,
+                sort=self.sort,
+                group_keys=self.group_keys,
+                squeeze=self.squeeze,
                 observed=self.observed,
+                mutated=self.mutated,
+                dropna=self.dropna,
             )
         elif ndim == 1:
             if subset is None:
                 subset = self.obj[key]
             return SeriesGroupBy(
-                subset, selection=key, grouper=self.grouper, observed=self.observed
+                subset,
+                level=self.level,
+                grouper=self.grouper,
+                selection=key,
+                sort=self.sort,
+                group_keys=self.group_keys,
+                squeeze=self.squeeze,
+                observed=self.observed,
+                dropna=self.dropna,
             )
 
         raise AssertionError("invalid ndim for _gotitem")
diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py
index 8c51ebf89f5c0..bf1c58f771cd1 100644
--- a/pandas/tests/groupby/test_groupby.py
+++ b/pandas/tests/groupby/test_groupby.py
@@ -2069,3 +2069,36 @@ def test_group_on_two_row_multiindex_returns_one_tuple_key():
     assert len(result) == 1
     key = (1, 2)
     assert (result[key] == expected[key]).all()
+
+
+@pytest.mark.parametrize(
+    "attr, value",
+    [
+        ("axis", 1),
+        ("level", "a"),
+        ("as_index", False),
+        ("sort", False),
+        ("group_keys", False),
+        ("squeeze", True),
+        ("observed", True),
+        ("dropna", False),
+    ],
+)
+@pytest.mark.filterwarnings(
+    "ignore:The `squeeze` parameter is deprecated:FutureWarning"
+)
+def test_subsetting_columns_keeps_attrs(attr, value):
+    # GH 9959 - When subsetting columns, don't drop attributes
+    df = pd.DataFrame({"a": [1], "b": [2], "c": [3]})
+    if attr != "axis":
+        df = df.set_index("a")
+
+    expected = df.groupby("a", **{attr: value})
+
+    result = expected[["b"]]
+    assert getattr(result, attr) == getattr(expected, attr)
+
+    if attr in ("axis", "as_index"):
+        pytest.xfail("GH 35443: Attribute currently not passed on to series")
+    result = expected["b"]
+    assert getattr(result, attr) == getattr(expected, attr)


Code review comments:
diff_hunk: @@ -2069,3 +2069,36 @@ def test_group_on_two_row_multiindex_returns_one_tuple_key():
     assert len(result) == 1
     key = (1, 2)
     assert (result[key] == expected[key]).all()
+
+
+@pytest.mark.parametrize(
+    "attr, value",
+    [
+        ("axis", 1),
+        ("level", "a"),
+        ("as_index", False),
+        ("sort", False),
+        ("group_keys", False),
+        ("squeeze", True),
+        ("observed", True),
+        ("dropna", False),
+    ],
+)
+@pytest.mark.filterwarnings(
+    "ignore:The `squeeze` parameter is deprecated:FutureWarning"
+)
+def test_subsetting_columns_keeps_attrs(attr, value):
+    # GH 9959 - When subsetting columns, don't drop attributes
+    df = pd.DataFrame({"a": [1], "b": [2], "c": [3]})
+    if attr != "axis":
+        df = df.set_index("a")
+
+    expected = df.groupby("a", **{attr: value})
+
+    result = expected[["b"]]
path: pandas/tests/groupby/test_groupby.py
text: @user1:
and then here
```
result = df[['b']] if klass is DataFrame else df['b']
```