From d3428c13c27f0cc551bd92ce2dd51f9f481a34f2 Mon Sep 17 00:00:00 2001
From: jbrockmendel <jbrockmendel@gmail.com>
Date: Fri, 4 Dec 2020 08:21:23 -0800
Subject: [PATCH] Backport PR #38272: BUG: DataFrame.apply with axis=1 and EA
 dtype

---
 pandas/core/apply.py                         | 27 +++++++++++++++-----
 pandas/tests/frame/apply/test_frame_apply.py |  6 +++++
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/pandas/core/apply.py b/pandas/core/apply.py
index fd7ffd1b54a70..af4717498aa19 100644
--- a/pandas/core/apply.py
+++ b/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":
diff --git a/pandas/tests/frame/apply/test_frame_apply.py b/pandas/tests/frame/apply/test_frame_apply.py
index a89a20fc69ef8..73a80d048ba00 100644
--- a/pandas/tests/frame/apply/test_frame_apply.py
+++ b/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