diff --git a/pandas/core/arrays/base.py b/pandas/core/arrays/base.py index eae401f9744f0..c2fc72ff753a8 100644 --- a/pandas/core/arrays/base.py +++ b/pandas/core/arrays/base.py @@ -32,7 +32,7 @@ from pandas.core import ops from pandas.core.algorithms import factorize_array, unique -from pandas.core.missing import backfill_1d, pad_1d +from pandas.core.missing import get_fill_func from pandas.core.sorting import nargminmax, nargsort _extension_array_shared_docs: Dict[str, str] = dict() @@ -616,7 +616,7 @@ def fillna(self, value=None, method=None, limit=None): if mask.any(): if method is not None: - func = pad_1d if method == "pad" else backfill_1d + func = get_fill_func(method) new_values = func(self.astype(object), limit=limit, mask=mask) new_values = self._from_sequence(new_values, dtype=self.dtype) else: diff --git a/pandas/core/missing.py b/pandas/core/missing.py index c2926debcb6d6..f3229b2876e5d 100644 --- a/pandas/core/missing.py +++ b/pandas/core/missing.py @@ -562,20 +562,22 @@ def interpolate_2d( values = values.reshape(tuple((1,) + values.shape)) method = clean_fill_method(method) + tvalues = transf(values) if method == "pad": - values = transf(pad_2d(transf(values), limit=limit)) + result = _pad_2d(tvalues, limit=limit) else: - values = transf(backfill_2d(transf(values), limit=limit)) + result = _backfill_2d(tvalues, limit=limit) + result = transf(result) # reshape back if ndim == 1: - values = values[0] + result = result[0] if orig_values.dtype.kind == "M": # convert float back to datetime64 - values = values.astype(orig_values.dtype) + result = result.astype(orig_values.dtype) - return values + return result def _cast_values_for_fillna(values, dtype: DtypeObj, has_mask: bool): @@ -597,10 +599,9 @@ def _cast_values_for_fillna(values, dtype: DtypeObj, has_mask: bool): return values -def _fillna_prep(values, mask=None, dtype: Optional[DtypeObj] = None): - # boilerplate for pad_1d, backfill_1d, pad_2d, backfill_2d - if dtype is None: - dtype = values.dtype +def _fillna_prep(values, mask=None): + # boilerplate for _pad_1d, _backfill_1d, _pad_2d, _backfill_2d + dtype = values.dtype has_mask = mask is not None if not has_mask: @@ -613,20 +614,20 @@ def _fillna_prep(values, mask=None, dtype: Optional[DtypeObj] = None): return values, mask -def pad_1d(values, limit=None, mask=None, dtype: Optional[DtypeObj] = None): - values, mask = _fillna_prep(values, mask, dtype) +def _pad_1d(values, limit=None, mask=None): + values, mask = _fillna_prep(values, mask) algos.pad_inplace(values, mask, limit=limit) return values -def backfill_1d(values, limit=None, mask=None, dtype: Optional[DtypeObj] = None): - values, mask = _fillna_prep(values, mask, dtype) +def _backfill_1d(values, limit=None, mask=None): + values, mask = _fillna_prep(values, mask) algos.backfill_inplace(values, mask, limit=limit) return values -def pad_2d(values, limit=None, mask=None, dtype: Optional[DtypeObj] = None): - values, mask = _fillna_prep(values, mask, dtype) +def _pad_2d(values, limit=None, mask=None): + values, mask = _fillna_prep(values, mask) if np.all(values.shape): algos.pad_2d_inplace(values, mask, limit=limit) @@ -636,8 +637,8 @@ def pad_2d(values, limit=None, mask=None, dtype: Optional[DtypeObj] = None): return values -def backfill_2d(values, limit=None, mask=None, dtype: Optional[DtypeObj] = None): - values, mask = _fillna_prep(values, mask, dtype) +def _backfill_2d(values, limit=None, mask=None): + values, mask = _fillna_prep(values, mask) if np.all(values.shape): algos.backfill_2d_inplace(values, mask, limit=limit) @@ -647,7 +648,7 @@ def backfill_2d(values, limit=None, mask=None, dtype: Optional[DtypeObj] = None) return values -_fill_methods = {"pad": pad_1d, "backfill": backfill_1d} +_fill_methods = {"pad": _pad_1d, "backfill": _backfill_1d} def get_fill_func(method): @@ -724,7 +725,7 @@ def inner(invalid, limit): return f_idx & b_idx -def _rolling_window(a, window): +def _rolling_window(a: np.ndarray, window: int): """ [True, True, False, True, False], 2 ->