Skip to content

Commit 75f3804

Browse files
push limit_area down into interpolate_2d
1 parent 1b1ab72 commit 75f3804

File tree

3 files changed

+64
-36
lines changed

3 files changed

+64
-36
lines changed

pandas/core/arrays/categorical.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -1646,9 +1646,9 @@ def fillna(self, value=None, method=None, limit=None):
16461646

16471647
# TODO: dispatch when self.categories is EA-dtype
16481648
values = np.asarray(self).reshape(-1, len(self))
1649-
values = interpolate_2d(values, method, 0, None, value).astype(
1650-
self.categories.dtype
1651-
)[0]
1649+
values = interpolate_2d(
1650+
values, method=method, axis=0, limit=None, fill_value=value
1651+
).astype(self.categories.dtype)[0]
16521652
codes = _get_codes_for_values(values, self.categories)
16531653

16541654
else:

pandas/core/internals/blocks.py

+4-32
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from datetime import datetime, timedelta
2-
import functools
32
import inspect
43
import re
54
from typing import TYPE_CHECKING, Any, List, Optional
@@ -1238,43 +1237,16 @@ def _interpolate_with_fill(
12381237
# We only get here for non-ExtensionBlock
12391238
fill_value = convert_scalar_for_putitemlike(fill_value, self.values.dtype)
12401239

1241-
interpolate_2d = functools.partial(
1242-
missing.interpolate_2d,
1240+
values = missing.interpolate_2d(
1241+
values,
12431242
method=method,
1243+
axis=axis,
12441244
limit=limit,
1245+
limit_area=limit_area,
12451246
fill_value=fill_value,
12461247
dtype=self.dtype,
12471248
)
12481249

1249-
if limit_area is None:
1250-
values = interpolate_2d(values, axis=axis)
1251-
else:
1252-
1253-
def func(values):
1254-
invalid = isna(values)
1255-
1256-
if not invalid.any():
1257-
return values
1258-
1259-
if not invalid.all():
1260-
first = missing.find_valid_index(values, "first")
1261-
last = missing.find_valid_index(values, "last")
1262-
1263-
values = interpolate_2d(values)
1264-
1265-
if limit_area == "inside":
1266-
invalid[first : last + 1] = False
1267-
elif limit_area == "outside":
1268-
invalid[:first] = False
1269-
invalid[last + 1 :] = False
1270-
1271-
values[invalid] = np.nan
1272-
else:
1273-
values = interpolate_2d(values)
1274-
return values
1275-
1276-
values = np.apply_along_axis(func, axis, values)
1277-
12781250
blocks = [self.make_block_same_class(values, ndim=self.ndim)]
12791251
return self._maybe_downcast(blocks, downcast)
12801252

pandas/core/missing.py

+57-1
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
Routines for filling missing data.
33
"""
44

5+
import functools
56
from typing import Any, List, Optional, Set, Union
67

78
import numpy as np
@@ -543,12 +544,67 @@ def _cubicspline_interpolate(xi, yi, x, axis=0, bc_type="not-a-knot", extrapolat
543544

544545

545546
def interpolate_2d(
546-
values, method="pad", axis=0, limit=None, fill_value=None, dtype=None
547+
values,
548+
method: str = "pad",
549+
axis: int = 0,
550+
limit: Optional[int] = None,
551+
limit_area: Optional[str] = None,
552+
fill_value: Optional[Any] = None,
553+
dtype=None,
547554
):
548555
"""
549556
Perform an actual interpolation of values, values will be make 2-d if
550557
needed fills inplace, returns the result.
558+
559+
If `limit area` is not None, the interpolataion needs to be done either
560+
rowwise or columnwise
551561
"""
562+
int_2d = functools.partial(
563+
_interpolate_2d,
564+
method=method,
565+
limit=limit,
566+
fill_value=fill_value,
567+
dtype=dtype,
568+
)
569+
570+
if limit_area is None:
571+
return int_2d(values, axis=axis)
572+
else:
573+
574+
def func(values):
575+
invalid = isna(values)
576+
577+
if not invalid.any():
578+
return values
579+
580+
if not invalid.all():
581+
first = find_valid_index(values, "first")
582+
last = find_valid_index(values, "last")
583+
584+
values = int_2d(values)
585+
586+
if limit_area == "inside":
587+
invalid[first : last + 1] = False
588+
elif limit_area == "outside":
589+
invalid[:first] = False
590+
invalid[last + 1 :] = False
591+
592+
values[invalid] = np.nan
593+
else:
594+
values = int_2d(values)
595+
return values
596+
597+
return np.apply_along_axis(func, axis, values)
598+
599+
600+
def _interpolate_2d(
601+
values,
602+
method: str = "pad",
603+
axis: int = 0,
604+
limit: Optional[int] = None,
605+
fill_value: Optional[Any] = None,
606+
dtype=None,
607+
):
552608
orig_values = values
553609

554610
transf = (lambda x: x) if axis == 0 else (lambda x: x.T)

0 commit comments

Comments
 (0)