CLN: aggregation.transform

doc/source/whatsnew/v1.2.0.rst

@@ -349,6 +349,7 @@ Reshaping
 - Bug in :func:`union_indexes` where input index names are not preserved in some cases. Affects :func:`concat` and :class:`DataFrame` constructor (:issue:`13475`)
 - Bug in func :meth:`crosstab` when using multiple columns with ``margins=True`` and ``normalize=True`` (:issue:`35144`)
 - Bug in :meth:`DataFrame.agg` with ``func={'name':<FUNC>}`` incorrectly raising ``TypeError`` when ``DataFrame.columns==['Name']`` (:issue:`36212`)
+- Bug in :meth:`Series.transform` would give incorrect results or raise when the argument ``func`` was dictionary (:issue:`35811`)
 -
 
 Sparse
@@ -368,7 +369,6 @@ Other
 ^^^^^
 - Bug in :meth:`DataFrame.replace` and :meth:`Series.replace` incorrectly raising ``AssertionError`` instead of ``ValueError`` when invalid parameter combinations are passed (:issue:`36045`)
 - Bug in :meth:`DataFrame.replace` and :meth:`Series.replace` with numeric values and string ``to_replace`` (:issue:`34789`)
-- Bug in :meth:`Series.transform` would give incorrect results or raise when the argument ``func`` was dictionary (:issue:`35811`)
 - Bug in :meth:`Index.union` behaving differently depending on whether operand is a :class:`Index` or other list-like (:issue:`36384`)
 
 .. ---------------------------------------------------------------------------

pandas/core/aggregation.py

@@ -16,9 +16,17 @@
     Sequence,
     Tuple,
     Union,
+    cast,
 )
 
-from pandas._typing import AggFuncType, Axis, FrameOrSeries, Label
+from pandas._typing import (
+    AggFuncType,
+    AggFuncTypeBase,
+    Axis,
+    FrameOrSeries,
+    FrameOrSeriesUnion,
+    Label,
+)
 
 from pandas.core.dtypes.common import is_dict_like, is_list_like
 from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries
@@ -388,7 +396,7 @@ def validate_func_kwargs(
 
 def transform(
     obj: FrameOrSeries, func: AggFuncType, axis: Axis, *args, **kwargs,
-) -> FrameOrSeries:
+) -> FrameOrSeriesUnion:
     """
     Transform a DataFrame or Series
 
@@ -415,67 +423,94 @@ def transform(
     ValueError
         If the transform function fails or does not transform.
     """
-    from pandas.core.reshape.concat import concat
-
     is_series = obj.ndim == 1
-
     if obj._get_axis_number(axis) == 1:
         assert not is_series
         return transform(obj.T, func, 0, *args, **kwargs).T
 
-    if isinstance(func, list):
+    if is_list_like(func) and not is_dict_like(func):
+        func = cast(List[AggFuncTypeBase], func)
+        # Convert func equivalent dict
         if is_series:
             func = {com.get_callable_name(v) or v: v for v in func}
         else:
             func = {col: func for col in obj}
 
-    if isinstance(func, dict):
-        if not is_series:
-            cols = sorted(set(func.keys()) - set(obj.columns))
-            if len(cols) > 0:
-                raise SpecificationError(f"Column(s) {cols} do not exist")
-
-        if any(isinstance(v, dict) for v in func.values()):
-            # GH 15931 - deprecation of renaming keys
-            raise SpecificationError("nested renamer is not supported")
-
-        results = {}
-        for name, how in func.items():
-            colg = obj._gotitem(name, ndim=1)
-            try:
-                results[name] = transform(colg, how, 0, *args, **kwargs)
-            except Exception as e:
-                if str(e) == "Function did not transform":
-                    raise e
-
-        # combine results
-        if len(results) == 0:
+    if is_dict_like(func):
+        func = cast(Dict[Label, Union[AggFuncTypeBase, List[AggFuncTypeBase]]], func)
+        result = transform_dict_like(obj, func, *args, **kwargs)
+    else:
+        func = cast(AggFuncTypeBase, func)
+        try:
+            result = transform_str_or_callable(obj, func, *args, **kwargs)
+        except Exception:
             raise ValueError("Transform function failed")
-        return concat(results, axis=1)
-
-    # func is either str or callable
-    try:
-        if isinstance(func, str):
-            result = obj._try_aggregate_string_function(func, *args, **kwargs)
-        else:
-            f = obj._get_cython_func(func)
-            if f and not args and not kwargs:
-                result = getattr(obj, f)()
-            else:
-                try:
-                    result = obj.apply(func, args=args, **kwargs)
-                except Exception:
-                    result = func(obj, *args, **kwargs)
-    except Exception:
-        raise ValueError("Transform function failed")
 
     # Functions that transform may return empty Series/DataFrame
     # when the dtype is not appropriate
     if isinstance(result, (ABCSeries, ABCDataFrame)) and result.empty:
         raise ValueError("Transform function failed")
+
     if not isinstance(result, (ABCSeries, ABCDataFrame)) or not result.index.equals(
         obj.index
     ):
         raise ValueError("Function did not transform")
 
     return result
+
+
+def transform_dict_like(
+    obj: FrameOrSeries,
+    func: Dict[Label, Union[AggFuncTypeBase, List[AggFuncTypeBase]]],
+    *args,
+    **kwargs,
+):
+    """
+    Compute transform in the case of a dict-like func
+    """
+    from pandas.core.reshape.concat import concat
+
+    if obj.ndim != 1:
+        # Check for missing columns on a frame
+        cols = sorted(set(func.keys()) - set(obj.columns))
+        if len(cols) > 0:
+            raise SpecificationError(f"Column(s) {cols} do not exist")
+
+    if any(isinstance(v, dict) for v in func.values()):
+        # GH 15931 - deprecation of renaming keys
+        raise SpecificationError("nested renamer is not supported")
+
+    results: Dict[Label, FrameOrSeriesUnion] = {}
+    for name, how in func.items():
+        colg = obj._gotitem(name, ndim=1)
+        try:
+            results[name] = transform(colg, how, 0, *args, **kwargs)
+        except Exception as e:
+            if str(e) == "Function did not transform":
+                raise e
+
+    # combine results
+    if len(results) == 0:
+        raise ValueError("Transform function failed")
+    return concat(results, axis=1)
+
+
+def transform_str_or_callable(
+    obj: FrameOrSeries, func: AggFuncTypeBase, *args, **kwargs
+) -> FrameOrSeriesUnion:
+    """
+    Compute transform in the case of a string or callable func
+    """
+    if isinstance(func, str):
+        return obj._try_aggregate_string_function(func, *args, **kwargs)
+
+    if not args and not kwargs:
+        f = obj._get_cython_func(func)
+        if f:
+            return getattr(obj, f)()
+
+    # Two possible ways to use a UDF - apply or call directly
+    try:
+        return obj.apply(func, args=args, **kwargs)
+    except Exception:
+        return func(obj, *args, **kwargs)

pandas/core/frame.py

@@ -7255,7 +7255,7 @@ def diff(self, periods: int = 1, axis: Axis = 0) -> DataFrame:
 
     def _gotitem(
         self,
-        key: Union[str, List[str]],
+        key: Union[Label, List[Label]],
         ndim: int,
         subset: Optional[FrameOrSeriesUnion] = None,
     ) -> FrameOrSeriesUnion:

pandas/core/shared_docs.py

@@ -265,16 +265,17 @@
 
 Parameters
 ----------
-func : function, str, list or dict
+func : function, str, list-like or dict-like
     Function to use for transforming the data. If a function, must either
-    work when passed a {klass} or when passed to {klass}.apply.
+    work when passed a {klass} or when passed to {klass}.apply. If func
+    is both list-like and dict-like, dict-like behavior takes precedence.
 
     Accepted combinations are:
 
     - function
     - string function name
-    - list of functions and/or function names, e.g. ``[np.exp, 'sqrt']``
-    - dict of axis labels -> functions, function names or list of such.
+    - list-like of functions and/or function names, e.g. ``[np.exp, 'sqrt']``
+    - dict-like of axis labels -> functions, function names or list-like of such.
 {axis}
 *args
     Positional arguments to pass to `func`.