BUG: DataFrame.groupby(., dropna=True, axis=0) incorrectly throws ShapeError

doc/source/whatsnew/v1.2.0.rst

@@ -221,7 +221,7 @@ Indexing
 Missing
 ^^^^^^^
 
-- Bug in :meth:`SeriesGroupBy.transform` now correctly handles missing values for `dropna=False` (:issue:`35014`)
+- Bug in :meth:`SeriesGroupBy.transform` and :meth:`DataFrameGroupBy.transform` now correctly handle missing values (:issue:`35014` and :issue:`35612`)
 -
 
 MultiIndex

pandas/core/groupby/generic.py

@@ -538,14 +538,20 @@ def _transform_general(
             if isinstance(res, (ABCDataFrame, ABCSeries)):
                 res = res._values
 
-            results.append(klass(res, index=group.index))
+            indexer = self._get_index(name) if self.dropna else group.index
+            results.append(klass(res, index=indexer))
 
         # check for empty "results" to avoid concat ValueError
         if results:
             from pandas.core.reshape.concat import concat
 
             concatenated = concat(results)
-            result = self._set_result_index_ordered(concatenated)
+
+            if not self.dropna:
+                result = self._set_result_index_ordered(concatenated)
+            else:
+                result = concatenated.sort_index()
+                result.index = self._selected_obj.index[result.index.asi8]
         else:
             result = self.obj._constructor(dtype=np.float64)
         # we will only try to coerce the result type if
@@ -557,7 +563,6 @@ def _transform_general(
                 result = maybe_downcast_numeric(result, self._selected_obj.dtype)
 
         result.name = self._selected_obj.name
-        result.index = self._selected_obj.index
         return result
 
     def _transform_fast(self, result) -> Series:
@@ -1392,6 +1397,9 @@ def _wrap_applied_output(self, keys, values, not_indexed_same=False):
     def _transform_general(
         self, func, *args, engine="cython", engine_kwargs=None, **kwargs
     ):
+        """
+        Transform with a non-str `func`.
+        """
         from pandas.core.reshape.concat import concat
 
         applied = []
@@ -1404,7 +1412,9 @@ def _transform_general(
         else:
             fast_path, slow_path = self._define_paths(func, *args, **kwargs)
 
+        has_nan = False
         for name, group in gen:
+            has_nan = has_nan or isna(name)
             object.__setattr__(group, "name", name)
 
             if maybe_use_numba(engine):
@@ -1413,9 +1423,8 @@ def _transform_general(
                 if cache_key not in NUMBA_FUNC_CACHE:
                     NUMBA_FUNC_CACHE[cache_key] = numba_func
                 # Return the result as a DataFrame for concatenation later
-                res = self.obj._constructor(
-                    res, index=group.index, columns=group.columns
-                )
+                indexer = self._get_index(name) if self.dropna else group.index
+                res = self.obj._constructor(res, index=indexer, columns=group.columns)
             else:
                 # Try slow path and fast path.
                 try:
@@ -1449,7 +1458,6 @@ def _transform_general(
                 applied.append(r)
             else:
                 applied.append(res)
-
         concat_index = obj.columns if self.axis == 0 else obj.index
         other_axis = 1 if self.axis == 0 else 0  # switches between 0 & 1
         concatenated = concat(applied, axis=self.axis, verify_integrity=False)
@@ -1677,12 +1685,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

@@ -690,7 +690,10 @@ 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)
+        result_idx, obj_idx = result.index, self.obj._get_axis(self.axis)
+        intersection = result_idx.intersection(obj_idx)
+        indexer = obj_idx if intersection.empty else intersection
+        result.set_axis(indexer, axis=self.axis, inplace=True)
         return result
 
     def _dir_additions(self):

pandas/tests/groupby/test_apply.py

@@ -404,7 +404,7 @@ def trans2(group):
 
 
 def test_apply_transform(ts):
-    grouped = ts.groupby(lambda x: x.month)
+    grouped = ts.groupby(lambda x: x.month, dropna=False)
     result = grouped.apply(lambda x: x * 2)
     expected = grouped.transform(lambda x: x * 2)
     tm.assert_series_equal(result, expected)

pandas/tests/groupby/test_groupby_dropna.py

@@ -163,37 +163,40 @@ def test_groupby_dropna_series_by(dropna, expected):
 
 
 @pytest.mark.parametrize(
-    "dropna,df_expected,s_expected",
+    "dropna,expected",
     [
-        pytest.param(
-            True,
-            pd.DataFrame({"B": [2, 2, 1]}),
-            pd.Series(data=[2, 2, 1], name="B"),
-            marks=pytest.mark.xfail(raises=ValueError),
-        ),
-        (
-            False,
-            pd.DataFrame({"B": [2, 2, 1, 1]}),
-            pd.Series(data=[2, 2, 1, 1], name="B"),
-        ),
+        (True, pd.DataFrame({"B": [2, 2, 1]})),
+        (False, pd.DataFrame({"B": [2, 2, 1, 1]})),
     ],
 )
-def test_slice_groupby_then_transform(dropna, df_expected, s_expected):
-    # GH35014
+def test_groupby_dataframe_slice_then_transform(dropna, expected):
+    # GH35014 & GH35612
 
     df = pd.DataFrame({"A": [0, 0, 1, None], "B": [1, 2, 3, None]})
     gb = df.groupby("A", dropna=dropna)
 
-    res = gb.transform(len)
-    tm.assert_frame_equal(res, df_expected)
+    result = gb.transform(len)
+    tm.assert_frame_equal(result, expected)
+
+    result = gb[["B"]].transform(len)
+    tm.assert_frame_equal(result, expected)
 
-    gb_slice = gb[["B"]]
-    res = gb_slice.transform(len)
-    tm.assert_frame_equal(res, df_expected)
 
-    gb_slice = gb["B"]
-    res = gb["B"].transform(len)
-    tm.assert_series_equal(res, s_expected)
+@pytest.mark.parametrize(
+    "dropna,expected",
+    [
+        (True, pd.Series(data=[2, 2, 1], name="B")),
+        (False, pd.Series(data=[2, 2, 1, 1], name="B")),
+    ],
+)
+def test_groupby_series_slice_then_transform_(dropna, expected):
+    # GH35014 & GH35612
+
+    df = pd.DataFrame({"A": [0, 0, 1, None], "B": [1, 2, 3, None]})
+    gb = df.groupby("A", dropna=dropna)
+
+    result = gb["B"].transform(len)
+    tm.assert_series_equal(result, expected)
 
 
 @pytest.mark.parametrize(

pandas/tests/groupby/test_grouping.py

@@ -615,7 +615,7 @@ def test_list_grouper_with_nat(self):
         [
             (
                 "transform",
-                pd.Series(name=2, dtype=np.float64, index=pd.RangeIndex(0, 0, 1)),
+                pd.Series(name=2, dtype=np.float64, index=pd.Index([], dtype="object")),
             ),
             (
                 "agg",