Skip to content

Fix #60766:.map,.apply would convert element type for extension array #61396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions doc/source/whatsnew/v3.0.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -815,6 +815,8 @@ ExtensionArray
- Bug in comparison between object with :class:`ArrowDtype` and incompatible-dtyped (e.g. string vs bool) incorrectly raising instead of returning all-``False`` (for ``==``) or all-``True`` (for ``!=``) (:issue:`59505`)
- Bug in constructing pandas data structures when passing into ``dtype`` a string of the type followed by ``[pyarrow]`` while PyArrow is not installed would raise ``NameError`` rather than ``ImportError`` (:issue:`57928`)
- Bug in various :class:`DataFrame` reductions for pyarrow temporal dtypes returning incorrect dtype when result was null (:issue:`59234`)
- Bug in :meth:`Series.map` and :meth:`Series.apply` where applying functions to a Series with an :class:`Int32Dtype` or other :class:`ExtensionDtype` would convert elements to float and ``pd.NA`` to ``np.nan``, instead of preserving the original types (:issue:`60766`)


Styler
^^^^^^
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/arrays/masked.py
Original file line number Diff line number Diff line change
Expand Up @@ -1325,7 +1325,7 @@ def max(self, *, skipna: bool = True, axis: AxisInt | None = 0, **kwargs):
return self._wrap_reduction_result("max", result, skipna=skipna, axis=axis)

def map(self, mapper, na_action: Literal["ignore"] | None = None):
return map_array(self.to_numpy(), mapper, na_action=na_action)
return map_array(self, mapper, na_action=na_action)

@overload
def any(
Expand Down
18 changes: 18 additions & 0 deletions pandas/tests/arrays/masked/test_basemaskedarray_map.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
import pandas as pd

def test_basemaskedarray_map():
for dtype, data, expected_data in [
("Int32", [1, 2, None, 4], [2, 3, pd.NA, 5]),

]:
s = pd.Series(data, dtype=dtype)

def transform(x):
if x is None:
return x
return x + 1

result = s.map(transform)
expected = pd.Series(expected_data, dtype=result.dtype)

assert result.tolist() == expected.tolist()
11 changes: 11 additions & 0 deletions pandas/tests/extension/test_masked.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,12 @@ class TestMaskedArrays(base.ExtensionTests):
@pytest.mark.parametrize("na_action", [None, "ignore"])
def test_map(self, data_missing, na_action):
result = data_missing.map(lambda x: x, na_action=na_action)
if data_missing.dtype.kind != "b":
for i in range(len(result)):
if result[i] is pd.NA:
result[i] = "nan"
result = result.astype("float64")

if data_missing.dtype == Float32Dtype():
# map roundtrips through objects, which converts to float64
expected = data_missing.to_numpy(dtype="float64", na_value=np.nan)
Expand All @@ -181,10 +187,15 @@ def test_map(self, data_missing, na_action):
def test_map_na_action_ignore(self, data_missing_for_sorting):
zero = data_missing_for_sorting[2]
result = data_missing_for_sorting.map(lambda x: zero, na_action="ignore")

if data_missing_for_sorting.dtype.kind == "b":
expected = np.array([False, pd.NA, False], dtype=object)
else:
expected = np.array([zero, np.nan, zero])
for i in range(len(result)):
if result[i] is pd.NA:
result[i] = "nan"
result = result.astype("float64")
tm.assert_numpy_array_equal(result, expected)

def _get_expected_exception(self, op_name, obj, other):
Expand Down
Loading