Backport PR pandas-dev#40090: REG: DataFrame/Series.transform with list and non-list dict values

Author: rhshadrach

doc/source/whatsnew/v1.2.3.rst

@@ -24,6 +24,8 @@ Fixed regressions
   Passing ``ascending=None`` is still considered invalid,
   and the new error message suggests a proper usage
   (``ascending`` must be a boolean or a list-like boolean).
+- Fixed regression in :meth:`DataFrame.transform` and :meth:`Series.transform` giving incorrect column labels when passed a dictionary with a mix of list and non-list values (:issue:`40018`)
+-
 
 .. ---------------------------------------------------------------------------
 

pandas/core/aggregation.py

@@ -491,6 +491,22 @@ def transform_dict_like(
         # GH 15931 - deprecation of renaming keys
         raise SpecificationError("nested renamer is not supported")
 
+    is_aggregator = lambda x: isinstance(x, (list, tuple, dict))
+
+    # if we have a dict of any non-scalars
+    # eg. {'A' : ['mean']}, normalize all to
+    # be list-likes
+    # Cannot use func.values() because arg may be a Series
+    if any(is_aggregator(x) for _, x in func.items()):
+        new_func: AggFuncTypeDict = {}
+        for k, v in func.items():
+            if not is_aggregator(v):
+                # mypy can't realize v is not a list here
+                new_func[k] = [v]  # type:ignore[list-item]
+            else:
+                new_func[k] = v
+        func = new_func
+
     results: Dict[Label, FrameOrSeriesUnion] = {}
     for name, how in func.items():
         colg = obj._gotitem(name, ndim=1)

pandas/tests/frame/apply/test_frame_transform.py

@@ -99,6 +99,17 @@ def test_transform_dictlike(axis, float_frame, box):
     tm.assert_frame_equal(result, expected)
 
 
+def test_transform_dictlike_mixed():
+    # GH 40018 - mix of lists and non-lists in values of a dictionary
+    df = DataFrame({"a": [1, 2], "b": [1, 4], "c": [1, 4]})
+    result = df.transform({"b": ["sqrt", "abs"], "c": "sqrt"})
+    expected = DataFrame(
+        [[1.0, 1, 1.0], [2.0, 4, 2.0]],
+        columns=MultiIndex([("b", "c"), ("sqrt", "abs")], [(0, 0, 1), (0, 1, 0)]),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
 @pytest.mark.parametrize(
     "ops",
     [

pandas/tests/series/apply/test_series_transform.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from pandas import DataFrame, Series, concat
+from pandas import DataFrame, MultiIndex, Series, concat
 import pandas._testing as tm
 from pandas.core.base import SpecificationError
 from pandas.core.groupby.base import transformation_kernels
@@ -52,6 +52,17 @@ def test_transform_dictlike(string_series, box):
     tm.assert_frame_equal(result, expected)
 
 
+def test_transform_dictlike_mixed():
+    # GH 40018 - mix of lists and non-lists in values of a dictionary
+    df = Series([1, 4])
+    result = df.transform({"b": ["sqrt", "abs"], "c": "sqrt"})
+    expected = DataFrame(
+        [[1.0, 1, 1.0], [2.0, 4, 2.0]],
+        columns=MultiIndex([("b", "c"), ("sqrt", "abs")], [(0, 0, 1), (0, 1, 0)]),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
 def test_transform_wont_agg(string_series):
     # GH 35964
     # we are trying to transform with an aggregator