PERF/CLN: Use Numpy C Iterator API to unify isnaobj 1D/2D

pandas/_libs/algos.pyx

@@ -1366,7 +1366,7 @@ def rank_2d(
         nan_fill_val = get_rank_nan_fill_val(nans_rank_highest, <numeric_object_t>0)
 
         if numeric_object_t is object:
-            mask = missing.isnaobj2d(values).view(np.uint8)
+            mask = missing.isnaobj(values).view(np.uint8)
         elif numeric_object_t is float64_t or numeric_object_t is float32_t:
             mask = np.isnan(values).view(np.uint8)
         else:

pandas/_libs/missing.pyx

@@ -4,10 +4,16 @@ from sys import maxsize
 
 cimport cython
 from cython cimport Py_ssize_t
+
 import numpy as np
 
 cimport numpy as cnp
 from numpy cimport (
+    PyArray_GETITEM,
+    PyArray_ITER_DATA,
+    PyArray_ITER_NEXT,
+    PyArray_IterNew,
+    flatiter,
     float64_t,
     int64_t,
     ndarray,
@@ -197,56 +203,22 @@ cpdef ndarray[uint8_t] isnaobj(ndarray arr, bint inf_as_na=False):
     result : ndarray (dtype=np.bool_)
     """
     cdef:
-        Py_ssize_t i, n
+        Py_ssize_t i, n = arr.size
         object val
-        ndarray[uint8_t] result
-
-    assert arr.ndim == 1, "'arr' must be 1-D."
-
-    n = len(arr)
-    result = np.empty(n, dtype=np.uint8)
-    for i in range(n):
-        val = arr[i]
-        result[i] = checknull(val, inf_as_na=inf_as_na)
-    return result.view(np.bool_)
-
-
-@cython.wraparound(False)
-@cython.boundscheck(False)
-def isnaobj2d(arr: ndarray, inf_as_na: bool = False) -> ndarray:
-    """
-    Return boolean mask denoting which elements of a 2-D array are na-like,
-    according to the criteria defined in `checknull`:
-     - None
-     - nan
-     - NaT
-     - np.datetime64 representation of NaT
-     - np.timedelta64 representation of NaT
-     - NA
-     - Decimal("NaN")
-
-    Parameters
-    ----------
-    arr : ndarray
-
-    Returns
-    -------
-    result : ndarray (dtype=np.bool_)
-    """
-    cdef:
-        Py_ssize_t i, j, n, m
-        object val
-        ndarray[uint8_t, ndim=2] result
-
-    assert arr.ndim == 2, "'arr' must be 2-D."
+        bint is_null
+        ndarray result = np.empty((<object>arr).shape, dtype=np.uint8)
+        flatiter it = PyArray_IterNew(arr)
+        flatiter it2 = PyArray_IterNew(result)
 
-    n, m = (<object>arr).shape
-    result = np.zeros((n, m), dtype=np.uint8)
     for i in range(n):
-        for j in range(m):
-            val = arr[i, j]
-            if checknull(val, inf_as_na=inf_as_na):
-                result[i, j] = 1
+        # The PyArray_GETITEM and PyArray_ITER_NEXT are faster
+        #  equivalents to `val = values[i]`
+        val = PyArray_GETITEM(arr, PyArray_ITER_DATA(it))
+        PyArray_ITER_NEXT(it)
+        is_null = checknull(val, inf_as_na=inf_as_na)
+        # Dereference pointer (set value)
+        (<uint8_t *>(PyArray_ITER_DATA(it2)))[0] = <uint8_t>is_null
+        PyArray_ITER_NEXT(it2)
     return result.view(np.bool_)
 
 

pandas/core/dtypes/missing.py

@@ -313,10 +313,8 @@ def _isna_string_dtype(values: np.ndarray, inf_as_na: bool) -> npt.NDArray[np.bo
         result = np.zeros(values.shape, dtype=bool)
     else:
 
-        if values.ndim == 1:
+        if values.ndim == 1 or values.ndim == 2:
             result = libmissing.isnaobj(values, inf_as_na=inf_as_na)
-        elif values.ndim == 2:
-            result = libmissing.isnaobj2d(values, inf_as_na=inf_as_na)
         else:
             # 0-D, reached via e.g. mask_missing
             result = libmissing.isnaobj(values.ravel(), inf_as_na=inf_as_na)

pandas/tests/dtypes/test_missing.py

@@ -624,9 +624,9 @@ def _check_behavior(self, arr, expected):
         arr = np.atleast_2d(arr)
         expected = np.atleast_2d(expected)
 
-        result = libmissing.isnaobj2d(arr)
+        result = libmissing.isnaobj(arr)
         tm.assert_numpy_array_equal(result, expected)
-        result = libmissing.isnaobj2d(arr, inf_as_na=True)
+        result = libmissing.isnaobj(arr, inf_as_na=True)
         tm.assert_numpy_array_equal(result, expected)
 
     def test_basic(self):