diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index 35cb247e96bc3..8bb8f00b4c406 100644 --- a/pandas/core/groupby/generic.py +++ b/pandas/core/groupby/generic.py @@ -354,35 +354,17 @@ def array_func(values: ArrayLike) -> ArrayLike: ) return self._reindex_output(ser) - def _wrap_aggregated_output( - self, - output: Mapping[base.OutputKey, Series | ArrayLike], + def _indexed_output_to_ndframe( + self, output: Mapping[base.OutputKey, ArrayLike] ) -> Series: """ - Wraps the output of a SeriesGroupBy aggregation into the expected result. - - Parameters - ---------- - output : Mapping[base.OutputKey, Union[Series, ArrayLike]] - Data to wrap. - - Returns - ------- - Series - - Notes - ----- - In the vast majority of cases output will only contain one element. - The exception is operations that expand dimensions, like ohlc. + Wrap the dict result of a GroupBy aggregation into a Series. """ assert len(output) == 1 - - name = self.obj.name - index = self.grouper.result_index values = next(iter(output.values())) - - result = self.obj._constructor(values, index=index, name=name) - return self._reindex_output(result) + result = self.obj._constructor(values) + result.name = self.obj.name + return result def _wrap_transformed_output( self, output: Mapping[base.OutputKey, Series | ArrayLike] @@ -1614,46 +1596,19 @@ def _insert_inaxis_grouper_inplace(self, result: DataFrame) -> None: if in_axis and name not in columns: result.insert(0, name, lev) - def _wrap_aggregated_output( - self, - output: Mapping[base.OutputKey, Series | ArrayLike], + def _indexed_output_to_ndframe( + self, output: Mapping[base.OutputKey, ArrayLike] ) -> DataFrame: """ - Wraps the output of DataFrameGroupBy aggregations into the expected result. - - Parameters - ---------- - output : Mapping[base.OutputKey, Union[Series, np.ndarray]] - Data to wrap. - - Returns - ------- - DataFrame + Wrap the dict result of a GroupBy aggregation into a DataFrame. """ - if isinstance(output, DataFrame): - result = output - else: - indexed_output = {key.position: val for key, val in output.items()} - columns = Index([key.label for key in output]) - columns._set_names(self._obj_with_exclusions._get_axis(1 - self.axis).names) - - result = self.obj._constructor(indexed_output) - result.columns = columns - - if not self.as_index: - self._insert_inaxis_grouper_inplace(result) - result = result._consolidate() - else: - result.index = self.grouper.result_index - - if self.axis == 1: - result = result.T - if result.index.equals(self.obj.index): - # Retain e.g. DatetimeIndex/TimedeltaIndex freq - result.index = self.obj.index.copy() - # TODO: Do this more systematically + indexed_output = {key.position: val for key, val in output.items()} + columns = Index([key.label for key in output]) + columns._set_names(self._obj_with_exclusions._get_axis(1 - self.axis).names) - return self._reindex_output(result) + result = self.obj._constructor(indexed_output) + result.columns = columns + return result def _wrap_transformed_output( self, output: Mapping[base.OutputKey, Series | ArrayLike] diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py index 0e358e611f418..a60ec29581337 100644 --- a/pandas/core/groupby/groupby.py +++ b/pandas/core/groupby/groupby.py @@ -1095,9 +1095,54 @@ def _set_result_index_ordered( return result - def _wrap_aggregated_output(self, output: Mapping[base.OutputKey, ArrayLike]): + def _indexed_output_to_ndframe( + self, result: Mapping[base.OutputKey, ArrayLike] + ) -> Series | DataFrame: raise AbstractMethodError(self) + def _wrap_aggregated_output( + self, output: Series | DataFrame | Mapping[base.OutputKey, ArrayLike] + ): + """ + Wraps the output of GroupBy aggregations into the expected result. + + Parameters + ---------- + output : Series, DataFrame, or Mapping[base.OutputKey, ArrayLike] + Data to wrap. + + Returns + ------- + Series or DataFrame + """ + + if isinstance(output, (Series, DataFrame)): + # We get here (for DataFrameGroupBy) if we used Manager.grouped_reduce, + # in which case our columns are already set correctly. + # ATM we do not get here for SeriesGroupBy; when we do, we will + # need to require that result.name already match self.obj.name + result = output + else: + result = self._indexed_output_to_ndframe(output) + + if not self.as_index: + # `not self.as_index` is only relevant for DataFrameGroupBy, + # enforced in __init__ + self._insert_inaxis_grouper_inplace(result) + result = result._consolidate() + else: + result.index = self.grouper.result_index + + if self.axis == 1: + # Only relevant for DataFrameGroupBy, no-op for SeriesGroupBy + result = result.T + if result.index.equals(self.obj.index): + # Retain e.g. DatetimeIndex/TimedeltaIndex freq + result.index = self.obj.index.copy() + # TODO: Do this more systematically + + return self._reindex_output(result) + def _wrap_transformed_output(self, output: Mapping[base.OutputKey, ArrayLike]): raise AbstractMethodError(self)