move DataFrameGroupBy._transform_general logic to _set_result_index_ordered

Author: arw2019

pandas/core/groupby/generic.py

@@ -1397,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 = []
@@ -1455,17 +1458,11 @@ 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)
         concatenated = concatenated.reindex(concat_index, axis=other_axis, copy=False)
-        if not self.dropna or not has_nan:
-            return self._set_result_index_ordered(concatenated)
-        else:
-            concatenated.sort_index(inplace=True)
-            concatenated.index = obj.index[concatenated.index.asi8]
-            return concatenated
+        return self._set_result_index_ordered(concatenated)
 
     @Substitution(klass="DataFrame")
     @Appender(_transform_template)

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):