diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 4615cb4ec7abd..d54d1855ac2f8 100644 --- a/pandas/core/arrays/_mixins.py +++ b/pandas/core/arrays/_mixins.py @@ -278,8 +278,12 @@ def fillna( if mask.any(): if method is not None: - func = missing.get_fill_func(method) - new_values, _ = func(self._ndarray.copy(), limit=limit, mask=mask) + # TODO: check value is None + # (for now) when self.ndim == 2, we assume axis=0 + func = missing.get_fill_func(method, ndim=self.ndim) + new_values, _ = func(self._ndarray.T.copy(), limit=limit, mask=mask.T) + new_values = new_values.T + # TODO: PandasArray didn't used to copy, need tests for this new_values = self._from_backing_data(new_values) else: diff --git a/pandas/core/arrays/period.py b/pandas/core/arrays/period.py index 96a159c0804c9..7e9e13400e11f 100644 --- a/pandas/core/arrays/period.py +++ b/pandas/core/arrays/period.py @@ -639,6 +639,14 @@ def searchsorted(self, value, side="left", sorter=None) -> np.ndarray: m8arr = self._ndarray.view("M8[ns]") return m8arr.searchsorted(value, side=side, sorter=sorter) + def fillna(self, value=None, method=None, limit=None) -> PeriodArray: + if method is not None: + # view as dt64 so we get treated as timelike in core.missing + dta = self.view("M8[ns]") + result = dta.fillna(value=value, method=method, limit=limit) + return result.view(self.dtype) + return super().fillna(value=value, method=method, limit=limit) + # ------------------------------------------------------------------ # Arithmetic Methods diff --git a/pandas/core/missing.py b/pandas/core/missing.py index 1b5a7237b5287..dc42a175409c2 100644 --- a/pandas/core/missing.py +++ b/pandas/core/missing.py @@ -646,8 +646,6 @@ def interpolate_2d( values, ) - orig_values = values - transf = (lambda x: x) if axis == 0 else (lambda x: x.T) # reshape a 1 dim if needed @@ -669,10 +667,6 @@ def interpolate_2d( if ndim == 1: result = result[0] - if orig_values.dtype.kind in ["m", "M"]: - # convert float back to datetime64/timedelta64 - result = result.view(orig_values.dtype) - return result @@ -755,9 +749,11 @@ def _backfill_2d(values, limit=None, mask=None): _fill_methods = {"pad": _pad_1d, "backfill": _backfill_1d} -def get_fill_func(method): +def get_fill_func(method, ndim: int = 1): method = clean_fill_method(method) - return _fill_methods[method] + if ndim == 1: + return _fill_methods[method] + return {"pad": _pad_2d, "backfill": _backfill_2d}[method] def clean_reindex_fill_method(method): diff --git a/pandas/tests/arrays/test_datetimes.py b/pandas/tests/arrays/test_datetimes.py index d159d76030250..8e6c330475e68 100644 --- a/pandas/tests/arrays/test_datetimes.py +++ b/pandas/tests/arrays/test_datetimes.py @@ -195,6 +195,37 @@ def test_fillna_preserves_tz(self, method): assert arr[2] is pd.NaT assert dti[2] == pd.Timestamp("2000-01-03", tz="US/Central") + def test_fillna_2d(self): + dti = pd.date_range("2016-01-01", periods=6, tz="US/Pacific") + dta = dti._data.reshape(3, 2).copy() + dta[0, 1] = pd.NaT + dta[1, 0] = pd.NaT + + res1 = dta.fillna(method="pad") + expected1 = dta.copy() + expected1[1, 0] = dta[0, 0] + tm.assert_extension_array_equal(res1, expected1) + + res2 = dta.fillna(method="backfill") + expected2 = dta.copy() + expected2 = dta.copy() + expected2[1, 0] = dta[2, 0] + expected2[0, 1] = dta[1, 1] + tm.assert_extension_array_equal(res2, expected2) + + # with different ordering for underlying ndarray; behavior should + # be unchanged + dta2 = dta._from_backing_data(dta._ndarray.copy(order="F")) + assert dta2._ndarray.flags["F_CONTIGUOUS"] + assert not dta2._ndarray.flags["C_CONTIGUOUS"] + tm.assert_extension_array_equal(dta, dta2) + + res3 = dta2.fillna(method="pad") + tm.assert_extension_array_equal(res3, expected1) + + res4 = dta2.fillna(method="backfill") + tm.assert_extension_array_equal(res4, expected2) + def test_array_interface_tz(self): tz = "US/Central" data = DatetimeArray(pd.date_range("2017", periods=2, tz=tz)) diff --git a/pandas/tests/extension/base/dim2.py b/pandas/tests/extension/base/dim2.py index fbe2537e8a7bf..073880d79d872 100644 --- a/pandas/tests/extension/base/dim2.py +++ b/pandas/tests/extension/base/dim2.py @@ -131,6 +131,17 @@ def test_concat_2d(self, data): with pytest.raises(ValueError): left._concat_same_type([left, right], axis=2) + @pytest.mark.parametrize("method", ["backfill", "pad"]) + def test_fillna_2d_method(self, data_missing, method): + arr = data_missing.repeat(2).reshape(2, 2) + assert arr[0].isna().all() + assert not arr[1].isna().any() + + result = arr.fillna(method=method) + + expected = data_missing.fillna(method=method).repeat(2).reshape(2, 2) + self.assert_extension_array_equal(result, expected) + @pytest.mark.parametrize("method", ["mean", "median", "var", "std", "sum", "prod"]) def test_reductions_2d_axis_none(self, data, method, request): if not hasattr(data, method):