Skip to content

REF: simplify maybe_upcast_putmask #38487

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
Dec 16, 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
51 changes: 9 additions & 42 deletions pandas/core/dtypes/cast.py
Original file line number Diff line number Diff line change
Expand Up @@ -419,9 +419,7 @@ def maybe_cast_to_extension_array(
return result


def maybe_upcast_putmask(
result: np.ndarray, mask: np.ndarray, other: Scalar
) -> Tuple[np.ndarray, bool]:
def maybe_upcast_putmask(result: np.ndarray, mask: np.ndarray) -> np.ndarray:
"""
A safe version of putmask that potentially upcasts the result.

Expand All @@ -435,69 +433,38 @@ def maybe_upcast_putmask(
The destination array. This will be mutated in-place if no upcasting is
necessary.
mask : boolean ndarray
other : scalar
The source value.

Returns
-------
result : ndarray
changed : bool
Set to true if the result array was upcasted.

Examples
--------
>>> arr = np.arange(1, 6)
>>> mask = np.array([False, True, False, True, True])
>>> result, _ = maybe_upcast_putmask(arr, mask, False)
>>> result = maybe_upcast_putmask(arr, mask)
>>> result
array([1, 0, 3, 0, 0])
array([ 1., nan, 3., nan, nan])
"""
if not isinstance(result, np.ndarray):
raise ValueError("The result input must be a ndarray.")
if not is_scalar(other):
# We _could_ support non-scalar other, but until we have a compelling
# use case, we assume away the possibility.
raise ValueError("other must be a scalar")

# NB: we never get here with result.dtype.kind in ["m", "M"]

if mask.any():
# Two conversions for date-like dtypes that can't be done automatically
# in np.place:
# NaN -> NaT
# integer or integer array -> date-like array
if result.dtype.kind in ["m", "M"]:
if isna(other):
other = result.dtype.type("nat")
elif is_integer(other):
other = np.array(other, dtype=result.dtype)

def changeit():
# we are forced to change the dtype of the result as the input
# isn't compatible
r, _ = maybe_upcast(result, fill_value=other, copy=True)
np.place(r, mask, other)

return r, True

# we want to decide whether place will work
# if we have nans in the False portion of our mask then we need to
# upcast (possibly), otherwise we DON't want to upcast (e.g. if we
# have values, say integers, in the success portion then it's ok to not
# upcast)
new_dtype, _ = maybe_promote(result.dtype, other)
new_dtype, _ = maybe_promote(result.dtype, np.nan)
if new_dtype != result.dtype:
result = result.astype(new_dtype, copy=True)

# we have a scalar or len 0 ndarray
# and its nan and we are changing some values
if isna(other):
return changeit()

try:
np.place(result, mask, other)
except TypeError:
# e.g. int-dtype result and float-dtype other
return changeit()
np.place(result, mask, np.nan)

return result, False
return result


def maybe_promote(dtype, fill_value=np.nan):
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/ops/array_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def _masked_arith_op(x: np.ndarray, y, op):
with np.errstate(all="ignore"):
result[mask] = op(xrav[mask], y)

result, _ = maybe_upcast_putmask(result, ~mask, np.nan)
result = maybe_upcast_putmask(result, ~mask)
result = result.reshape(x.shape) # 2D compat
return result

Expand Down
52 changes: 9 additions & 43 deletions pandas/tests/dtypes/cast/test_upcast.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,61 +11,27 @@
def test_upcast_error(result):
# GH23823 require result arg to be ndarray
mask = np.array([False, True, False])
other = np.array([61, 62, 63])
with pytest.raises(ValueError, match="The result input must be a ndarray"):
result, _ = maybe_upcast_putmask(result, mask, other)


@pytest.mark.parametrize(
"arr, other",
[
(np.arange(1, 6), np.array([61, 62, 63])),
(np.arange(1, 6), np.array([61.1, 62.2, 63.3])),
(np.arange(10, 15), np.array([61, 62])),
(np.arange(10, 15), np.array([61, np.nan])),
(
np.arange("2019-01-01", "2019-01-06", dtype="datetime64[D]"),
np.arange("2018-01-01", "2018-01-04", dtype="datetime64[D]"),
),
(
np.arange("2019-01-01", "2019-01-06", dtype="datetime64[D]"),
np.arange("2018-01-01", "2018-01-03", dtype="datetime64[D]"),
),
],
)
def test_upcast_scalar_other(arr, other):
# for now we do not support non-scalar `other`
mask = np.array([False, True, False, True, True])
with pytest.raises(ValueError, match="other must be a scalar"):
maybe_upcast_putmask(arr, mask, other)
result = maybe_upcast_putmask(result, mask)


def test_upcast():
# GH23823
arr = np.arange(1, 6)
mask = np.array([False, True, False, True, True])
result, changed = maybe_upcast_putmask(arr, mask, other=np.nan)
result = maybe_upcast_putmask(arr, mask)

expected = np.array([1, np.nan, 3, np.nan, np.nan])
assert changed
tm.assert_numpy_array_equal(result, expected)


def test_upcast_datetime():
# GH23823
arr = np.arange("2019-01-01", "2019-01-06", dtype="datetime64[D]")
def test_maybe_upcast_putmask_bool():
# a case where maybe_upcast_putmask is *not* equivalent to
# try: np.putmask(result, mask, np.nan)
# except (ValueError, TypeError): result = np.where(mask, result, np.nan)
arr = np.array([True, False, True, False, True], dtype=bool)
mask = np.array([False, True, False, True, True])
result, changed = maybe_upcast_putmask(arr, mask, other=np.nan)
result = maybe_upcast_putmask(arr, mask)

expected = np.array(
[
"2019-01-01",
np.datetime64("NaT"),
"2019-01-03",
np.datetime64("NaT"),
np.datetime64("NaT"),
],
dtype="datetime64[D]",
)
assert not changed
expected = np.array([True, np.nan, True, np.nan, np.nan], dtype=object)
tm.assert_numpy_array_equal(result, expected)