PERF: avoid object conversion in fillna(method=pad|backfill) for masked arrays

asv_bench/benchmarks/frame_methods.py

@@ -351,15 +351,27 @@ def time_isnull_obj(self):
 
 class Fillna:
 
-    params = ([True, False], ["pad", "bfill"])
-    param_names = ["inplace", "method"]
-
-    def setup(self, inplace, method):
+    params = (
+        [True, False],
+        ["pad", "bfill"],
+        [
+            "float64",
+            "float32",
+            "object",
+            "Int64",
+            "Float64",
+        ],
+    )
+    param_names = ["inplace", "method", "dtype"]
+
+    def setup(self, inplace, method, dtype):
         values = np.random.randn(10000, 100)
         values[::2] = np.nan
-        self.df = DataFrame(values)
+        if dtype == "Int64":
+            values = values.round()
+        self.df = DataFrame(values, dtype=dtype)
 
-    def time_frame_fillna(self, inplace, method):
+    def time_frame_fillna(self, inplace, method, dtype):
         self.df.fillna(inplace=inplace, method=method)
 
 

pandas/_libs/algos.pyx

@@ -623,6 +623,38 @@ def pad_inplace(algos_t[:] values, const uint8_t[:] mask, limit=None):
             val = values[i]
 
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def pad_inplace_masked(algos_t[:] values, uint8_t[:] mask, limit=None):
+    cdef:
+        Py_ssize_t i, N
+        algos_t val
+        uint8_t prev_mask
+        int lim, fill_count = 0
+
+    N = len(values)
+
+    # GH#2778
+    if N == 0:
+        return
+
+    lim = validate_limit(N, limit)
+
+    val = values[0]
+    prev_mask = mask[0]
+    for i in range(1, N):
+        if mask[i]:
+            if fill_count >= lim:
+                continue
+            fill_count += 1
+            values[i] = val
+            mask[i] = prev_mask
+        else:
+            fill_count = 0
+            val = values[i]
+            prev_mask = mask[i]
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def pad_2d_inplace(algos_t[:, :] values, const uint8_t[:, :] mask, limit=None):

pandas/core/arrays/masked.py

@@ -28,6 +28,7 @@
     cache_readonly,
     doc,
 )
+from pandas.util._validators import validate_fillna_kwargs
 
 from pandas.core.dtypes.base import ExtensionDtype
 from pandas.core.dtypes.common import (
@@ -38,12 +39,16 @@
     is_string_dtype,
     pandas_dtype,
 )
+from pandas.core.dtypes.inference import is_array_like
 from pandas.core.dtypes.missing import (
     isna,
     notna,
 )
 
-from pandas.core import nanops
+from pandas.core import (
+    missing,
+    nanops,
+)
 from pandas.core.algorithms import (
     factorize_array,
     isin,
@@ -144,6 +149,40 @@ def __getitem__(
 
         return type(self)(self._data[item], self._mask[item])
 
+    @doc(ExtensionArray.fillna)
+    def fillna(
+        self: BaseMaskedArrayT, value=None, method=None, limit=None
+    ) -> BaseMaskedArrayT:
+        value, method = validate_fillna_kwargs(value, method)
+
+        mask = self._mask
+
+        if is_array_like(value):
+            if len(value) != len(self):
+                raise ValueError(
+                    f"Length of 'value' does not match. Got ({len(value)}) "
+                    f" expected {len(self)}"
+                )
+            value = value[mask]
+
+        if mask.any():
+            if method is not None:
+                func = missing.get_fill_func(method)
+                new_values, new_mask = func(
+                    self._data.copy(),
+                    limit=limit,
+                    mask=mask.copy(),
+                    update_mask=True,
+                )
+                return type(self)(new_values, new_mask.view(np.bool_))
+            else:
+                # fill with value
+                new_values = self.copy()
+                new_values[mask] = value
+        else:
+            new_values = self.copy()
+        return new_values
+
     def _coerce_to_array(self, values) -> Tuple[np.ndarray, np.ndarray]:
         raise AbstractMethodError(self)
 

pandas/core/missing.py

@@ -692,16 +692,34 @@ def _fillna_prep(values, mask=None):
     return values, mask
 
 
-def _pad_1d(values, limit=None, mask=None):
+def _pad_1d(
+    values: np.ndarray,
+    limit: int | None = None,
+    mask: np.ndarray | None = None,
+    update_mask: bool = False,
+) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
     values, mask = _fillna_prep(values, mask)
-    algos.pad_inplace(values, mask, limit=limit)
-    return values
+    if update_mask:
+        algos.pad_inplace_masked(values, mask, limit=limit)
+        return values, mask
+    else:
+        algos.pad_inplace(values, mask, limit=limit)
+        return values
 
 
-def _backfill_1d(values, limit=None, mask=None):
+def _backfill_1d(
+    values: np.ndarray,
+    limit: int | None = None,
+    mask: np.ndarray | None = None,
+    update_mask: bool = False,
+) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
     values, mask = _fillna_prep(values, mask)
-    algos.backfill_inplace(values, mask, limit=limit)
-    return values
+    if update_mask:
+        algos.pad_inplace_masked(values[::-1], mask[::-1], limit=limit)
+        return values, mask
+    else:
+        algos.backfill_inplace(values, mask, limit=limit)
+        return values
 
 
 def _pad_2d(values, limit=None, mask=None):