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