diff --git a/pandas/core/aggregation.py b/pandas/core/aggregation.py index e2374b81ca13b..7ca68d8289bd5 100644 --- a/pandas/core/aggregation.py +++ b/pandas/core/aggregation.py @@ -10,6 +10,7 @@ Callable, DefaultDict, Dict, + Iterable, List, Optional, Sequence, @@ -17,14 +18,14 @@ Union, ) -from pandas._typing import AggFuncType, Label +from pandas._typing import AggFuncType, FrameOrSeries, Label from pandas.core.dtypes.common import is_dict_like, is_list_like from pandas.core.base import SpecificationError import pandas.core.common as com from pandas.core.indexes.api import Index -from pandas.core.series import FrameOrSeriesUnion, Series +from pandas.core.series import Series def reconstruct_func( @@ -276,12 +277,13 @@ def maybe_mangle_lambdas(agg_spec: Any) -> Any: def relabel_result( - result: FrameOrSeriesUnion, + result: FrameOrSeries, func: Dict[str, List[Union[Callable, str]]], - columns: Tuple, - order: List[int], + columns: Iterable[Label], + order: Iterable[int], ) -> Dict[Label, Series]: - """Internal function to reorder result if relabelling is True for + """ + Internal function to reorder result if relabelling is True for dataframe.agg, and return the reordered result in dict. Parameters: diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 606bd4cc3b52d..fe6fb97012fac 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -7415,6 +7415,12 @@ def aggregate(self, func=None, axis=0, *args, **kwargs): if relabeling: # This is to keep the order to columns occurrence unchanged, and also # keep the order of new columns occurrence unchanged + + # For the return values of reconstruct_func, if relabeling is + # False, columns and order will be None. + assert columns is not None + assert order is not None + result_in_dict = relabel_result(result, func, columns, order) result = DataFrame(result_in_dict, index=columns)