Skip to content

Commit 828e743

Browse files
authored
ENH: Implement PandasArray, DTA, TDA interpolate (#53594)
* ENH: implement PandasArray.interpolate * Implement DTA/TDA.interpolate
1 parent 7369347 commit 828e743

File tree

3 files changed

+101
-14
lines changed

3 files changed

+101
-14
lines changed

pandas/core/arrays/datetimelike.py

+42
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@
106106

107107
from pandas.core import (
108108
algorithms,
109+
missing,
109110
nanops,
110111
)
111112
from pandas.core.algorithms import (
@@ -142,6 +143,7 @@
142143
from pandas.tseries import frequencies
143144

144145
if TYPE_CHECKING:
146+
from pandas import Index
145147
from pandas.core.arrays import (
146148
DatetimeArray,
147149
PeriodArray,
@@ -2228,6 +2230,46 @@ def copy(self, order: str = "C") -> Self:
22282230
new_obj._freq = self.freq
22292231
return new_obj
22302232

2233+
def interpolate(
2234+
self,
2235+
*,
2236+
method,
2237+
axis: int,
2238+
index: Index | None,
2239+
limit,
2240+
limit_direction,
2241+
limit_area,
2242+
fill_value,
2243+
inplace: bool,
2244+
**kwargs,
2245+
) -> Self:
2246+
"""
2247+
See NDFrame.interpolate.__doc__.
2248+
"""
2249+
# NB: we return type(self) even if inplace=True
2250+
if method != "linear":
2251+
raise NotImplementedError
2252+
2253+
if inplace:
2254+
out_data = self._ndarray
2255+
else:
2256+
out_data = self._ndarray.copy()
2257+
2258+
missing.interpolate_array_2d(
2259+
out_data,
2260+
method=method,
2261+
axis=axis,
2262+
index=index,
2263+
limit=limit,
2264+
limit_direction=limit_direction,
2265+
limit_area=limit_area,
2266+
fill_value=fill_value,
2267+
**kwargs,
2268+
)
2269+
if inplace:
2270+
return self
2271+
return type(self)._simple_new(out_data, dtype=self.dtype)
2272+
22312273

22322274
# -------------------------------------------------------------------
22332275
# Shared Constructor Helpers

pandas/core/arrays/numpy_.py

+41
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
from pandas.core import (
2121
arraylike,
22+
missing,
2223
nanops,
2324
ops,
2425
)
@@ -33,9 +34,12 @@
3334
Dtype,
3435
NpDtype,
3536
Scalar,
37+
Self,
3638
npt,
3739
)
3840

41+
from pandas import Index
42+
3943

4044
# error: Definition of "_concat_same_type" in base class "NDArrayBacked" is
4145
# incompatible with definition in base class "ExtensionArray"
@@ -220,6 +224,43 @@ def _values_for_factorize(self) -> tuple[np.ndarray, float | None]:
220224
fv = np.nan
221225
return self._ndarray, fv
222226

227+
def interpolate(
228+
self,
229+
*,
230+
method,
231+
axis: int,
232+
index: Index | None,
233+
limit,
234+
limit_direction,
235+
limit_area,
236+
fill_value,
237+
inplace: bool,
238+
**kwargs,
239+
) -> Self:
240+
"""
241+
See NDFrame.interpolate.__doc__.
242+
"""
243+
# NB: we return type(self) even if inplace=True
244+
if inplace:
245+
out_data = self._ndarray
246+
else:
247+
out_data = self._ndarray.copy()
248+
249+
missing.interpolate_array_2d(
250+
out_data,
251+
method=method,
252+
axis=axis,
253+
index=index,
254+
limit=limit,
255+
limit_direction=limit_direction,
256+
limit_area=limit_area,
257+
fill_value=fill_value,
258+
**kwargs,
259+
)
260+
if inplace:
261+
return self
262+
return type(self)._simple_new(out_data, dtype=self.dtype)
263+
223264
# ------------------------------------------------------------------------
224265
# Reductions
225266

pandas/core/internals/blocks.py

+18-14
Original file line numberDiff line numberDiff line change
@@ -1394,27 +1394,27 @@ def interpolate(
13941394
)
13951395

13961396
refs = None
1397+
arr_inplace = inplace
13971398
if inplace:
13981399
if using_cow and self.refs.has_reference():
1399-
data = self.values.copy()
1400+
arr_inplace = False
14001401
else:
1401-
data = self.values
14021402
refs = self.refs
1403-
else:
1404-
data = self.values.copy()
1405-
data = cast(np.ndarray, data) # bc overridden by ExtensionBlock
14061403

1407-
missing.interpolate_array_2d(
1408-
data,
1404+
# Dispatch to the PandasArray method.
1405+
# We know self.array_values is a PandasArray bc EABlock overrides
1406+
new_values = cast(PandasArray, self.array_values).interpolate(
14091407
method=method,
14101408
axis=axis,
14111409
index=index,
14121410
limit=limit,
14131411
limit_direction=limit_direction,
14141412
limit_area=limit_area,
14151413
fill_value=fill_value,
1414+
inplace=arr_inplace,
14161415
**kwargs,
14171416
)
1417+
data = new_values._ndarray
14181418

14191419
nb = self.make_block_same_class(data, refs=refs)
14201420
return nb._maybe_downcast([nb], downcast, using_cow)
@@ -2262,18 +2262,22 @@ def interpolate(
22622262
if method == "linear": # type: ignore[comparison-overlap]
22632263
# TODO: GH#50950 implement for arbitrary EAs
22642264
refs = None
2265+
arr_inplace = inplace
22652266
if using_cow:
22662267
if inplace and not self.refs.has_reference():
2267-
data_out = values._ndarray
22682268
refs = self.refs
22692269
else:
2270-
data_out = values._ndarray.copy()
2271-
else:
2272-
data_out = values._ndarray if inplace else values._ndarray.copy()
2273-
missing.interpolate_array_2d(
2274-
data_out, method=method, limit=limit, index=index, axis=axis
2270+
arr_inplace = False
2271+
2272+
new_values = self.values.interpolate(
2273+
method=method,
2274+
index=index,
2275+
axis=axis,
2276+
inplace=arr_inplace,
2277+
limit=limit,
2278+
fill_value=fill_value,
2279+
**kwargs,
22752280
)
2276-
new_values = type(values)._simple_new(data_out, dtype=values.dtype)
22772281
return self.make_block_same_class(new_values, refs=refs)
22782282

22792283
elif values.ndim == 2 and axis == 0:

0 commit comments

Comments
 (0)