Skip to content

Commit

Permalink
BUG: DataFrame.apply with axis=1 and EA dtype (#38272)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Dec 4, 2020
1 parent 4699c2b commit b5233c4
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 7 deletions.
27 changes: 20 additions & 7 deletions pandas/core/apply.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -392,12 +397,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":
Expand Down
6 changes: 6 additions & 0 deletions pandas/tests/frame/apply/test_frame_apply.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit b5233c4

Please sign in to comment.