REF: do masking in checked_add_with_arr (#47396)

* REF: do masking in checked_add_with_arr * mypy fixup

Author: jbrockmendel

pandas/core/algorithms.py

@@ -1017,10 +1017,10 @@ def rank(
 
 def checked_add_with_arr(
  arr: npt.NDArray[np.int64],
- b,
+ b: int | npt.NDArray[np.int64],
  arr_mask: npt.NDArray[np.bool_] | None = None,
  b_mask: npt.NDArray[np.bool_] | None = None,
-) -> np.ndarray:
+) -> npt.NDArray[np.int64]:
  """
  Perform array addition that checks for underflow and overflow.
 
@@ -1098,7 +1098,12 @@ def checked_add_with_arr(
 
  if to_raise:
  raise OverflowError("Overflow in int64 addition")
- return arr + b
+
+ result = arr + b
+ if arr_mask is not None or b2_mask is not None:
+ np.putmask(result, ~not_nan, iNaT)
+
+ return result
 
 
 # --------------- #

pandas/core/arrays/datetimelike.py

@@ -1105,8 +1105,12 @@ def _add_datetimelike_scalar(self, other):
  return DatetimeArray(result)
 
  i8 = self.asi8
- result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
- result = self._maybe_mask_results(result)
+ # Incompatible types in assignment (expression has type "ndarray[Any,
+ # dtype[signedinteger[_64Bit]]]", variable has type
+ # "ndarray[Any, dtype[datetime64]]")
+ result = checked_add_with_arr( # type: ignore[assignment]
+ i8, other.value, arr_mask=self._isnan
+ )
  dtype = DatetimeTZDtype(tz=other.tz) if other.tz else DT64NS_DTYPE
  return DatetimeArray(result, dtype=dtype, freq=self.freq)
 
@@ -1147,7 +1151,6 @@ def _sub_datetimelike_scalar(self, other: datetime | np.datetime64):
 
  i8 = self.asi8
  result = checked_add_with_arr(i8, -other.value, arr_mask=self._isnan)
- result = self._maybe_mask_results(result)
  return result.view("timedelta64[ns]")
 
  @final
@@ -1169,31 +1172,30 @@ def _sub_datetime_arraylike(self, other):
 
  self_i8 = self.asi8
  other_i8 = other.asi8
- arr_mask = self._isnan | other._isnan
- new_values = checked_add_with_arr(self_i8, -other_i8, arr_mask=arr_mask)
- if self._hasna or other._hasna:
- np.putmask(new_values, arr_mask, iNaT)
+ new_values = checked_add_with_arr(
+ self_i8, -other_i8, arr_mask=self._isnan, b_mask=other._isnan
+ )
  return new_values.view("timedelta64[ns]")
 
  @final
- def _sub_period(self, other: Period):
+ def _sub_period(self, other: Period) -> npt.NDArray[np.object_]:
  if not is_period_dtype(self.dtype):
  raise TypeError(f"cannot subtract Period from a {type(self).__name__}")
 
  # If the operation is well-defined, we return an object-dtype ndarray
  # of DateOffsets. Null entries are filled with pd.NaT
  self._check_compatible_with(other)
  asi8 = self.asi8
- new_data = asi8 - other.ordinal
- new_data = np.array([self.freq.base * x for x in new_data])
+ new_i8_data = asi8 - other.ordinal # TODO: checked_add_with_arr
+ new_data = np.array([self.freq.base * x for x in new_i8_data])
 
  if self._hasna:
  new_data[self._isnan] = NaT
 
  return new_data
 
  @final
- def _add_period(self, other: Period):
+ def _add_period(self, other: Period) -> PeriodArray:
  if not is_timedelta64_dtype(self.dtype):
  raise TypeError(f"cannot add Period to a {type(self).__name__}")
 
@@ -1226,8 +1228,6 @@ def _add_timedeltalike_scalar(self, other):
  inc = delta_to_nanoseconds(other, reso=self._reso) # type: ignore[attr-defined]
 
  new_values = checked_add_with_arr(self.asi8, inc, arr_mask=self._isnan)
- new_values = new_values.view("i8")
- new_values = self._maybe_mask_results(new_values)
  new_values = new_values.view(self._ndarray.dtype)
 
  new_freq = None
@@ -1263,10 +1263,6 @@ def _add_timedelta_arraylike(
  new_values = checked_add_with_arr(
  self_i8, other_i8, arr_mask=self._isnan, b_mask=other._isnan
  )
- if self._hasna or other._hasna:
- mask = self._isnan | other._isnan
- np.putmask(new_values, mask, iNaT)
-
  return type(self)(new_values, dtype=self.dtype)
 
  @final
@@ -1310,11 +1306,11 @@ def _sub_period_array(self, other: PeriodArray) -> npt.NDArray[np.object_]:
  self = cast("PeriodArray", self)
  self._require_matching_freq(other)
 
- new_values = checked_add_with_arr(
+ new_i8_values = checked_add_with_arr(
  self.asi8, -other.asi8, arr_mask=self._isnan, b_mask=other._isnan
  )
 
- new_values = np.array([self.freq.base * x for x in new_values])
+ new_values = np.array([self.freq.base * x for x in new_i8_values])
  if self._hasna or other._hasna:
  mask = self._isnan | other._isnan
  new_values[mask] = NaT

pandas/core/arrays/period.py

@@ -734,8 +734,6 @@ def _addsub_int_array_or_scalar(
  if op is operator.sub:
  other = -other
  res_values = algos.checked_add_with_arr(self.asi8, other, arr_mask=self._isnan)
- res_values = res_values.view("i8")
- np.putmask(res_values, self._isnan, iNaT)
  return type(self)(res_values, freq=self.freq)
 
  def _add_offset(self, other: BaseOffset):