Skip to content

REF: Block.putmask dont use split_and_operate #40281

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 2 commits into from
Mar 8, 2021
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
2 changes: 1 addition & 1 deletion pandas/core/array_algos/putmask.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ def putmask_smart(values: np.ndarray, mask: np.ndarray, new) -> np.ndarray:

# n should be the length of the mask or a scalar here
if not is_list_like(new):
new = np.repeat(new, len(mask))
new = np.broadcast_to(new, mask.shape)

# see if we are only masking values that if putted
# will work in the current dtype
Expand Down
61 changes: 23 additions & 38 deletions pandas/core/internals/blocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1061,64 +1061,49 @@ def putmask(self, mask, new) -> List[Block]:
-------
List[Block]
"""
transpose = self.ndim == 2
orig_mask = mask
mask, noop = validate_putmask(self.values.T, mask)
assert not isinstance(new, (ABCIndex, ABCSeries, ABCDataFrame))

new_values = self.values # delay copy if possible.
# if we are passed a scalar None, convert it here
if not is_list_like(new) and isna(new) and not self.is_object:
# FIXME: make sure we have compatible NA
if not self.is_object and is_valid_na_for_dtype(new, self.dtype):
new = self.fill_value

if self._can_hold_element(new):
if transpose:
new_values = new_values.T

putmask_without_repeat(new_values, mask, new)
putmask_without_repeat(self.values.T, mask, new)
return [self]

elif noop:
return [self]

dtype, _ = infer_dtype_from(new)
if dtype.kind in ["m", "M"]:
# using putmask with object dtype will incorrect cast to object
# using putmask with object dtype will incorrectly cast to object
# Having excluded self._can_hold_element, we know we cannot operate
# in-place, so we are safe using `where`
return self.where(new, ~mask)

else:
# may need to upcast
if transpose:
mask = mask.T
if isinstance(new, np.ndarray):
new = new.T

# operate column-by-column
def f(mask, val, idx):

if idx is None:
# ndim==1 case.
n = new
else:

if isinstance(new, np.ndarray):
n = np.squeeze(new[idx % new.shape[0]])
else:
n = np.array(new)

# type of the new block
dtype = find_common_type([n.dtype, val.dtype])

# we need to explicitly astype here to make a copy
n = n.astype(dtype)

nv = putmask_smart(val, mask, n)
return nv
elif self.ndim == 1 or self.shape[0] == 1:
# no need to split columns
nv = putmask_smart(self.values.T, mask, new).T
return [self.make_block(nv)]

new_blocks = self.split_and_operate(mask, f, True)
return new_blocks
else:
is_array = isinstance(new, np.ndarray)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

prob should add a comment that splitting a block


res_blocks = []
nbs = self._split()
for i, nb in enumerate(nbs):
n = new
if is_array:
# we have a different value per-column
n = new[:, i : i + 1]

submask = orig_mask[:, i : i + 1]
rbs = nb.putmask(submask, n)
res_blocks.extend(rbs)
return res_blocks

@final
def coerce_to_target_dtype(self, other) -> Block:
Expand Down