Skip to content

Implement add_offset_array for PeriodIndex #19826

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 22, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions pandas/core/indexes/period.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -746,6 +747,28 @@ def _sub_period(self, other):
# result must be Int64Index or Float64Index
return Index(new_data, name=self.name)

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
Expand Down
54 changes: 40 additions & 14 deletions pandas/tests/indexes/period/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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)
Expand Down