Skip to content

CLN: simplify Block.putmask #38677

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 24, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 7 additions & 25 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1029,9 +1029,7 @@ def _putmask_simple(self, mask: np.ndarray, value: Any):
# GH#37833 np.putmask is more performant than __setitem__
np.putmask(values, mask, value)

def putmask(
self, mask, new, inplace: bool = False, axis: int = 0, transpose: bool = False
) -> List["Block"]:
def putmask(self, mask, new, axis: int = 0) -> List["Block"]:
"""
putmask the data to the block; it is possible that we may create a
new dtype of block
Expand All @@ -1042,16 +1040,13 @@ def putmask(
----------
mask : np.ndarray[bool], SparseArray[bool], or BooleanArray
new : a ndarray/object
inplace : bool, default False
Perform inplace modification.
axis : int
transpose : bool, default False
Set to True if self is stored with axes reversed.

Returns
-------
List[Block]
"""
transpose = self.ndim == 2
mask = _extract_bool_array(mask)
assert not isinstance(new, (ABCIndex, ABCSeries, ABCDataFrame))

Expand All @@ -1077,8 +1072,6 @@ def putmask(
new = np.repeat(new, new_values.shape[-1]).reshape(self.shape)
new = new.astype(new_values.dtype)

if new_values is self.values and not inplace:
new_values = new_values.copy()
# we require exact matches between the len of the
# values we are setting (or is compat). np.putmask
# doesn't check this and will simply truncate / pad
Expand All @@ -1099,6 +1092,7 @@ def putmask(
# `np.place` on the other hand uses the ``new`` values at it is
# to place in the masked locations of ``new_values``
np.place(new_values, mask, new)
# i.e. new_values[mask] = new
elif mask.shape[-1] == len(new) or len(new) == 1:
np.putmask(new_values, mask, new)
else:
Expand Down Expand Up @@ -1143,18 +1137,10 @@ def f(mask, val, idx):
nv = _putmask_smart(val, mask, n)
return nv

new_blocks = self.split_and_operate(mask, f, inplace)
new_blocks = self.split_and_operate(mask, f, True)
return new_blocks

if inplace:
return [self]

if transpose:
if new_values is None:
new_values = self.values if inplace else self.values.copy()
new_values = new_values.T

return [self.make_block(new_values)]
return [self]

def coerce_to_target_dtype(self, other):
"""
Expand Down Expand Up @@ -1697,17 +1683,13 @@ def set_inplace(self, locs, values):
assert locs.tolist() == [0]
self.values = values

def putmask(
self, mask, new, inplace: bool = False, axis: int = 0, transpose: bool = False
) -> List["Block"]:
def putmask(self, mask, new, axis: int = 0) -> List["Block"]:
"""
See Block.putmask.__doc__
"""
inplace = validate_bool_kwarg(inplace, "inplace")

mask = _extract_bool_array(mask)

new_values = self.values if inplace else self.values.copy()
new_values = self.values

if isinstance(new, (np.ndarray, ExtensionArray)) and len(new) == len(mask):
new = new[mask]
Expand Down
3 changes: 0 additions & 3 deletions pandas/core/internals/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -565,7 +565,6 @@ def setitem(self, indexer, value) -> "BlockManager":
return self.apply("setitem", indexer=indexer, value=value)

def putmask(self, mask, new, align: bool = True, axis: int = 0):
transpose = self.ndim == 2

if align:
align_keys = ["new", "mask"]
Expand All @@ -578,9 +577,7 @@ def putmask(self, mask, new, align: bool = True, axis: int = 0):
align_keys=align_keys,
mask=mask,
new=new,
inplace=True,
axis=axis,
transpose=transpose,
)

def diff(self, n: int, axis: int) -> "BlockManager":
Expand Down