REF: simplify maybe_upcast_putmask (#38487)

Author: jbrockmendel

pandas/core/dtypes/cast.py

@@ -418,9 +418,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.
 
@@ -434,69 +432,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):

pandas/core/ops/array_ops.py

@@ -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
 

pandas/tests/dtypes/cast/test_upcast.py

@@ -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)