Backport PR pandas-dev#38272: BUG: DataFrame.apply with axis=1 and EA dtype

Author: jbrockmendel

pandas/core/apply.py

@@ -9,7 +9,12 @@
 from pandas._typing import Axis
 from pandas.util._decorators import cache_readonly
 
-from pandas.core.dtypes.common import is_dict_like, is_list_like, is_sequence
+from pandas.core.dtypes.common import (
+    is_dict_like,
+    is_extension_array_dtype,
+    is_list_like,
+    is_sequence,
+)
 from pandas.core.dtypes.generic import ABCSeries
 
 from pandas.core.construction import create_series_with_explicit_dtype
@@ -407,12 +412,20 @@ def series_generator(self):
         mgr = ser._mgr
         blk = mgr.blocks[0]
 
-        for (arr, name) in zip(values, self.index):
-            # GH#35462 re-pin mgr in case setitem changed it
-            ser._mgr = mgr
-            blk.values = arr
-            ser.name = name
-            yield ser
+        if is_extension_array_dtype(blk.dtype):
+            # values will be incorrect for this block
+            # TODO(EA2D): special case would be unnecessary with 2D EAs
+            obj = self.obj
+            for i in range(len(obj)):
+                yield obj._ixs(i, axis=0)
+
+        else:
+            for (arr, name) in zip(values, self.index):
+                # GH#35462 re-pin mgr in case setitem changed it
+                ser._mgr = mgr
+                blk.values = arr
+                ser.name = name
+                yield ser
 
     @property
     def result_index(self) -> "Index":

pandas/tests/frame/apply/test_frame_apply.py

@@ -60,6 +60,12 @@ def test_apply(self, float_frame):
         assert isinstance(df["c0"].dtype, CategoricalDtype)
         assert isinstance(df["c1"].dtype, CategoricalDtype)
 
+    def test_apply_axis1_with_ea(self):
+        # GH#36785
+        df = DataFrame({"A": [Timestamp("2013-01-01", tz="UTC")]})
+        result = df.apply(lambda x: x, axis=1)
+        tm.assert_frame_equal(result, df)
+
     def test_apply_mixed_datetimelike(self):
         # mixed datetimelike
         # GH 7778