Skip to content

REF: use np.where instead of maybe_upcast_putmask in nanops #38130

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
Nov 29, 2020
Merged
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
19 changes: 10 additions & 9 deletions pandas/core/nanops.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
from pandas._typing import ArrayLike, Dtype, DtypeObj, F, Scalar
from pandas.compat._optional import import_optional_dependency

from pandas.core.dtypes.cast import maybe_upcast_putmask
from pandas.core.dtypes.common import (
get_dtype,
is_any_int_dtype,
Expand Down Expand Up @@ -284,18 +283,20 @@ def _get_values(
"""
# In _get_values is only called from within nanops, and in all cases
# with scalar fill_value. This guarantee is important for the
# maybe_upcast_putmask call below
# np.where call below
assert is_scalar(fill_value)
values = extract_array(values, extract_numpy=True)

mask = _maybe_get_mask(values, skipna, mask)

dtype = values.dtype

datetimelike = False
if needs_i8_conversion(values.dtype):
# changing timedelta64/datetime64 to int64 needs to happen after
# finding `mask` above
values = np.asarray(values.view("i8"))
datetimelike = True

dtype_ok = _na_ok_dtype(dtype)

Expand All @@ -306,13 +307,13 @@ def _get_values(
)

if skipna and (mask is not None) and (fill_value is not None):
values = values.copy()
if dtype_ok and mask.any():
np.putmask(values, mask, fill_value)

# promote if needed
else:
values, _ = maybe_upcast_putmask(values, mask, fill_value)
if mask.any():
if dtype_ok or datetimelike:
values = values.copy()
np.putmask(values, mask, fill_value)
else:
# np.where will promote if needed
values = np.where(~mask, values, fill_value)

# return a platform independent precision dtype
dtype_max = dtype
Expand Down