BUG: Series.mask with small int dtypes raising (#45750)

Author: jbrockmendel

doc/source/whatsnew/v1.5.0.rst

@@ -271,6 +271,7 @@ Indexing
 - Bug in :meth:`loc.__setitem__` treating ``range`` keys as positional instead of label-based (:issue:`45479`)
 - Bug in :meth:`Series.__setitem__` when setting ``boolean`` dtype values containing ``NA`` incorrectly raising instead of casting to ``boolean`` dtype (:issue:`45462`)
 - Bug in :meth:`Series.__setitem__` where setting :attr:`NA` into a numeric-dtpye :class:`Series` would incorrectly upcast to object-dtype rather than treating the value as ``np.nan`` (:issue:`44199`)
+- Bug in :meth:`Series.mask` with ``inplace=True`` or setting values with a boolean mask with small integer dtypes incorrectly raising (:issue:`45750`)
 - Bug in :meth:`DataFrame.mask` with ``inplace=True`` and ``ExtensionDtype`` columns incorrectly raising (:issue:`45577`)
 - Bug in getting a column from a DataFrame with an object-dtype row index with datetime-like values: the resulting Series now preserves the exact object-dtype Index from the parent DataFrame (:issue:`42950`)
 - Bug in indexing on a :class:`DatetimeIndex` with a ``np.str_`` key incorrectly raising (:issue:`45580`)

pandas/core/array_algos/putmask.py

@@ -15,9 +15,7 @@
 from pandas.compat import np_version_under1p20
 
 from pandas.core.dtypes.cast import (
-    can_hold_element,
     convert_scalar_for_putitemlike,
-    find_common_type,
     infer_dtype_from,
 )
 from pandas.core.dtypes.common import is_list_like
@@ -61,59 +59,6 @@ def putmask_inplace(values: ArrayLike, mask: npt.NDArray[np.bool_], value: Any)
         np.putmask(values, mask, value)
 
 
-def putmask_smart(values: np.ndarray, mask: npt.NDArray[np.bool_], new) -> np.ndarray:
-    """
-    Return a new ndarray, try to preserve dtype if possible.
-
-    Parameters
-    ----------
-    values : np.ndarray
-        `values`, updated in-place.
-    mask : np.ndarray[bool]
-        Applies to both sides (array like).
-    new : listlike `new values` aligned with `values`
-
-    Returns
-    -------
-    values : ndarray with updated values
-        this *may* be a copy of the original
-
-    See Also
-    --------
-    np.putmask
-    """
-    # we cannot use np.asarray() here as we cannot have conversions
-    # that numpy does when numeric are mixed with strings
-
-    # see if we are only masking values that if putted
-    # will work in the current dtype
-    try:
-        nn = new[mask]
-    except TypeError:
-        # TypeError: only integer scalar arrays can be converted to a scalar index
-        pass
-    else:
-        # We only get to putmask_smart when we cannot hold 'new' in values.
-        #  The "smart" part of putmask_smart is checking if we can hold new[mask]
-        #  in values, in which case we can still avoid the need to cast.
-        if can_hold_element(values, nn):
-            values[mask] = nn
-            return values
-
-    new = np.asarray(new)
-
-    if values.dtype.kind == new.dtype.kind:
-        # preserves dtype if possible
-        np.putmask(values, mask, new)
-        return values
-
-    dtype = find_common_type([values.dtype, new.dtype])
-    values = values.astype(dtype)
-
-    np.putmask(values, mask, new)
-    return values
-
-
 def putmask_without_repeat(
     values: np.ndarray, mask: npt.NDArray[np.bool_], new: Any
 ) -> None:

pandas/core/dtypes/cast.py

@@ -1980,7 +1980,6 @@ def np_can_hold_element(dtype: np.dtype, element: Any) -> Any:
             if tipo.kind not in ["i", "u"]:
                 if isinstance(element, np.ndarray) and element.dtype.kind == "f":
                     # If all can be losslessly cast to integers, then we can hold them
-                    #  We do something similar in putmask_smart
                     casted = element.astype(dtype)
                     comp = casted == element
                     if comp.all():

pandas/core/internals/blocks.py

@@ -75,7 +75,6 @@
 from pandas.core.array_algos.putmask import (
     extract_bool_array,
     putmask_inplace,
-    putmask_smart,
     putmask_without_repeat,
     setitem_datetimelike_compat,
     validate_putmask,
@@ -978,15 +977,12 @@ def putmask(self, mask, new) -> list[Block]:
             # no need to split columns
 
             if not is_list_like(new):
-                # putmask_smart can't save us the need to cast
+                # using just new[indexer] can't save us the need to cast
                 return self.coerce_to_target_dtype(new).putmask(mask, new)
-
-            # This differs from
-            #  `self.coerce_to_target_dtype(new).putmask(mask, new)`
-            # because putmask_smart will check if new[mask] may be held
-            # by our dtype.
-            nv = putmask_smart(values.T, mask, new).T
-            return [self.make_block(nv)]
+            else:
+                indexer = mask.nonzero()[0]
+                nb = self.setitem(indexer, new[indexer])
+                return [nb]
 
         else:
             is_array = isinstance(new, np.ndarray)

pandas/tests/series/indexing/test_setitem.py

@@ -364,6 +364,48 @@ def test_setitem_nan_with_bool(self):
         expected = Series([np.nan, False, True], dtype=object)
         tm.assert_series_equal(result, expected)
 
+    def test_setitem_mask_smallint_upcast(self):
+        orig = Series([1, 2, 3], dtype="int8")
+        alt = np.array([999, 1000, 1001], dtype=np.int64)
+
+        mask = np.array([True, False, True])
+
+        ser = orig.copy()
+        ser[mask] = Series(alt)
+        expected = Series([999, 2, 1001])
+        tm.assert_series_equal(ser, expected)
+
+        ser2 = orig.copy()
+        ser2.mask(mask, alt, inplace=True)
+        tm.assert_series_equal(ser2, expected)
+
+        ser3 = orig.copy()
+        res = ser3.where(~mask, Series(alt))
+        tm.assert_series_equal(res, expected)
+
+    def test_setitem_mask_smallint_no_upcast(self):
+        # like test_setitem_mask_smallint_upcast, but while we can't hold 'alt',
+        #  we *can* hold alt[mask] without casting
+        orig = Series([1, 2, 3], dtype="uint8")
+        alt = Series([245, 1000, 246], dtype=np.int64)
+
+        mask = np.array([True, False, True])
+
+        ser = orig.copy()
+        ser[mask] = alt
+        expected = Series([245, 2, 246], dtype="uint8")
+        tm.assert_series_equal(ser, expected)
+
+        ser2 = orig.copy()
+        ser2.mask(mask, alt, inplace=True)
+        tm.assert_series_equal(ser2, expected)
+
+        # FIXME: don't leave commented-out
+        # FIXME: ser.where(~mask, alt) unnecessarily upcasts to int64
+        # ser3 = orig.copy()
+        # res = ser3.where(~mask, alt)
+        # tm.assert_series_equal(res, expected)
+
 
 class TestSetitemViewCopySemantics:
     def test_setitem_invalidates_datetime_index_freq(self):