From ea71de1f7b3966655d17f7c28ea5474bc6db2c6c Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 3 Dec 2020 10:49:54 -0800 Subject: [PATCH 1/2] BUG: DataFrame.apply with axis=1 and EA dtype --- pandas/core/apply.py | 13 ++++++++++++- pandas/tests/frame/apply/test_frame_apply.py | 6 ++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/pandas/core/apply.py b/pandas/core/apply.py index c5260deafc0c3..80ebc861cbbfe 100644 --- a/pandas/core/apply.py +++ b/pandas/core/apply.py @@ -9,7 +9,12 @@ from pandas._typing import Axis, FrameOrSeriesUnion 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 @@ -392,6 +397,12 @@ def series_generator(self): mgr = ser._mgr blk = mgr.blocks[0] + if is_extension_array_dtype(blk.dtype): + # values will be incorrect for this block + vals = blk.values + # TODO(EA2D): iterator would be unnecessary with 2D EAs + values = (vals[slice(i, i + 1)] for i in range(len(vals))) + for (arr, name) in zip(values, self.index): # GH#35462 re-pin mgr in case setitem changed it ser._mgr = mgr diff --git a/pandas/tests/frame/apply/test_frame_apply.py b/pandas/tests/frame/apply/test_frame_apply.py index 15952f36b0fae..9ec56c3429b22 100644 --- a/pandas/tests/frame/apply/test_frame_apply.py +++ b/pandas/tests/frame/apply/test_frame_apply.py @@ -58,6 +58,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 From 4db8fbdda6f1c1df647abcfcfdf220e28b369d5c Mon Sep 17 00:00:00 2001 From: Brock Date: Thu, 3 Dec 2020 19:26:10 -0800 Subject: [PATCH 2/2] more robust solution --- pandas/core/apply.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/pandas/core/apply.py b/pandas/core/apply.py index 80ebc861cbbfe..6d9e11ecb824f 100644 --- a/pandas/core/apply.py +++ b/pandas/core/apply.py @@ -399,16 +399,18 @@ def series_generator(self): if is_extension_array_dtype(blk.dtype): # values will be incorrect for this block - vals = blk.values - # TODO(EA2D): iterator would be unnecessary with 2D EAs - values = (vals[slice(i, i + 1)] for i in range(len(vals))) - - 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 + # 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":