implement add_offset_array for PeriodIndex (#19826)

Author: jbrockmendel

pandas/core/indexes/period.py

@@ -44,6 +44,7 @@
 from pandas.util._decorators import (Appender, Substitution, cache_readonly,
                                      deprecate_kwarg)
 from pandas.compat import zip, u
+from pandas.errors import PerformanceWarning
 
 import pandas.core.indexes.base as ibase
 _index_doc_kwargs = dict(ibase._index_doc_kwargs)
@@ -746,6 +747,28 @@ def _sub_period(self, other):
         # result must be Int64Index or Float64Index
         return Index(new_data)
 
+    def _add_offset_array(self, other):
+        # Array/Index of DateOffset objects
+        if len(other) == 1:
+            return self + other[0]
+        else:
+            warnings.warn("Adding/subtracting array of DateOffsets to "
+                          "{cls} not vectorized"
+                          .format(cls=type(self).__name__), PerformanceWarning)
+            res_values = self.astype('O').values + np.array(other)
+            return self.__class__(res_values)
+
+    def _sub_offset_array(self, other):
+        # Array/Index of DateOffset objects
+        if len(other) == 1:
+            return self - other[0]
+        else:
+            warnings.warn("Adding/subtracting array of DateOffsets to "
+                          "{cls} not vectorized"
+                          .format(cls=type(self).__name__), PerformanceWarning)
+            res_values = self.astype('O').values - np.array(other)
+            return self.__class__(res_values)
+
     def shift(self, n):
         """
         Specialized shift which produces an PeriodIndex

pandas/tests/indexes/period/test_arithmetic.py

@@ -9,6 +9,7 @@
                     period_range, Period, PeriodIndex,
                     _np_version_under1p10)
 import pandas.core.indexes.period as period
+from pandas.errors import PerformanceWarning
 
 
 _common_mismatch = [pd.offsets.YearBegin(2),
@@ -254,32 +255,57 @@ def test_comp_nat(self, dtype):
 
 
 class TestPeriodIndexArithmetic(object):
-    def test_pi_add_offset_array(self):
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_pi_add_offset_array(self, box):
         # GH#18849
         pi = pd.PeriodIndex([pd.Period('2015Q1'), pd.Period('2016Q2')])
-        offs = np.array([pd.offsets.QuarterEnd(n=1, startingMonth=12),
-                         pd.offsets.QuarterEnd(n=-2, startingMonth=12)])
-        res = pi + offs
+        offs = box([pd.offsets.QuarterEnd(n=1, startingMonth=12),
+                    pd.offsets.QuarterEnd(n=-2, startingMonth=12)])
         expected = pd.PeriodIndex([pd.Period('2015Q2'), pd.Period('2015Q4')])
+
+        with tm.assert_produces_warning(PerformanceWarning):
+            res = pi + offs
         tm.assert_index_equal(res, expected)
 
+        with tm.assert_produces_warning(PerformanceWarning):
+            res2 = offs + pi
+        tm.assert_index_equal(res2, expected)
+
         unanchored = np.array([pd.offsets.Hour(n=1),
                                pd.offsets.Minute(n=-2)])
+        # addition/subtraction ops with incompatible offsets should issue
+        # a PerformanceWarning and _then_ raise a TypeError.
         with pytest.raises(period.IncompatibleFrequency):
-            pi + unanchored
-        with pytest.raises(TypeError):
-            unanchored + pi
+            with tm.assert_produces_warning(PerformanceWarning):
+                pi + unanchored
+        with pytest.raises(period.IncompatibleFrequency):
+            with tm.assert_produces_warning(PerformanceWarning):
+                unanchored + pi
 
-    @pytest.mark.xfail(reason='GH#18824 radd doesnt implement this case')
-    def test_pi_radd_offset_array(self):
-        # GH#18849
+    @pytest.mark.parametrize('box', [np.array, pd.Index])
+    def test_pi_sub_offset_array(self, box):
+        # GH#18824
         pi = pd.PeriodIndex([pd.Period('2015Q1'), pd.Period('2016Q2')])
-        offs = np.array([pd.offsets.QuarterEnd(n=1, startingMonth=12),
-                         pd.offsets.QuarterEnd(n=-2, startingMonth=12)])
-        res = offs + pi
-        expected = pd.PeriodIndex([pd.Period('2015Q2'), pd.Period('2015Q4')])
+        other = box([pd.offsets.QuarterEnd(n=1, startingMonth=12),
+                     pd.offsets.QuarterEnd(n=-2, startingMonth=12)])
+
+        expected = PeriodIndex([pi[n] - other[n] for n in range(len(pi))])
+
+        with tm.assert_produces_warning(PerformanceWarning):
+            res = pi - other
         tm.assert_index_equal(res, expected)
 
+        anchored = box([pd.offsets.MonthEnd(), pd.offsets.Day(n=2)])
+
+        # addition/subtraction ops with anchored offsets should issue
+        # a PerformanceWarning and _then_ raise a TypeError.
+        with pytest.raises(period.IncompatibleFrequency):
+            with tm.assert_produces_warning(PerformanceWarning):
+                pi - anchored
+        with pytest.raises(period.IncompatibleFrequency):
+            with tm.assert_produces_warning(PerformanceWarning):
+                anchored - pi
+
     def test_pi_add_iadd_pi_raises(self):
         rng = pd.period_range('1/1/2000', freq='D', periods=5)
         other = pd.period_range('1/6/2000', freq='D', periods=5)