implement integer array add/sub for datetimelike indexes

Author: jbrockmendel

pandas/core/indexes/datetimelike.py

@@ -12,7 +12,7 @@
 
 import numpy as np
 
-from pandas._libs import lib, iNaT, NaT
+from pandas._libs import lib, iNaT, NaT, Timedelta
 from pandas._libs.tslibs.period import Period
 from pandas._libs.tslibs.timedeltas import delta_to_nanoseconds
 from pandas._libs.tslibs.timestamps import round_ns
@@ -34,6 +34,7 @@
     is_string_dtype,
     is_datetime64_dtype,
     is_datetime64tz_dtype,
+    is_datetime64_any_dtype,
     is_period_dtype,
     is_timedelta64_dtype)
 from pandas.core.dtypes.generic import (
@@ -48,6 +49,7 @@
 from pandas.util._decorators import Appender, cache_readonly
 import pandas.core.dtypes.concat as _concat
 import pandas.tseries.frequencies as frequencies
+from pandas.tseries.offsets import Tick
 
 import pandas.core.indexes.base as ibase
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
@@ -694,6 +696,47 @@ def _addsub_offset_array(self, other, op):
             kwargs['freq'] = 'infer'
         return self._constructor(res_values, **kwargs)
 
+    def _addsub_int_array(self, other, op):
+        """
+        Add or subtract array-like of integers equivalent to applying
+        `shift` pointwise.
+
+        Parameters
+        ----------
+        other : Index, np.ndarray
+            integer-dtype
+        op : {operator.add, operator.sub}
+
+        Returns
+        -------
+        result : same class as self
+        """
+        assert op in [operator.add, operator.sub]
+        if is_period_dtype(self):
+            # easy case for PeriodIndex
+            if op is operator.sub:
+                other = -other
+            res_values = checked_add_with_arr(self.asi8, other,
+                                              arr_mask=self._isnan)
+            res_values = res_values.view('i8')
+            res_values[self._isnan] = iNaT
+            return self._from_ordinals(res_values, freq=self.freq)
+
+        elif self.freq is None:
+            # GH#19123
+            raise NullFrequencyError("Cannot shift with no freq")
+
+        elif isinstance(self.freq, Tick):
+            # easy case where we can convert to timedelta64 operation
+            td = Timedelta(self.freq)
+            return op(self, td * other)
+
+        else:
+            # We should only get here with DatetimeIndex; dispatch
+            # to _addsub_offset_array
+            assert not is_timedelta64_dtype(self)
+            return op(self, np.array(other) * self.freq)
+
     @classmethod
     def _add_datetimelike_methods(cls):
         """
@@ -730,9 +773,8 @@ def __add__(self, other):
             elif is_datetime64_dtype(other) or is_datetime64tz_dtype(other):
                 # DatetimeIndex, ndarray[datetime64]
                 return self._add_datelike(other)
-            elif is_integer_dtype(other) and self.freq is None:
-                # GH#19123
-                raise NullFrequencyError("Cannot shift with no freq")
+            elif is_integer_dtype(other):
+                result = self._addsub_int_array(other, operator.add)
             else:  # pragma: no cover
                 return NotImplemented
 
@@ -783,13 +825,12 @@ def __sub__(self, other):
             elif is_datetime64_dtype(other) or is_datetime64tz_dtype(other):
                 # DatetimeIndex, ndarray[datetime64]
                 result = self._sub_datelike(other)
+            elif is_integer_dtype(other):
+                result = self._addsub_int_array(other, operator.sub)
             elif isinstance(other, Index):
                 raise TypeError("cannot subtract {cls} and {typ}"
                                 .format(cls=type(self).__name__,
                                         typ=type(other).__name__))
-            elif is_integer_dtype(other) and self.freq is None:
-                # GH#19123
-                raise NullFrequencyError("Cannot shift with no freq")
             else:  # pragma: no cover
                 return NotImplemented
 
@@ -810,6 +851,11 @@ def __rsub__(self, other):
                 # we need to wrap in DatetimeIndex and flip the operation
                 from pandas import DatetimeIndex
                 return DatetimeIndex(other) - self
+            elif (is_datetime64_any_dtype(self) and hasattr(other, 'dtype') and
+                  not is_datetime64_any_dtype(other)):
+                raise TypeError("cannot subtract {cls} from {typ}"
+                                .format(cls=type(self).__name__,
+                                        typ=type(other).__name__))
             return -(self - other)
         cls.__rsub__ = __rsub__
 

pandas/tests/indexes/datetimes/test_arithmetic.py

@@ -355,6 +355,46 @@ def test_dti_isub_int(self, tz, one):
         rng -= one
         tm.assert_index_equal(rng, expected)
 
+    # -------------------------------------------------------------
+    # __add__/__sub__ with integer arrays
+
+    @pytest.mark.parametrize('freq', ['H', 'D'])
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_dti_add_intarray_tick(self, box, freq):
+        dti = pd.date_range('2016-01-01', periods=2, freq=freq)
+        other = box([4, -1])
+        expected = DatetimeIndex([dti[n] + other[n] for n in range(len(dti))])
+        result = dti + other
+        tm.assert_index_equal(result, expected)
+        result = other + dti
+        tm.assert_index_equal(result, expected)
+
+    @pytest.mark.parametrize('freq', ['W', 'M', 'MS', 'Q'])
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_dti_add_intarray_non_tick(self, box, freq):
+        dti = pd.date_range('2016-01-01', periods=2, freq=freq)
+        other = box([4, -1])
+        expected = DatetimeIndex([dti[n] + other[n] for n in range(len(dti))])
+        with tm.assert_produces_warning(PerformanceWarning):
+            result = dti + other
+        tm.assert_index_equal(result, expected)
+        with tm.assert_produces_warning(PerformanceWarning):
+            result = other + dti
+        tm.assert_index_equal(result, expected)
+
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_dti_add_intarray_no_freq(self, box):
+        dti = pd.DatetimeIndex(['2016-01-01', 'NaT', '2017-04-05 06:07:08'])
+        other = box([9, 4, -1])
+        with pytest.raises(NullFrequencyError):
+            dti + other
+        with pytest.raises(NullFrequencyError):
+            other + dti
+        with pytest.raises(NullFrequencyError):
+            dti - other
+        with pytest.raises(TypeError):
+            other - dti
+
     # -------------------------------------------------------------
     # DatetimeIndex.shift is used in integer addition
 
@@ -516,7 +556,7 @@ def test_dti_sub_tdi(self, tz):
         result = dti - tdi.values
         tm.assert_index_equal(result, expected)
 
-        msg = 'cannot perform __neg__ with this index type:'
+        msg = 'cannot subtract DatetimeIndex from'
         with tm.assert_raises_regex(TypeError, msg):
             tdi.values - dti
 
@@ -541,7 +581,8 @@ def test_dti_isub_tdi(self, tz):
         tm.assert_index_equal(result, expected)
 
         msg = '|'.join(['cannot perform __neg__ with this index type:',
-                        'ufunc subtract cannot use operands with types'])
+                        'ufunc subtract cannot use operands with types',
+                        'cannot subtract DatetimeIndex from'])
         with tm.assert_raises_regex(TypeError, msg):
             tdi.values -= dti
 

pandas/tests/indexes/period/test_arithmetic.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 from datetime import timedelta
+import operator
+
 import pytest
 import numpy as np
 
@@ -9,6 +11,7 @@
                     period_range, Period, PeriodIndex,
                     _np_version_under1p10)
 import pandas.core.indexes.period as period
+from pandas.core import ops
 from pandas.errors import PerformanceWarning
 
 
@@ -434,6 +437,29 @@ def test_pi_sub_isub_offset(self):
         rng -= pd.offsets.MonthEnd(5)
         tm.assert_index_equal(rng, expected)
 
+    # ---------------------------------------------------------------
+    # __add__/__sub__ with integer arrays
+
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    @pytest.mark.parametrize('op', [operator.add, ops.radd])
+    def test_pi_add_intarray(self, box, op):
+        pi = pd.PeriodIndex([pd.Period('2015Q1'), pd.Period('NaT')])
+        other = box([4, -1])
+        result = op(pi, other)
+        expected = pd.PeriodIndex([pd.Period('2016Q1'), pd.Period('NaT')])
+        tm.assert_index_equal(result, expected)
+
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_pi_sub_intarray(self, box):
+        pi = pd.PeriodIndex([pd.Period('2015Q1'), pd.Period('NaT')])
+        other = box([4, -1])
+        result = pi - other
+        expected = pd.PeriodIndex([pd.Period('2014Q1'), pd.Period('NaT')])
+        tm.assert_index_equal(result, expected)
+
+        with pytest.raises(TypeError):
+            other - pi
+
     # ---------------------------------------------------------------
     # Timedelta-like (timedelta, timedelta64, Timedelta, Tick)
     # TODO: Some of these are misnomers because of non-Tick DateOffsets

pandas/tests/indexes/timedeltas/test_arithmetic.py

@@ -530,6 +530,42 @@ def test_tdi_isub_int(self, one):
         rng -= one
         tm.assert_index_equal(rng, expected)
 
+    # -------------------------------------------------------------
+    # __add__/__sub__ with integer arrays
+
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_tdi_add_integer_array(self, box):
+        rng = timedelta_range('1 days 09:00:00', freq='H', periods=3)
+        other = box([4, 3, 2])
+        expected = TimedeltaIndex(['1 day 13:00:00'] * 3)
+        result = rng + other
+        tm.assert_index_equal(result, expected)
+        result = other + rng
+        tm.assert_index_equal(result, expected)
+
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_tdi_sub_integer_array(self, box):
+        rng = timedelta_range('9H', freq='H', periods=3)
+        other = box([4, 3, 2])
+        expected = TimedeltaIndex(['5H', '7H', '9H'])
+        result = rng - other
+        tm.assert_index_equal(result, expected)
+        result = other - rng
+        tm.assert_index_equal(result, -expected)
+
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_tdi_addsub_integer_array_no_freq(self, box):
+        tdi = TimedeltaIndex(['1 Day', 'NaT', '3 Hours'])
+        other = box([14, -1, 16])
+        with pytest.raises(NullFrequencyError):
+            tdi + other
+        with pytest.raises(NullFrequencyError):
+            other + tdi
+        with pytest.raises(NullFrequencyError):
+            tdi - other
+        with pytest.raises(NullFrequencyError):
+            other - tdi
+
     # -------------------------------------------------------------
     # Binary operations TimedeltaIndex and timedelta-like