BUG: identity checking NA in map incorrect

doc/source/whatsnew/v3.0.0.rst

@@ -412,6 +412,7 @@ Numeric
 
 Conversion
 ^^^^^^^^^^
+- Bug in :meth:`BaseMaskedArray.map` was casting ``pd.NA`` to ``np.nan``. (:issue:`57390`)
 - Bug in :meth:`DataFrame.astype` not casting ``values`` for Arrow-based dictionary dtype correctly (:issue:`58479`)
 - Bug in :meth:`DataFrame.update` bool dtype being converted to object (:issue:`55509`)
 - Bug in :meth:`Series.astype` might modify read-only array inplace when casting to a string dtype (:issue:`57212`)

pandas/_libs/lib.pyi

@@ -74,6 +74,8 @@ def map_infer(
     *,
     convert: Literal[False],
     ignore_na: bool = ...,
+    mask: npt.NDArray[np.bool_] | None = ...,
+    na_value: Any = ...,
 ) -> np.ndarray: ...
 @overload
 def map_infer(
@@ -82,6 +84,8 @@ def map_infer(
     *,
     convert: bool = ...,
     ignore_na: bool = ...,
+    mask: npt.NDArray[np.bool_] | None = ...,
+    na_value: Any = ...,
 ) -> ArrayLike: ...
 @overload
 def maybe_convert_objects(

pandas/_libs/lib.pyx

@@ -2931,7 +2931,13 @@ def map_infer_mask(
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def map_infer(
-    ndarray arr, object f, *, bint convert=True, bint ignore_na=False
+    ndarray arr,
+    object f,
+    *,
+    bint convert=True,
+    bint ignore_na=False,
+    const uint8_t[:] mask=None,
+    object na_value=None,
 ) -> "ArrayLike":
     """
     Substitute for np.vectorize with pandas-friendly dtype inference.
@@ -2943,6 +2949,10 @@ def map_infer(
     convert : bint
     ignore_na : bint
         If True, NA values will not have f applied
+    mask : ndarray, optional
+        uint8 dtype ndarray indicating na_value to apply `f` to.
+    na_value : Any, optional
+        The input value to use for masked values.
 
     Returns
     -------
@@ -2959,7 +2969,10 @@ def map_infer(
         if ignore_na and checknull(arr[i]):
             result[i] = arr[i]
             continue
-        val = f(arr[i])
+        elif mask is not None and na_value is not None and mask[i]:
+            val = f(na_value)
+        else:
+            val = f(arr[i])
 
         if cnp.PyArray_IsZeroDim(val):
             # unbox 0-dim arrays, GH#690

pandas/core/algorithms.py

@@ -1678,9 +1678,23 @@ def map_array(
     if not len(arr):
         return arr.copy()
 
-    # we must convert to python types
-    values = arr.astype(object, copy=False)
+    na_value = None
+    mask = isna(arr)
+    if isinstance(arr.dtype, BaseMaskedDtype) and na_action is None:
+        arr = cast("BaseMaskedArray", arr)
+        values = arr._data
+        if arr._hasna:
+            na_value = arr.dtype.na_value
+    else:
+        # we must convert to python types
+        values = arr.astype(object, copy=False)
+
     if na_action is None:
-        return lib.map_infer(values, mapper)
+        return lib.map_infer(
+            values,
+            mapper,
+            mask=mask,
+            na_value=na_value,
+        )
     else:
-        return lib.map_infer_mask(values, mapper, mask=isna(values).view(np.uint8))
+        return lib.map_infer_mask(values, mapper, mask=mask)

pandas/core/arrays/masked.py

@@ -1319,7 +1319,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(

pandas/tests/extension/test_masked.py

@@ -17,6 +17,7 @@
 import numpy as np
 import pytest
 
+from pandas._libs import missing as libmissing
 from pandas.compat import (
     IS64,
     is_platform_windows,
@@ -171,17 +172,13 @@ 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 == Float32Dtype():
-            # map roundtrips through objects, which converts to float64
-            expected = data_missing.to_numpy(dtype="float64", na_value=np.nan)
-        else:
-            expected = data_missing.to_numpy()
+        expected = data_missing.astype(object, copy=True)
         tm.assert_numpy_array_equal(result, expected)
 
     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":
+        if data_missing_for_sorting.dtype.na_value is libmissing.NA:
             expected = np.array([False, pd.NA, False], dtype=object)
         else:
             expected = np.array([zero, np.nan, zero])

pandas/tests/series/methods/test_map.py

@@ -261,6 +261,13 @@ def test_map_int():
     assert not isna(merged["c"])
 
 
+def test_map_int_with_pd_na():
+    s = Series([pd.NA, 11, 22, pd.NA], dtype="Int64")
+    result = s.map(lambda x: 5 if x is pd.NA else 2 * x)
+    expected = Series([5, 22, 44, 5])
+    tm.assert_series_equal(result, expected)
+
+
 def test_map_type_inference():
     s = Series(range(3))
     s2 = s.map(lambda x: np.where(x == 0, 0, 1))