diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index fcd5cd0979252..bc820b5f01f37 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -89,6 +89,7 @@ new_block, to_native_types, ) +from pandas.core.missing import interpolate_array if TYPE_CHECKING: from pandas import Float64Index @@ -367,8 +368,13 @@ def diff(self: T, n: int, axis: int) -> T: axis = 0 return self.apply(algos.diff, n=n, axis=axis) - def interpolate(self: T, **kwargs) -> T: - return self.apply_with_block("interpolate", swap_axis=False, **kwargs) + def interpolate(self: T, axis: int = 0, **kwargs) -> T: + if axis == 0: + return self.apply(interpolate_array, **kwargs) + else: + return self.apply_with_block( + "interpolate", swap_axis=False, axis=axis, **kwargs + ) def shift(self: T, periods: int, axis: int, fill_value) -> T: if fill_value is lib.no_default: diff --git a/pandas/core/missing.py b/pandas/core/missing.py index 8a3d892876b5c..9d328975501a9 100644 --- a/pandas/core/missing.py +++ b/pandas/core/missing.py @@ -23,11 +23,17 @@ ArrayLike, Axis, F, + FillnaOptions, npt, ) from pandas.compat._optional import import_optional_dependency +from pandas.util._validators import validate_bool_kwarg -from pandas.core.dtypes.cast import infer_dtype_from +from pandas.core.dtypes.cast import ( + infer_dtype_from, + maybe_downcast_to_dtype, + soft_convert_objects, +) from pandas.core.dtypes.common import ( is_array_like, is_numeric_v_string_like, @@ -975,3 +981,65 @@ def _rolling_window(a: npt.NDArray[np.bool_], window: int) -> npt.NDArray[np.boo shape = a.shape[:-1] + (a.shape[-1] - window + 1, window) strides = a.strides + (a.strides[-1],) return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides) + + +def _maybe_downcast(arr: np.ndarray, downcast=None): + if arr.dtype == np.dtype(object) and downcast is None: + return soft_convert_objects(arr, datetime=True, numeric=False) + + if downcast: + return maybe_downcast_to_dtype(arr, downcast) + return arr + + +def interpolate_array( + arr: ArrayLike, + method: FillnaOptions = "pad", + axis: int = 0, + index: Index | None = None, + inplace: bool = False, + limit: int | None = None, + limit_direction: str = "forward", + limit_area: str | None = None, + fill_value: Any | None = None, + coerce: bool = False, + downcast: str | None = None, + **kwargs, +) -> ArrayLike: + + inplace = validate_bool_kwarg(inplace, "inplace") + + # first check for extensionarrays + if not isinstance(arr, np.ndarray): + return arr.fillna(value=fill_value, method=method, limit=limit) + + if arr.dtype.kind in ["b", "i", "u"]: + # those dtypes can never hold NAs + # If there are no NAs, then interpolate is a no-op + return arr if inplace else arr.copy() + + try: + m = clean_fill_method(method) + except ValueError: + m = None + if m is None and arr.dtype.kind != "f": + # only deal with floats + # bc we already checked that can_hold_na, we dont have int dtype here + # TODO: make a copy if not inplace? + return arr + + data = arr if inplace else arr.copy() + + interpolate_array_2d( + data, + method=method, + axis=axis, + index=index, + limit=limit, + limit_direction=limit_direction, + limit_area=limit_area, + fill_value=fill_value, + **kwargs, + ) + + return _maybe_downcast(data, downcast)