ENH: DataFrame.map uses .map method of ExtensionArrays

Author: topper-123

doc/source/whatsnew/v2.1.0.rst

@@ -34,6 +34,7 @@ Other enhancements
 - Improve error message when setting :class:`DataFrame` with wrong number of columns through :meth:`DataFrame.isetitem` (:issue:`51701`)
 - Let :meth:`DataFrame.to_feather` accept a non-default :class:`Index` and non-string column names (:issue:`51787`)
 - :class:`api.extensions.ExtensionArray` now has a :meth:`~api.extensions.ExtensionArray.map` method (:issue:`51809`)
+- :meth:`DataFrame.applymap` now uses the :meth:`~api.extensions.ExtensionArray.map` method for of the underlying arrays for :class:`api.extensions.ExtensionArray` instances (:issue:`xxxxx`)
 - Improve error message when having incompatible columns using :meth:`DataFrame.merge` (:issue:`51861`)
 - Improved error message when creating a DataFrame with empty data (0 rows), no index and an incorrect number of columns. (:issue:`52084`)
 - :meth:`arrays.SparseArray.map` now supports ``na_action`` (:issue:`52096`).

pandas/core/frame.py

@@ -9955,14 +9955,14 @@ def applymap(
             raise ValueError(
                 f"na_action must be 'ignore' or None. Got {repr(na_action)}"
             )
-        ignore_na = na_action == "ignore"
+
+        if self.empty:
+            return self.copy()
+
         func = functools.partial(func, **kwargs)
 
-        # if we have a dtype == 'M8[ns]', provide boxed values
         def infer(x):
-            if x.empty:
-                return lib.map_infer(x, func, ignore_na=ignore_na)
-            return lib.map_infer(x.astype(object)._values, func, ignore_na=ignore_na)
+            return x._map_values(func, na_action=na_action)
 
         return self.apply(infer).__finalize__(self, "applymap")
 

pandas/tests/apply/test_frame_apply.py

@@ -546,6 +546,29 @@ def test_applymap_float_object_conversion(val):
     assert result == object
 
 
+@pytest.mark.parametrize("na_action", ["ignore", None])
+def test_applymap_keeps_dtype(na_action):
+    # GHxxxxx
+    arr = pd.Series(["a", np.nan, "b"])
+    sparse_arr = arr.astype(pd.SparseDtype(object))
+    df = pd.DataFrame(data={"a": arr, "b": sparse_arr})
+
+    def func(x):
+        return str.upper(x) if not pd.isna(x) else x
+
+    result = df.applymap(func, na_action=na_action)
+
+    expected_sparse = pd.array(["A", np.nan, "B"], dtype=pd.SparseDtype(object))
+    expected_arr = expected_sparse.astype(object)
+    expected = pd.DataFrame({"a": expected_arr, "b": expected_sparse})
+
+    tm.assert_frame_equal(result, expected)
+
+    result_empty = df.iloc[:0, :].applymap(func, na_action=na_action)
+    expected_empty = expected.iloc[:0, :]
+    tm.assert_frame_equal(result_empty, expected_empty)
+
+
 def test_applymap_str():
     # GH 2786
     df = DataFrame(np.random.random((3, 4)))