WIP: PERF: Cythonize fillna

asv_bench/benchmarks/replace.py

@@ -14,10 +14,20 @@ def setup(self, inplace):
         data = np.random.randn(N)
         data[::2] = np.nan
         self.ts = pd.Series(data, index=rng)
+        self.df = pd.DataFrame(np.random.randn(10 ** 3, 10 ** 3))
 
     def time_fillna(self, inplace):
         self.ts.fillna(0.0, inplace=inplace)
 
+    def peakmem_fillna(self, inplace):
+        self.ts.fillna(0.0, inplace=inplace)
+
+    def time_fillna_limit(self, inplace):
+        self.ts.fillna(0.0, inplace=inplace, limit=10 ** 5)
+
+    def time_fillna_df(self, inplace):
+        self.df.fillna(0.0, inplace=inplace)
+
     def time_replace(self, inplace):
         self.ts.replace(np.nan, 0.0, inplace=inplace)
 

pandas/_libs/algos.pyx

@@ -52,6 +52,10 @@ from pandas._libs.khash cimport (
     kh_resize_int64,
     khiter_t,
 )
+from pandas._libs.missing cimport (
+    checknull,
+    checknull_old,
+)
 from pandas._libs.util cimport (
     get_nat,
     numeric,
@@ -63,6 +67,8 @@ cdef:
     float64_t FP_ERR = 1e-13
     float64_t NaN = <float64_t>np.NaN
     int64_t NPY_NAT = get_nat()
+    float64_t INF = <float64_t> np.inf
+    float64_t NEGINF = -INF
 
 cdef enum TiebreakEnumType:
     TIEBREAK_AVERAGE
@@ -832,6 +838,202 @@ def backfill_2d_inplace(algos_t[:, :] values,
     pad_2d_inplace(values[:, ::-1], mask[:, ::-1], limit)
 
 
+# Fillna logic
+# We have our own fused type instead of algos_t
+# since we don't need to support types that can't hold NAs(ints, etc)
+ctypedef fused fillna_t:
+    float64_t
+    float32_t
+    object
+    # TODO: Maybe add datetime64 support through viewing data as int64?
+    # But datetime64 seems to be handled elsewhere
+    int64_t  # Datetime64
+    # TODO: Complex support?
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def fillna1d(fillna_t[:] arr,
+             fillna_t value,
+             Py_ssize_t limit,
+             bint inf_as_na=False
+             ) -> ndarray:
+    """
+    Fills na-like elements inplace for a 1D array
+    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
+    value : object
+        The value to use to replace nans
+    limit : int, default None
+        The number of elements to fill. If None, fills all NaN values
+    inf_as_na:
+        Whether to consider INF and NEGINF as NA
+    """
+    cdef:
+        Py_ssize_t i, N, lim
+        Py_ssize_t count=0
+        fillna_t val
+        bint result
+
+    assert arr.ndim == 1, "'arr' must be 1-D."
+
+    N = len(arr)
+    for i in range(N):
+        val = arr[i]
+        if fillna_t is object:
+            if inf_as_na:
+                result = checknull_old(val)
+            else:
+                result = checknull(val)
+        elif fillna_t is float32_t or fillna_t is float64_t:
+            result = val != val
+            if inf_as_na:
+                result = result and (val == INF or val == NEGINF)
+        else:
+            result = val == NPY_NAT
+        if result and count < limit:
+            arr[i] = value
+            count+=1
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def fillna1d_multi_values(fillna_t[:] arr,
+                          algos_t[:] value,
+                          Py_ssize_t limit,
+                          bint inf_as_na=False
+                          ) -> ndarray:
+    """
+    Fills na-like elements inplace for a 1D array
+    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
+    value : ndarray/ExtensionArray
+        A ndarray/ExtensionArray with same length as arr
+        describing which fill value to use at each position,
+        with a value of np.nan indicating that a position should
+        not be filled
+    limit : int, default None
+        The number of elements to fill. If None, fills all NaN values
+    inf_as_na:
+        Whether to consider INF and NEGINF as NA
+    """
+    cdef:
+        Py_ssize_t i, N
+        Py_ssize_t count=0
+        fillna_t val
+        algos_t fill_value
+        bint result
+
+    assert arr.ndim == 1, "'arr' must be 1-D."
+
+    N = len(arr)
+    for i in range(N):
+        fill_value = value[i]
+        if algos_t is object or algos_t is float64_t or algos_t is float32_t:
+            if fill_value != fill_value:
+                # np.nan don't fill
+                continue
+        val = arr[i]
+        if fillna_t is object:
+            if inf_as_na:
+                result = checknull_old(val)
+            else:
+                result = checknull(val)
+        elif fillna_t is float32_t or fillna_t is float64_t:
+            result = val != val
+            if inf_as_na:
+                result = result and (val == INF or val == NEGINF)
+        else:
+            result = val == NPY_NAT
+        if result and count < limit:
+            # Ugh... We have to cast here since technically could have a int64->float32
+            # There shouldn't be any risk here since BlockManager should check
+            # that the element can be held
+            arr[i] = <fillna_t>fill_value
+            count+=1
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def fillna2d(fillna_t[:, :] arr,
+             fillna_t value,
+             Py_ssize_t limit,
+             bint inf_as_na=False
+             ) -> ndarray:
+    """
+    Fills na-like elements inplace for a 2D array
+    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
+    value : object
+        The value to use to replace nans
+    limit : int, default None
+        The number of elements to fill. If None, fills all NaN values
+    inf_as_na:
+        Whether to consider INF and NEGINF as NA
+    """
+    cdef:
+        Py_ssize_t i, j, n, m
+        Py_ssize_t count=0
+        fillna_t val
+        bint result
+
+    assert arr.ndim == 2, "'arr' must be 2-D."
+
+    n, m = (<object>arr).shape
+    if inf_as_na:
+        check_func = checknull_old
+    else:
+        check_func = checknull
+    for i in range(n):
+        count = 0  # Limit is per axis
+        for j in range(m):
+            val = arr[i, j]
+            if fillna_t is object:
+                if inf_as_na:
+                    result = checknull_old(val)
+                else:
+                    result = checknull(val)
+            elif fillna_t is float32_t or fillna_t is float64_t:
+                result = val != val
+                if inf_as_na:
+                    result = result and (val == INF or val == NEGINF)
+            else:
+                result = val == NPY_NAT
+            if result and count < limit:
+                arr[i, j] = value
+                count+=1
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def is_monotonic(ndarray[algos_t, ndim=1] arr, bint timelike):

pandas/core/internals/blocks.py

@@ -434,40 +434,56 @@ def fillna(
         fillna on the block with the value. If we fail, then convert to
         ObjectBlock and try again
         """
+        # TODO: Handle inf_as_na, we need to get option and pass to cython funcs
         inplace = validate_bool_kwarg(inplace, "inplace")
-
-        mask = isna(self.values)
-        mask, noop = validate_putmask(self.values, mask)
-
-        if limit is not None:
-            limit = libalgos.validate_limit(None, limit=limit)
-            mask[mask.cumsum(self.ndim - 1) > limit] = False
+        arr = self if inplace else self.copy()
+        limit = libalgos.validate_limit(
+            len(self) if self.ndim == 1 else self.shape[1], limit=limit
+        )
 
         if not self._can_hold_na:
-            if inplace:
-                return [self]
+            return [arr]
+
+        if not self.is_extension:
+            if self._can_hold_element(value):
+                if self.ndim == 1:
+                    if is_list_like(value):
+                        # TODO: Verify EA case
+                        if is_extension_array_dtype(value):
+                            mask = value.isna()
+                            value = np.asarray(value[mask], dtype=object)
+                            libalgos.fillna1d_multi_values(
+                                arr.values[mask], value=value, limit=limit
+                            )
+                        else:
+                            libalgos.fillna1d_multi_values(
+                                arr.values, value=value, limit=limit
+                            )
+                    else:
+                        libalgos.fillna1d(arr.values, value=value, limit=limit)
+                else:
+                    libalgos.fillna2d(arr.values, value=value, limit=limit)
+                return arr._maybe_downcast([arr], downcast)
+            elif self.ndim == 1 or self.shape[0] == 1:
+                blk = self.coerce_to_target_dtype(value)
+                # bc we have already cast, inplace=True may avoid an extra copy
+                return blk.fillna(value, limit=limit, inplace=True, downcast=None)
             else:
-                return [self.copy()]
-
-        if self._can_hold_element(value):
-            nb = self if inplace else self.copy()
-            putmask_inplace(nb.values, mask, value)
-            return nb._maybe_downcast([nb], downcast)
-
-        if noop:
-            # we can't process the value, but nothing to do
-            return [self] if inplace else [self.copy()]
-
-        elif self.ndim == 1 or self.shape[0] == 1:
-            blk = self.coerce_to_target_dtype(value)
-            # bc we have already cast, inplace=True may avoid an extra copy
-            return blk.fillna(value, limit=limit, inplace=True, downcast=None)
-
+                # operate column-by-column
+                return self.split_and_operate(
+                    type(self).fillna,
+                    value,
+                    limit=limit,
+                    inplace=inplace,
+                    downcast=None,
+                )
         else:
-            # operate column-by-column
-            return self.split_and_operate(
-                type(self).fillna, value, limit=limit, inplace=inplace, downcast=None
-            )
+            # TODO: This seems to work for EAS, verify it does
+            return [
+                self.make_block_same_class(
+                    values=self.values.fillna(value=value, limit=limit)
+                )
+            ]
 
     @final
     def _split(self) -> list[Block]: