Skip to content

CLN: remove index keyword from _wrap_aggregated_output #41128

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.

Sign up for GitHub

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 6 commits into from
Apr 26, 2021
Merged

CLN: remove index keyword from _wrap_aggregated_output #41128

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
81897cd
REF: remove index kwarg from _wrap_aggregated_output
jbrockmendel Apr 23, 2021
6328016
REF: remove _wrap_series_output
jbrockmendel Apr 23, 2021
f3b1baa
REF: remove unnecessary layer
jbrockmendel Apr 23, 2021
90b662a
revert incorrect comment
jbrockmendel Apr 23, 2021
941763b
PERF: next(iter(...))
jbrockmendel Apr 24, 2021
1ab60ba
next(iter(...))
jbrockmendel Apr 25, 2021
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
98 changes: 23 additions & 75 deletions pandas/core/groupby/generic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -336,14 +336,12 @@ def _aggregate_multiple_funcs(self, arg):
# let higher level handle
return results

# Argument 1 to "_wrap_aggregated_output" of "SeriesGroupBy" has
# incompatible type "Dict[OutputKey, Union[DataFrame,
# Series]]";
# expected "Mapping[OutputKey, Union[Series, ndarray]]"
output = self._wrap_aggregated_output(
results, index=None # type: ignore[arg-type]
)
return self.obj._constructor_expanddim(output, columns=columns)
indexed_output = {key.position: val for key, val in results.items()}
output = self.obj._constructor_expanddim(indexed_output, index=None)
output.columns = Index(key.label for key in results)

output = self._reindex_output(output)
return output

def _cython_agg_general(
self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1
Expand Down Expand Up @@ -371,78 +369,36 @@ def _cython_agg_general(
if not output:
raise DataError("No numeric types to aggregate")

# error: Argument 1 to "_wrap_aggregated_output" of "BaseGroupBy" has
# incompatible type "Dict[OutputKey, Union[ndarray, DatetimeArray]]";
# expected "Mapping[OutputKey, ndarray]"
return self._wrap_aggregated_output(
output, index=self.grouper.result_index # type: ignore[arg-type]
)

# TODO: index should not be Optional - see GH 35490
def _wrap_series_output(
self,
output: Mapping[base.OutputKey, Series | ArrayLike],
index: Index | None,
) -> FrameOrSeriesUnion:
"""
Wraps the output of a SeriesGroupBy operation into the expected result.

Parameters
----------
output : Mapping[base.OutputKey, Union[Series, np.ndarray, ExtensionArray]]
Data to wrap.
index : pd.Index or None
Index to apply to the output.

Returns
-------
Series or DataFrame

Notes
-----
In the vast majority of cases output and columns will only contain one
element. The exception is operations that expand dimensions, like ohlc.
"""
indexed_output = {key.position: val for key, val in output.items()}
columns = Index(key.label for key in output)
return self._wrap_aggregated_output(output)

result: FrameOrSeriesUnion
if len(output) > 1:
result = self.obj._constructor_expanddim(indexed_output, index=index)
result.columns = columns
elif not columns.empty:
result = self.obj._constructor(
indexed_output[0], index=index, name=columns[0]
)
else:
result = self.obj._constructor_expanddim()

return result

# TODO: Remove index argument, use self.grouper.result_index, see GH 35490
def _wrap_aggregated_output(
self,
output: Mapping[base.OutputKey, Series | np.ndarray],
index: Index | None,
) -> FrameOrSeriesUnion:
output: Mapping[base.OutputKey, Series | ArrayLike],
) -> Series:
"""
Wraps the output of a SeriesGroupBy aggregation into the expected result.

Parameters
----------
output : Mapping[base.OutputKey, Union[Series, np.ndarray]]
output : Mapping[base.OutputKey, Union[Series, ArrayLike]]
Data to wrap.

Returns
-------
Series or DataFrame
Series

Notes
-----
In the vast majority of cases output will only contain one element.
The exception is operations that expand dimensions, like ohlc.
"""
result = self._wrap_series_output(output=output, index=index)
assert len(output) == 1

name = self.obj.name
index = self.grouper.result_index
values = list(output.values())[0]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tiny perf boost by not creating the intermediate list

d = {'a': 1}
%timeit list(d.values())[0]
%timeit next(iter(d.values()))

gives

219 ns ± 0.363 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
146 ns ± 22.8 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good idea, updated


result = self.obj._constructor(values, index=index, name=name)
return self._reindex_output(result)

def _wrap_transformed_output(
Expand All @@ -466,7 +422,10 @@ def _wrap_transformed_output(
for consistency with DataFrame methods and _wrap_aggregated_output.
"""
assert len(output) == 1
result = self._wrap_series_output(output=output, index=self.obj.index)

name = self.obj.name
values = list(output.values())[0]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like the next trick can work here too.

result = self.obj._constructor(values, index=self.obj.index, name=name)

# No transformations increase the ndim of the result
assert isinstance(result, Series)
Expand Down Expand Up @@ -1115,14 +1074,6 @@ def _iterate_slices(self) -> Iterable[Series]:
def _cython_agg_general(
self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1
) -> DataFrame:
agg_mgr = self._cython_agg_manager(
how, alt=alt, numeric_only=numeric_only, min_count=min_count
)
return self._wrap_agged_manager(agg_mgr)

def _cython_agg_manager(
self, how: str, alt=None, numeric_only: bool = True, min_count: int = -1
) -> Manager2D:
# Note: we never get here with how="ohlc"; that goes through SeriesGroupBy

data: Manager2D = self._get_data_to_aggregate()
Expand Down Expand Up @@ -1186,11 +1137,9 @@ def py_fallback(values: ArrayLike) -> ArrayLike:
sgb = get_groupby(obj, self.grouper, observed=True)
result = sgb.aggregate(lambda x: alt(x, axis=self.axis))

assert isinstance(result, (Series, DataFrame)) # for mypy
# In the case of object dtype block, it may have been split
# in the operation. We un-split here.
result = result._consolidate()
assert isinstance(result, (Series, DataFrame)) # for mypy
# unwrap DataFrame/Series to get array
mgr = result._mgr
arrays = mgr.arrays
Expand Down Expand Up @@ -1226,7 +1175,7 @@ def array_func(values: ArrayLike) -> ArrayLike:
if not len(new_mgr):
raise DataError("No numeric types to aggregate")

return new_mgr
return self._wrap_agged_manager(new_mgr)

def _aggregate_frame(self, func, *args, **kwargs) -> DataFrame:
if self.grouper.nkeys != 1:
Expand Down Expand Up @@ -1733,7 +1682,6 @@ def _insert_inaxis_grouper_inplace(self, result: DataFrame) -> None:
def _wrap_aggregated_output(
self,
output: Mapping[base.OutputKey, Series | np.ndarray],
index: Index | None,
) -> DataFrame:
"""
Wraps the output of DataFrameGroupBy aggregations into the expected result.
Expand Down
8 changes: 3 additions & 5 deletions pandas/core/groupby/groupby.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1056,9 +1056,7 @@ def _set_result_index_ordered(

return result

def _wrap_aggregated_output(
self, output: Mapping[base.OutputKey, np.ndarray], index: Index | None
):
def _wrap_aggregated_output(self, output: Mapping[base.OutputKey, np.ndarray]):
raise AbstractMethodError(self)

def _wrap_transformed_output(self, output: Mapping[base.OutputKey, ArrayLike]):
Expand Down Expand Up @@ -1259,7 +1257,7 @@ def _python_agg_general(self, func, *args, **kwargs):
if not output:
return self._python_apply_general(f, self._selected_obj)

return self._wrap_aggregated_output(output, index=self.grouper.result_index)
return self._wrap_aggregated_output(output)

@final
def _agg_general(
Expand Down Expand Up @@ -2786,7 +2784,7 @@ def _get_cythonized_result(
raise TypeError(error_msg)

if aggregate:
return self._wrap_aggregated_output(output, index=self.grouper.result_index)
return self._wrap_aggregated_output(output)
else:
return self._wrap_transformed_output(output)

Expand Down