BUG: Prevent addition overflow with TimedeltaIndex

Expands checked-add array addition introduced in
gh-14237 to include all other addition cases (i.e.
TimedeltaIndex and Timedelta). Follow-up to gh-14453.

Author: gfyoung

asv_bench/benchmarks/algorithms.py

@@ -24,6 +24,9 @@ def setup(self):
         self.arrneg = np.arange(-1000000, 0)
         self.arrmixed = np.array([1, -1]).repeat(500000)
 
+        self.arr_nan = np.random.choice([True, False], size=1000000)
+        self.arrmixed_nan = np.random.choice([True, False], size=1000000)
+
         # match
         self.uniques = tm.makeStringIndex(1000).values
         self.all = self.uniques.repeat(10)
@@ -64,6 +67,16 @@ def time_add_overflow_neg_arr(self):
     def time_add_overflow_mixed_arr(self):
         self.checked_add(self.arr, self.arrmixed)
 
+    def time_add_overflow_first_arg_nan(self):
+        self.checked_add(self.arr, self.arrmixed, arr_mask=self.arr_nan)
+
+    def time_add_overflow_second_arg_nan(self):
+        self.checked_add(self.arr, self.arrmixed, b_mask=self.arrmixed_arr_nan)
+
+    def time_add_overflow_both_arg_nan(self):
+        self.checked_add(self.arr, self.arrmixed, arr_mask=self.arr_nan,
+                         b_mask=self.arrmixed_arr_nan)
+
 
 class Hashing(object):
     goal_time = 0.2

pandas/core/nanops.py

@@ -812,15 +812,23 @@ def unique1d(values):
     return uniques
 
 
-def _checked_add_with_arr(arr, b):
+def _checked_add_with_arr(arr, b, arr_mask=None, b_mask=None):
     """
+    Perform array addition that checks for underflow and overflow.
+
     Performs the addition of an int64 array and an int64 integer (or array)
-    but checks that they do not result in overflow first.
+    but checks that they do not result in overflow first. For elements that
+    are indicated to be NaN, whether or not there is overflow for that element
+    is automatically ignored.
 
     Parameters
     ----------
     arr : array addend.
     b : array or scalar addend.
+    arr_mask : boolean array or None
+        array indicating which elements to exclude from checking
+    b_mask : boolean array or boolean or None
+        array or scalar indicating which element(s) to exclude from checking
 
     Returns
     -------
@@ -832,16 +840,40 @@ def _checked_add_with_arr(arr, b):
     ------
     OverflowError if any x + y exceeds the maximum or minimum int64 value.
     """
-    # For performance reasons, we broadcast 'b' to the new array 'b2'
-    # so that it has the same size as 'arr'.
-    if _np_version_under1p10:
-        if lib.isscalar(b):
-            b2 = np.empty(arr.shape)
-            b2.fill(b)
+    def _broadcast(arr_or_scalar, shape):
+        """
+        Helper function to broadcast arrays / scalars to the desired shape.
+
+        This function is compatible with different versions of NumPy and is
+        implemented for performance reasons.
+        """
+        if _np_version_under1p10:
+            if lib.isscalar(arr_or_scalar):
+                out = np.empty(shape)
+                out.fill(arr_or_scalar)
+            else:
+                out = arr_or_scalar
         else:
-            b2 = b
+            out = np.broadcast_to(arr_or_scalar, shape)
+        return out
+
+    b2 = _broadcast(b, arr.shape)
+    if b_mask is not None:
+        b2_mask = _broadcast(b_mask, arr.shape)
+    else:
+        b2_mask = None
+
+    # For elements that are NaN, regardless of their value, we should
+    # ignore whether they overflow or not when doing the checked add.
+    if arr_mask is not None and b2_mask is not None:
+        not_nan = np.logical_not(arr_mask | b2_mask)
+    elif arr_mask is not None:
+        not_nan = np.logical_not(arr_mask)
+    elif b_mask is not None:
+        not_nan = np.logical_not(b2_mask)
     else:
-        b2 = np.broadcast_to(b, arr.shape)
+        not_nan = np.empty(arr.shape, dtype=bool)
+        not_nan.fill(True)
 
     # gh-14324: For each element in 'arr' and its corresponding element
     # in 'b2', we check the sign of the element in 'b2'. If it is positive,
@@ -854,12 +886,14 @@ def _checked_add_with_arr(arr, b):
     mask2 = b2 < 0
 
     if not mask1.any():
-        to_raise = (np.iinfo(np.int64).min - b2 > arr).any()
+        to_raise = ((np.iinfo(np.int64).min - b2 > arr) & not_nan).any()
     elif not mask2.any():
-        to_raise = (np.iinfo(np.int64).max - b2 < arr).any()
+        to_raise = ((np.iinfo(np.int64).max - b2 < arr) & not_nan).any()
     else:
-        to_raise = ((np.iinfo(np.int64).max - b2[mask1] < arr[mask1]).any() or
-                    (np.iinfo(np.int64).min - b2[mask2] > arr[mask2]).any())
+        to_raise = (((np.iinfo(np.int64).max -
+                      b2[mask1] < arr[mask1]) & not_nan[mask1]).any() or
+                    ((np.iinfo(np.int64).min -
+                      b2[mask2] > arr[mask2]) & not_nan[mask2]).any())
 
     if to_raise:
         raise OverflowError("Overflow in int64 addition")

pandas/tests/test_nanops.py

@@ -1018,11 +1018,38 @@ def test_int64_add_overflow():
         nanops._checked_add_with_arr(np.array([n, n]), np.array([n, n]))
     with tm.assertRaisesRegexp(OverflowError, msg):
         nanops._checked_add_with_arr(np.array([m, n]), np.array([n, n]))
+    with tm.assertRaisesRegexp(OverflowError, msg):
+        nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]),
+                                     arr_mask=np.array([False, True]))
+    with tm.assertRaisesRegexp(OverflowError, msg):
+        nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]),
+                                     b_mask=np.array([False, True]))
+    with tm.assertRaisesRegexp(OverflowError, msg):
+        nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]),
+                                     arr_mask=np.array([False, True]),
+                                     b_mask=np.array([False, True]))
     with tm.assertRaisesRegexp(OverflowError, msg):
         with tm.assert_produces_warning(RuntimeWarning):
             nanops._checked_add_with_arr(np.array([m, m]),
                                          np.array([np.nan, m]))
 
+    # Check that the nan boolean arrays override whether or not
+    # the addition overflows. We don't check the result but just
+    # the fact that an OverflowError is not raised.
+    with tm.assertRaises(AssertionError):
+        with tm.assertRaisesRegexp(OverflowError, msg):
+            nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]),
+                                         arr_mask=np.array([True, True]))
+    with tm.assertRaises(AssertionError):
+        with tm.assertRaisesRegexp(OverflowError, msg):
+            nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]),
+                                         b_mask=np.array([True, True]))
+    with tm.assertRaises(AssertionError):
+        with tm.assertRaisesRegexp(OverflowError, msg):
+            nanops._checked_add_with_arr(np.array([m, m]), np.array([m, m]),
+                                         arr_mask=np.array([True, False]),
+                                         b_mask=np.array([False, True]))
+
 
 if __name__ == '__main__':
     import nose

pandas/tseries/base.py

@@ -17,6 +17,7 @@
 from pandas.types.missing import isnull
 from pandas.core import common as com, algorithms
 from pandas.core.common import AbstractMethodError
+from pandas.core.nanops import _checked_add_with_arr
 
 import pandas.formats.printing as printing
 import pandas.tslib as tslib
@@ -684,7 +685,8 @@ def _add_delta_td(self, other):
         # return the i8 result view
 
         inc = tslib._delta_to_nanoseconds(other)
-        new_values = (self.asi8 + inc).view('i8')
+        new_values = _checked_add_with_arr(self.asi8, inc,
+                                           arr_mask=self._isnan).view('i8')
         if self.hasnans:
             new_values[self._isnan] = tslib.iNaT
         return new_values.view('i8')
@@ -699,7 +701,9 @@ def _add_delta_tdi(self, other):
 
         self_i8 = self.asi8
         other_i8 = other.asi8
-        new_values = self_i8 + other_i8
+        new_values = _checked_add_with_arr(self_i8, other_i8,
+                                           arr_mask=self._isnan,
+                                           b_mask=other._isnan)
         if self.hasnans or other.hasnans:
             mask = (self._isnan) | (other._isnan)
             new_values[mask] = tslib.iNaT

pandas/tseries/tests/test_timedeltas.py

@@ -1964,6 +1964,20 @@ def test_add_overflow(self):
         with tm.assertRaisesRegexp(OverflowError, msg):
             Timestamp('2000') + to_timedelta([106580], 'D')
 
+        # These should not overflow!
+        exp = TimedeltaIndex([pd.NaT])
+        result = to_timedelta([pd.NaT]) - Timedelta('1 days')
+        tm.assert_index_equal(result, exp)
+
+        exp = TimedeltaIndex(['4 days', pd.NaT])
+        result = to_timedelta(['5 days', pd.NaT]) - Timedelta('1 days')
+        tm.assert_index_equal(result, exp)
+
+        exp = TimedeltaIndex([pd.NaT, pd.NaT, '5 hours'])
+        result = (to_timedelta([pd.NaT, '5 days', '1 hours']) +
+                  to_timedelta(['7 seconds', pd.NaT, '4 hours']))
+        tm.assert_index_equal(result, exp)
+
 if __name__ == '__main__':
     nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
                    exit=False)