ENH add na_action to DataFrame.applymap

doc/source/whatsnew/v1.2.0.rst

@@ -50,6 +50,7 @@ For example:
 
 Other enhancements
 ^^^^^^^^^^^^^^^^^^
+- :meth:`applymap` now supports ``na_action`` (:issue:`23803`)
 - :class:`Index` with object dtype supports division and multiplication (:issue:`34160`)
 -
 -

pandas/_libs/lib.pyx

@@ -2360,14 +2360,17 @@ def map_infer_mask(ndarray arr, object f, const uint8_t[:] mask, bint convert=Tr
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def map_infer(ndarray arr, object f, bint convert=True):
+def map_infer(ndarray arr, object f, bint convert=True, bint ignore_na=False):
     """
     Substitute for np.vectorize with pandas-friendly dtype inference.
 
     Parameters
     ----------
     arr : ndarray
     f : function
+    convert : bint
+    ignore_na : bint
+        If True, NA values will not have f applied
 
     Returns
     -------
@@ -2381,6 +2384,9 @@ def map_infer(ndarray arr, object f, bint convert=True):
     n = len(arr)
     result = np.empty(n, dtype=object)
     for i in range(n):
+        if ignore_na and checknull(arr[i]):
+            result[i] = arr[i]
+            continue
         val = f(arr[i])
 
         if cnp.PyArray_IsZeroDim(val):

pandas/core/frame.py

@@ -7584,7 +7584,7 @@ def apply(self, func, axis=0, raw=False, result_type=None, args=(), **kwds):
         )
         return op.get_result()
 
-    def applymap(self, func) -> "DataFrame":
+    def applymap(self, func, na_action=None) -> "DataFrame":
         """
         Apply a function to a Dataframe elementwise.
 
@@ -7595,6 +7595,10 @@ def applymap(self, func) -> "DataFrame":
         ----------
         func : callable
             Python function, returns a single value from a single value.
+        na_action : {None, 'ignore'}, default None
+            If ‘ignore’, propagate NaN values, without passing them to func.
+
+            .. versionadded:: 1.1
 
         Returns
         -------
@@ -7618,6 +7622,15 @@ def applymap(self, func) -> "DataFrame":
         0  3  4
         1  5  5
 
+        Like Series.map, NA values can be ignored:
+
+        >>> df_copy = df.copy()
+        >>> df_copy.iloc[0, 0] = pd.NA
+        >>> df_copy.applymap(lambda x: len(str(x)), na_action='ignore')
+              0  1
+        0  <NA>  4
+        1     5  5
+
         Note that a vectorized version of `func` often exists, which will
         be much faster. You could square each number elementwise.
 
@@ -7633,11 +7646,17 @@ def applymap(self, func) -> "DataFrame":
         0   1.000000   4.494400
         1  11.262736  20.857489
         """
+        if na_action not in {"ignore", None}:
+            raise ValueError(
+                f"na_action must be 'ignore' or None. Got {repr(na_action)}"
+            )
+        ignore_na = na_action == "ignore"
+
         # if we have a dtype == 'M8[ns]', provide boxed values
         def infer(x):
             if x.empty:
-                return lib.map_infer(x, func)
-            return lib.map_infer(x.astype(object)._values, func)
+                return lib.map_infer(x, func, ignore_na=ignore_na)
+            return lib.map_infer(x.astype(object)._values, func, ignore_na=ignore_na)
 
         return self.apply(infer)
 

pandas/tests/frame/apply/test_frame_apply.py

@@ -587,6 +587,20 @@ def test_applymap(self, float_frame):
         tm.assert_frame_equal(applied, float_frame * 2)
         float_frame.applymap(type)
 
+        # GH 23803: na_ignore
+        strlen_frame = float_frame.applymap(lambda x: len(str(x)))
+        float_frame_with_na = float_frame.copy()
+        mask = np.random.randint(0, 2, size=float_frame.shape, dtype=bool)
+        float_frame_with_na[mask] = pd.NA
+        strlen_frame_na_ignore = float_frame_with_na.applymap(
+            lambda x: len(str(x)), na_action="ignore"
+        )
+        strlen_frame_with_na = strlen_frame.copy()
+        strlen_frame_with_na[mask] = pd.NA
+        print(strlen_frame_na_ignore)
+        print(strlen_frame_with_na)
+        tm.assert_frame_equal(strlen_frame_na_ignore, strlen_frame_with_na)
+
         # GH 465: function returning tuples
         result = float_frame.applymap(lambda x: (x, x))
         assert isinstance(result["A"][0], tuple)