handle dropna=False in _selected_obj, _set_result_index_ordered

Author: arw2019

pandas/core/groupby/generic.py

@@ -1677,12 +1677,17 @@ def _gotitem(self, key, ndim: int, subset=None):
                 exclusions=self.exclusions,
                 as_index=self.as_index,
                 observed=self.observed,
+                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,
+                selection=key,
+                grouper=self.grouper,
+                observed=self.observed,
+                dropna=self.dropna,
             )
 
         raise AssertionError("invalid ndim for _gotitem")

pandas/core/groupby/groupby.py

@@ -637,10 +637,12 @@ def _selected_obj(self):
 
         if self._selection is None or isinstance(self.obj, Series):
             if self._group_selection is not None:
-                return self.obj[self._group_selection]
-            return self.obj
+                result = self.obj[self._group_selection]
+            result = self.obj
         else:
-            return self.obj[self._selection]
+            result = self.obj[self._selection]
+
+        return result.dropna() if self.dropna else result
 
     def _reset_group_selection(self):
         """
@@ -690,7 +692,12 @@ def _set_result_index_ordered(self, result):
             result.set_axis(index, axis=self.axis, inplace=True)
             result = result.sort_index(axis=self.axis)
 
-        result.set_axis(self.obj._get_axis(self.axis), axis=self.axis, inplace=True)
+        if hasattr(self, "_selected_obj"):
+            labels = self._selected_obj._get_axis(self.axis)
+        else:
+            labels = self.obj._get_axis(self.axis)
+
+        result.set_axis(labels, axis=self.axis, inplace=True)
         return result
 
     def _dir_additions(self):

pandas/tests/groupby/test_groupby_dropna.py

@@ -165,12 +165,7 @@ def test_groupby_dropna_series_by(dropna, expected):
 @pytest.mark.parametrize(
     "dropna,df_expected,s_expected",
     [
-        pytest.param(
-            True,
-            pd.DataFrame({"B": [2, 2, 1]}),
-            pd.Series(data=[2, 2, 1], name="B"),
-            marks=pytest.mark.xfail(raises=ValueError),
-        ),
+        (True, pd.DataFrame({"B": [2, 2, 1]}), pd.Series(data=[2, 2, 1], name="B"),),
         (
             False,
             pd.DataFrame({"B": [2, 2, 1, 1]}),
@@ -179,7 +174,7 @@ def test_groupby_dropna_series_by(dropna, expected):
     ],
 )
 def test_slice_groupby_then_transform(dropna, df_expected, s_expected):
-    # GH35014
+    # GH35014 & GH35612
 
     df = pd.DataFrame({"A": [0, 0, 1, None], "B": [1, 2, 3, None]})
     gb = df.groupby("A", dropna=dropna)