Skip to content

REF: share _wrap_aggregated_output #43446

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 7, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 15 additions & 60 deletions pandas/core/groupby/generic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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]
Expand Down
47 changes: 46 additions & 1 deletion pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down