BUG: NumericArray * td64_array (#45622)

Author: jbrockmendel

doc/source/whatsnew/v1.5.0.rst

@@ -234,6 +234,7 @@ Timezones
 Numeric
 ^^^^^^^
 - Bug in operations with array-likes with ``dtype="boolean"`` and :attr:`NA` incorrectly altering the array in-place (:issue:`45421`)
+- Bug in multiplying a :class:`Series` with ``IntegerDtype`` or ``FloatingDtype`` by an arraylike with ``timedelta64[ns]`` dtype incorrectly raising (:issue:`45622`)
 -
 
 Conversion

pandas/core/arrays/boolean.py

@@ -349,11 +349,10 @@ def _coerce_to_array(
     def _logical_method(self, other, op):
 
         assert op.__name__ in {"or_", "ror_", "and_", "rand_", "xor", "rxor"}
-        other_is_booleanarray = isinstance(other, BooleanArray)
         other_is_scalar = lib.is_scalar(other)
         mask = None
 
-        if other_is_booleanarray:
+        if isinstance(other, BooleanArray):
             other, mask = other._data, other._mask
         elif is_list_like(other):
             other = np.asarray(other, dtype="bool")
@@ -370,7 +369,7 @@ def _logical_method(self, other, op):
             )
 
         if not other_is_scalar and len(self) != len(other):
-            raise ValueError("Lengths must match to compare")
+            raise ValueError("Lengths must match")
 
         if op.__name__ in {"or_", "ror_"}:
             result, mask = ops.kleene_or(self._data, other, self._mask, mask)
@@ -387,7 +386,7 @@ def _arith_method(self, other, op):
         mask = None
         op_name = op.__name__
 
-        if isinstance(other, BooleanArray):
+        if isinstance(other, BaseMaskedArray):
             other, mask = other._data, other._mask
 
         elif is_list_like(other):
@@ -397,14 +396,7 @@ def _arith_method(self, other, op):
             if len(self) != len(other):
                 raise ValueError("Lengths must match")
 
-        # nans propagate
-        if mask is None:
-            mask = self._mask
-            if other is libmissing.NA:
-                # GH#45421 don't alter inplace
-                mask = mask | True
-        else:
-            mask = self._mask | mask
+        mask = self._propagate_mask(mask, other)
 
         if other is libmissing.NA:
             # if other is NA, the result will be all NA and we can't run the
@@ -425,14 +417,6 @@ def _arith_method(self, other, op):
             with np.errstate(all="ignore"):
                 result = op(self._data, other)
 
-        # divmod returns a tuple
-        if op_name == "divmod":
-            div, mod = result
-            return (
-                self._maybe_mask_result(div, mask, other, "floordiv"),
-                self._maybe_mask_result(mod, mask, other, "mod"),
-            )
-
         return self._maybe_mask_result(result, mask, other, op_name)
 
     def __abs__(self):

pandas/core/arrays/masked.py

@@ -13,7 +13,6 @@
 import numpy as np
 
 from pandas._libs import (
-    iNaT,
     lib,
     missing as libmissing,
 )
@@ -582,6 +581,18 @@ def _hasna(self) -> bool:
         # error: Incompatible return value type (got "bool_", expected "bool")
         return self._mask.any()  # type: ignore[return-value]
 
+    def _propagate_mask(
+        self, mask: npt.NDArray[np.bool_] | None, other
+    ) -> npt.NDArray[np.bool_]:
+        if mask is None:
+            mask = self._mask.copy()  # TODO: need test for BooleanArray needing a copy
+            if other is libmissing.NA:
+                # GH#45421 don't alter inplace
+                mask = mask | True
+        else:
+            mask = self._mask | mask
+        return mask
+
     def _cmp_method(self, other, op) -> BooleanArray:
         from pandas.core.arrays import BooleanArray
 
@@ -619,12 +630,7 @@ def _cmp_method(self, other, op) -> BooleanArray:
                 if result is NotImplemented:
                     result = invalid_comparison(self._data, other, op)
 
-        # nans propagate
-        if mask is None:
-            mask = self._mask.copy()
-        else:
-            mask = self._mask | mask
-
+        mask = self._propagate_mask(mask, other)
         return BooleanArray(result, mask, copy=False)
 
     def _maybe_mask_result(self, result, mask, other, op_name: str):
@@ -636,6 +642,14 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
         other : scalar or array-like
         op_name : str
         """
+        if op_name == "divmod":
+            # divmod returns a tuple
+            div, mod = result
+            return (
+                self._maybe_mask_result(div, mask, other, "floordiv"),
+                self._maybe_mask_result(mod, mask, other, "mod"),
+            )
+
         # if we have a float operand we are by-definition
         # a float result
         # or our op is a divide
@@ -657,8 +671,11 @@ def _maybe_mask_result(self, result, mask, other, op_name: str):
             # e.g. test_numeric_arr_mul_tdscalar_numexpr_path
             from pandas.core.arrays import TimedeltaArray
 
-            result[mask] = iNaT
-            return TimedeltaArray._simple_new(result)
+            if not isinstance(result, TimedeltaArray):
+                result = TimedeltaArray._simple_new(result)
+
+            result[mask] = result.dtype.type("NaT")
+            return result
 
         elif is_integer_dtype(result):
             from pandas.core.arrays import IntegerArray

pandas/core/arrays/numeric.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import datetime
 import numbers
 from typing import (
     TYPE_CHECKING,
@@ -10,7 +9,6 @@
 import numpy as np
 
 from pandas._libs import (
-    Timedelta,
     lib,
     missing as libmissing,
 )
@@ -23,20 +21,21 @@
 
 from pandas.core.dtypes.common import (
     is_bool_dtype,
-    is_float,
     is_float_dtype,
-    is_integer,
     is_integer_dtype,
     is_list_like,
     is_object_dtype,
     is_string_dtype,
     pandas_dtype,
 )
 
+from pandas.core import ops
+from pandas.core.arrays.base import ExtensionArray
 from pandas.core.arrays.masked import (
     BaseMaskedArray,
     BaseMaskedDtype,
 )
+from pandas.core.construction import ensure_wrapped_if_datetimelike
 
 if TYPE_CHECKING:
     import pyarrow
@@ -219,34 +218,40 @@ def _arith_method(self, other, op):
         op_name = op.__name__
         omask = None
 
-        if getattr(other, "ndim", 0) > 1:
-            raise NotImplementedError("can only perform ops with 1-d structures")
-
-        if isinstance(other, NumericArray):
+        if isinstance(other, BaseMaskedArray):
             other, omask = other._data, other._mask
 
         elif is_list_like(other):
-            other = np.asarray(other)
+            if not isinstance(other, ExtensionArray):
+                other = np.asarray(other)
             if other.ndim > 1:
                 raise NotImplementedError("can only perform ops with 1-d structures")
-            if len(self) != len(other):
-                raise ValueError("Lengths must match")
-            if not (is_float_dtype(other) or is_integer_dtype(other)):
-                raise TypeError("can only perform ops with numeric values")
 
-        elif isinstance(other, (datetime.timedelta, np.timedelta64)):
-            other = Timedelta(other)
+        # We wrap the non-masked arithmetic logic used for numpy dtypes
+        #  in Series/Index arithmetic ops.
+        other = ops.maybe_prepare_scalar_for_op(other, (len(self),))
+        pd_op = ops.get_array_op(op)
+        other = ensure_wrapped_if_datetimelike(other)
 
-        else:
-            if not (is_float(other) or is_integer(other) or other is libmissing.NA):
-                raise TypeError("can only perform ops with numeric values")
+        mask = self._propagate_mask(omask, other)
 
-        if omask is None:
-            mask = self._mask.copy()
-            if other is libmissing.NA:
-                mask |= True
+        if other is libmissing.NA:
+            result = np.ones_like(self._data)
+            if "truediv" in op_name and self.dtype.kind != "f":
+                # The actual data here doesn't matter since the mask
+                #  will be all-True, but since this is division, we want
+                #  to end up with floating dtype.
+                result = result.astype(np.float64)
         else:
-            mask = self._mask | omask
+            # Make sure we do this before the "pow" mask checks
+            #  to get an expected exception message on shape mismatch.
+            if self.dtype.kind in ["i", "u"] and op_name in ["floordiv", "mod"]:
+                # ATM we don't match the behavior of non-masked types with
+                #  respect to floordiv-by-zero
+                pd_op = op
+
+            with np.errstate(all="ignore"):
+                result = pd_op(self._data, other)
 
         if op_name == "pow":
             # 1 ** x is 1.
@@ -266,25 +271,6 @@ def _arith_method(self, other, op):
             # x ** 0 is 1.
             mask = np.where((self._data == 0) & ~self._mask, False, mask)
 
-        if other is libmissing.NA:
-            result = np.ones_like(self._data)
-            if "truediv" in op_name and self.dtype.kind != "f":
-                # The actual data here doesn't matter since the mask
-                #  will be all-True, but since this is division, we want
-                #  to end up with floating dtype.
-                result = result.astype(np.float64)
-        else:
-            with np.errstate(all="ignore"):
-                result = op(self._data, other)
-
-        # divmod returns a tuple
-        if op_name == "divmod":
-            div, mod = result
-            return (
-                self._maybe_mask_result(div, mask, other, "floordiv"),
-                self._maybe_mask_result(mod, mask, other, "mod"),
-            )
-
         return self._maybe_mask_result(result, mask, other, op_name)
 
     _HANDLED_TYPES = (np.ndarray, numbers.Number)

pandas/tests/arithmetic/test_datetime64.py

@@ -1097,6 +1097,7 @@ def test_dt64arr_addsub_intlike(
                 # IntegerArray
                 "can only perform ops with numeric values",
                 "unsupported operand type.*Categorical",
+                r"unsupported operand type\(s\) for -: 'int' and 'Timestamp'",
             ]
         )
         assert_invalid_addsub_type(obj, 1, msg)

pandas/tests/arrays/boolean/test_arithmetic.py

@@ -103,19 +103,24 @@ def test_error_invalid_values(data, all_arithmetic_operators):
     )
     with pytest.raises(TypeError, match=msg):
         ops("foo")
-    msg = (
-        r"unsupported operand type\(s\) for|"
-        "Concatenation operation is not implemented for NumPy arrays"
+    msg = "|".join(
+        [
+            r"unsupported operand type\(s\) for",
+            "Concatenation operation is not implemented for NumPy arrays",
+        ]
     )
     with pytest.raises(TypeError, match=msg):
         ops(pd.Timestamp("20180101"))
 
     # invalid array-likes
     if op not in ("__mul__", "__rmul__"):
         # TODO(extension) numpy's mul with object array sees booleans as numbers
-        msg = (
-            r"unsupported operand type\(s\) for|can only concatenate str|"
-            "not all arguments converted during string formatting"
+        msg = "|".join(
+            [
+                r"unsupported operand type\(s\) for",
+                "can only concatenate str",
+                "not all arguments converted during string formatting",
+            ]
         )
         with pytest.raises(TypeError, match=msg):
             ops(pd.Series("foo", index=s.index))

pandas/tests/arrays/boolean/test_logical.py

@@ -63,7 +63,7 @@ def test_eq_mismatched_type(self, other):
     def test_logical_length_mismatch_raises(self, all_logical_operators):
         op_name = all_logical_operators
         a = pd.array([True, False, None], dtype="boolean")
-        msg = "Lengths must match to compare"
+        msg = "Lengths must match"
 
         with pytest.raises(ValueError, match=msg):
             getattr(a, op_name)([True, False])

pandas/tests/arrays/floating/test_arithmetic.py

@@ -129,9 +129,19 @@ def test_error_invalid_values(data, all_arithmetic_operators):
     ops = getattr(s, op)
 
     # invalid scalars
-    msg = (
-        r"(:?can only perform ops with numeric values)"
-        r"|(:?FloatingArray cannot perform the operation mod)"
+    msg = "|".join(
+        [
+            r"can only perform ops with numeric values",
+            r"FloatingArray cannot perform the operation mod",
+            "unsupported operand type",
+            "not all arguments converted during string formatting",
+            "can't multiply sequence by non-int of type 'float'",
+            "ufunc 'subtract' cannot use operands with types dtype",
+            r"can only concatenate str \(not \"float\"\) to str",
+            "ufunc '.*' not supported for the input types, and the inputs could not",
+            "ufunc '.*' did not contain a loop with signature matching types",
+            "Concatenation operation is not implemented for NumPy arrays",
+        ]
     )
     with pytest.raises(TypeError, match=msg):
         ops("foo")
@@ -148,6 +158,13 @@ def test_error_invalid_values(data, all_arithmetic_operators):
             "cannot perform .* with this index type: DatetimeArray",
             "Addition/subtraction of integers and integer-arrays "
             "with DatetimeArray is no longer supported. *",
+            "unsupported operand type",
+            "not all arguments converted during string formatting",
+            "can't multiply sequence by non-int of type 'float'",
+            "ufunc 'subtract' cannot use operands with types dtype",
+            r"ufunc 'add' cannot use operands with types dtype\('<M8\[ns\]'\)",
+            r"ufunc 'add' cannot use operands with types dtype\('float\d{2}'\)",
+            "cannot subtract DatetimeArray from ndarray",
         ]
     )
     with pytest.raises(TypeError, match=msg):