-
-
Notifications
You must be signed in to change notification settings - Fork 17.8k
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
Implement integer array add/sub for datetimelike indexes #19959
Changes from 14 commits
f1147af
39d8ee7
8780466
f4e8e01
70f359d
79e7575
ef6acdb
d477304
df84dc6
e4b6ec8
be737e0
650b1ef
671a008
bcf2419
8cc5270
6cb1b43
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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 ( | ||
|
@@ -814,6 +815,46 @@ 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) | ||
|
||
# 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): | ||
""" | ||
|
@@ -822,8 +863,6 @@ def _add_datetimelike_methods(cls): | |
""" | ||
|
||
def __add__(self, other): | ||
from pandas import DateOffset | ||
|
||
other = lib.item_from_zerodim(other) | ||
if isinstance(other, (ABCSeries, ABCDataFrame)): | ||
return NotImplemented | ||
|
@@ -853,9 +892,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) | ||
elif is_float_dtype(other): | ||
# Explicitly catch invalid dtypes | ||
raise TypeError("cannot add {dtype}-dtype to {cls}" | ||
|
@@ -915,14 +953,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") | ||
|
||
elif is_float_dtype(other): | ||
# Explicitly catch invalid dtypes | ||
raise TypeError("cannot subtract {dtype}-dtype from {cls}" | ||
|
@@ -948,6 +984,13 @@ 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)): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. can you add a comment on why this is There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Addressed, but GitHub UI not surfacing. |
||
# GH#19959 datetime - datetime is well-defined as timedelta, | ||
# but any other type - datetime is not well-defined. | ||
raise TypeError("cannot subtract {cls} from {typ}" | ||
.format(cls=type(self).__name__, | ||
typ=type(other).__name__)) | ||
return -(self - other) | ||
cls.__rsub__ = __rsub__ | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -367,6 +367,49 @@ 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): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. we must have some tests that do index arithmetic with integers, either move them here if they are not redundant or remove them. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There are tests for integers, but nothing ATM for integer-arrays. |
||
# GH#19959 | ||
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): | ||
# GH#19959 | ||
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): | ||
# GH#19959 | ||
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 | ||
|
||
|
@@ -528,7 +571,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 | ||
|
||
|
@@ -553,7 +596,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 | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -449,6 +449,31 @@ 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]) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. same comment as above |
||
@pytest.mark.parametrize('op', [operator.add, ops.radd]) | ||
def test_pi_add_intarray(self, box, op): | ||
# GH#19959 | ||
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): | ||
# GH#19959 | ||
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 | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
0.24 (this is an api change right?)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, just changed.