move DataFrameGroupBy._transform_general logic to _set_result_index_ordered

Author: arw2019

pandas/core/groupby/generic.py

@@ -517,6 +517,7 @@ def _transform_general(
         """
         Transform with a non-str `func`.
         """
+        print("Calling SeriesGroupBy._transform_general")
         if maybe_use_numba(engine):
             numba_func, cache_key = generate_numba_func(
                 func, engine_kwargs, kwargs, "groupby_transform"
@@ -1455,17 +1456,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

@@ -682,15 +682,17 @@ def _set_group_selection(self):
     def _set_result_index_ordered(self, result):
         # set the result index on the passed values object and
         # return the new object, xref 8046
-
         # the values/counts are repeated according to the group index
         # shortcut if we have an already ordered grouper
         if not self.grouper.is_monotonic:
             index = Index(np.concatenate(self._get_indices(self.grouper.result_index)))
             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):