REF: Move aggregation into apply (#38867)

Author: rhshadrach

pandas/core/apply.py

@@ -1,12 +1,18 @@
 import abc
 import inspect
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Type, cast
 
 import numpy as np
 
 from pandas._config import option_context
 
-from pandas._typing import AggFuncType, Axis, FrameOrSeriesUnion
+from pandas._typing import (
+    AggFuncType,
+    AggFuncTypeBase,
+    AggFuncTypeDict,
+    Axis,
+    FrameOrSeriesUnion,
+)
 from pandas.util._decorators import cache_readonly
 
 from pandas.core.dtypes.common import (
@@ -17,6 +23,7 @@
 )
 from pandas.core.dtypes.generic import ABCSeries
 
+from pandas.core.aggregation import agg_dict_like, agg_list_like
 from pandas.core.construction import create_series_with_explicit_dtype
 
 if TYPE_CHECKING:
@@ -27,6 +34,7 @@
 
 def frame_apply(
     obj: "DataFrame",
+    how: str,
     func: AggFuncType,
     axis: Axis = 0,
     raw: bool = False,
@@ -44,6 +52,7 @@ def frame_apply(
 
     return klass(
         obj,
+        how,
         func,
         raw=raw,
         result_type=result_type,
@@ -84,13 +93,16 @@ def wrap_results_for_axis(
     def __init__(
         self,
         obj: "DataFrame",
+        how: str,
         func,
         raw: bool,
         result_type: Optional[str],
         args,
         kwds,
     ):
+        assert how in ("apply", "agg")
         self.obj = obj
+        self.how = how
         self.raw = raw
         self.args = args or ()
         self.kwds = kwds or {}
@@ -104,15 +116,19 @@ def __init__(
         self.result_type = result_type
 
         # curry if needed
-        if (kwds or args) and not isinstance(func, (np.ufunc, str)):
+        if (
+            (kwds or args)
+            and not isinstance(func, (np.ufunc, str))
+            and not is_list_like(func)
+        ):
 
             def f(x):
                 return func(x, *args, **kwds)
 
         else:
             f = func
 
-        self.f = f
+        self.f: AggFuncType = f
 
     @property
     def res_columns(self) -> "Index":
@@ -139,6 +155,54 @@ def agg_axis(self) -> "Index":
         return self.obj._get_agg_axis(self.axis)
 
     def get_result(self):
+        if self.how == "apply":
+            return self.apply()
+        else:
+            return self.agg()
+
+    def agg(self) -> Tuple[Optional[FrameOrSeriesUnion], Optional[bool]]:
+        """
+        Provide an implementation for the aggregators.
+
+        Returns
+        -------
+        tuple of result, how.
+
+        Notes
+        -----
+        how can be a string describe the required post-processing, or
+        None if not required.
+        """
+        obj = self.obj
+        arg = self.f
+        args = self.args
+        kwargs = self.kwds
+
+        _axis = kwargs.pop("_axis", None)
+        if _axis is None:
+            _axis = getattr(obj, "axis", 0)
+
+        if isinstance(arg, str):
+            return obj._try_aggregate_string_function(arg, *args, **kwargs), None
+        elif is_dict_like(arg):
+            arg = cast(AggFuncTypeDict, arg)
+            return agg_dict_like(obj, arg, _axis), True
+        elif is_list_like(arg):
+            # we require a list, but not a 'str'
+            arg = cast(List[AggFuncTypeBase], arg)
+            return agg_list_like(obj, arg, _axis=_axis), None
+        else:
+            result = None
+
+        if callable(arg):
+            f = obj._get_cython_func(arg)
+            if f and not args and not kwargs:
+                return getattr(obj, f)(), None
+
+        # caller can react
+        return result, True
+
+    def apply(self) -> FrameOrSeriesUnion:
         """ compute the results """
         # dispatch to agg
         if is_list_like(self.f) or is_dict_like(self.f):
@@ -191,6 +255,8 @@ def apply_empty_result(self):
         we will try to apply the function to an empty
         series in order to see if this is a reduction function
         """
+        assert callable(self.f)
+
         # we are not asked to reduce or infer reduction
         # so just return a copy of the existing object
         if self.result_type not in ["reduce", None]:
@@ -246,6 +312,8 @@ def wrapper(*args, **kwargs):
             return self.obj._constructor_sliced(result, index=self.agg_axis)
 
     def apply_broadcast(self, target: "DataFrame") -> "DataFrame":
+        assert callable(self.f)
+
         result_values = np.empty_like(target.values)
 
         # axis which we want to compare compliance
@@ -279,6 +347,8 @@ def apply_standard(self):
         return self.wrap_results(results, res_index)
 
     def apply_series_generator(self) -> Tuple[ResType, "Index"]:
+        assert callable(self.f)
+
         series_gen = self.series_generator
         res_index = self.result_index
 

pandas/core/frame.py

@@ -121,12 +121,7 @@
 
 from pandas.core import algorithms, common as com, generic, nanops, ops
 from pandas.core.accessor import CachedAccessor
-from pandas.core.aggregation import (
-    aggregate,
-    reconstruct_func,
-    relabel_result,
-    transform,
-)
+from pandas.core.aggregation import reconstruct_func, relabel_result, transform
 from pandas.core.arraylike import OpsMixin
 from pandas.core.arrays import ExtensionArray
 from pandas.core.arrays.sparse import SparseFrameAccessor
@@ -7623,13 +7618,24 @@ def aggregate(self, func=None, axis: Axis = 0, *args, **kwargs):
         return result
 
     def _aggregate(self, arg, axis: Axis = 0, *args, **kwargs):
+        from pandas.core.apply import frame_apply
+
+        op = frame_apply(
+            self if axis == 0 else self.T,
+            how="agg",
+            func=arg,
+            axis=0,
+            args=args,
+            kwds=kwargs,
+        )
+        result, how = op.get_result()
+
         if axis == 1:
             # NDFrame.aggregate returns a tuple, and we need to transpose
             # only result
-            result, how = aggregate(self.T, arg, *args, **kwargs)
             result = result.T if result is not None else result
-            return result, how
-        return aggregate(self, arg, *args, **kwargs)
+
+        return result, how
 
     agg = aggregate
 
@@ -7789,6 +7795,7 @@ def apply(
 
         op = frame_apply(
             self,
+            how="apply",
             func=func,
             axis=axis,
             raw=raw,