Skip to content

Commit bd3f2a4

Browse files
authored
REF: move putmask internals in array_algos.putmask (#38833)
1 parent 59c5be2 commit bd3f2a4

File tree

2 files changed

+40
-33
lines changed

2 files changed

+40
-33
lines changed

pandas/core/array_algos/putmask.py

+34
Original file line numberDiff line numberDiff line change
@@ -120,3 +120,37 @@ def _putmask_preserve(new_values: np.ndarray, new, mask: np.ndarray):
120120
except (IndexError, ValueError):
121121
new_values[mask] = new
122122
return new_values
123+
124+
125+
def putmask_without_repeat(values: np.ndarray, mask: np.ndarray, new: Any) -> None:
126+
"""
127+
np.putmask will truncate or repeat if `new` is a listlike with
128+
len(new) != len(values). We require an exact match.
129+
130+
Parameters
131+
----------
132+
values : np.ndarray
133+
mask : np.ndarray[bool]
134+
new : Any
135+
"""
136+
if getattr(new, "ndim", 0) >= 1:
137+
new = new.astype(values.dtype, copy=False)
138+
139+
# TODO: this prob needs some better checking for 2D cases
140+
nlocs = mask.sum()
141+
if nlocs > 0 and is_list_like(new) and getattr(new, "ndim", 1) == 1:
142+
if nlocs == len(new):
143+
# GH#30567
144+
# If length of ``new`` is less than the length of ``values``,
145+
# `np.putmask` would first repeat the ``new`` array and then
146+
# assign the masked values hence produces incorrect result.
147+
# `np.place` on the other hand uses the ``new`` values at it is
148+
# to place in the masked locations of ``values``
149+
np.place(values, mask, new)
150+
# i.e. values[mask] = new
151+
elif mask.shape[-1] == len(new) or len(new) == 1:
152+
np.putmask(values, mask, new)
153+
else:
154+
raise ValueError("cannot assign mismatch length to masked array")
155+
else:
156+
np.putmask(values, mask, new)

pandas/core/internals/blocks.py

+6-33
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,11 @@
5454
from pandas.core.dtypes.missing import is_valid_nat_for_dtype, isna
5555

5656
import pandas.core.algorithms as algos
57-
from pandas.core.array_algos.putmask import putmask_inplace, putmask_smart
57+
from pandas.core.array_algos.putmask import (
58+
putmask_inplace,
59+
putmask_smart,
60+
putmask_without_repeat,
61+
)
5862
from pandas.core.array_algos.replace import compare_or_regex_search, replace_regex
5963
from pandas.core.array_algos.transforms import shift
6064
from pandas.core.arrays import (
@@ -1030,38 +1034,7 @@ def putmask(self, mask, new, axis: int = 0) -> List["Block"]:
10301034
if transpose:
10311035
new_values = new_values.T
10321036

1033-
# If the default repeat behavior in np.putmask would go in the
1034-
# wrong direction, then explicitly repeat and reshape new instead
1035-
if getattr(new, "ndim", 0) >= 1:
1036-
new = new.astype(new_values.dtype, copy=False)
1037-
1038-
# we require exact matches between the len of the
1039-
# values we are setting (or is compat). np.putmask
1040-
# doesn't check this and will simply truncate / pad
1041-
# the output, but we want sane error messages
1042-
#
1043-
# TODO: this prob needs some better checking
1044-
# for 2D cases
1045-
if (
1046-
is_list_like(new)
1047-
and np.any(mask[mask])
1048-
and getattr(new, "ndim", 1) == 1
1049-
):
1050-
if mask[mask].shape[-1] == len(new):
1051-
# GH 30567
1052-
# If length of ``new`` is less than the length of ``new_values``,
1053-
# `np.putmask` would first repeat the ``new`` array and then
1054-
# assign the masked values hence produces incorrect result.
1055-
# `np.place` on the other hand uses the ``new`` values at it is
1056-
# to place in the masked locations of ``new_values``
1057-
np.place(new_values, mask, new)
1058-
# i.e. new_values[mask] = new
1059-
elif mask.shape[-1] == len(new) or len(new) == 1:
1060-
np.putmask(new_values, mask, new)
1061-
else:
1062-
raise ValueError("cannot assign mismatch length to masked array")
1063-
else:
1064-
np.putmask(new_values, mask, new)
1037+
putmask_without_repeat(new_values, mask, new)
10651038

10661039
# maybe upcast me
10671040
elif mask.any():

0 commit comments

Comments
 (0)