Skip to content

Commit

Permalink
Masking and overflow checks for datetimeindex and timedeltaindex ops (p…
Browse files Browse the repository at this point in the history
…andas-dev#18020)

closes pandas-dev#17991

(cherry picked from commit 8388a47)
  • Loading branch information
jbrockmendel authored and TomAugspurger committed Dec 11, 2017
1 parent 8bede1a commit 8adc68f
Show file tree
Hide file tree
Showing 6 changed files with 64 additions and 4 deletions.
3 changes: 3 additions & 0 deletions doc/source/whatsnew/v0.21.1.txt
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,9 @@ Documentation Changes

Bug Fixes
~~~~~~~~~
- Bug in ``DataFrame.resample(...).apply(...)`` when there is a callable that returns different columns (:issue:`15169`)
- Bug in :class:`TimedeltaIndex` subtraction could incorrectly overflow when ``NaT`` is present (:issue:`17791`)
- Bug in :class:`DatetimeIndex` subtracting datetimelike from DatetimeIndex could fail to overflow (:issue:`18020`)

Conversion
^^^^^^^^^^
Expand Down
2 changes: 1 addition & 1 deletion pandas/core/indexes/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -681,7 +681,7 @@ def __sub__(self, other):
return self._add_delta(-other)
elif is_integer(other):
return self.shift(-other)
elif isinstance(other, datetime):
elif isinstance(other, (datetime, np.datetime64)):
return self._sub_datelike(other)
elif isinstance(other, Period):
return self._sub_period(other)
Expand Down
6 changes: 4 additions & 2 deletions pandas/core/indexes/datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
import pandas.core.dtypes.concat as _concat
from pandas.errors import PerformanceWarning
from pandas.core.common import _values_from_object, _maybe_box
from pandas.core.algorithms import checked_add_with_arr

from pandas.core.indexes.base import Index, _index_shared_docs
from pandas.core.indexes.numeric import Int64Index, Float64Index
Expand Down Expand Up @@ -762,7 +763,7 @@ def _sub_datelike(self, other):
raise TypeError("DatetimeIndex subtraction must have the same "
"timezones or no timezones")
result = self._sub_datelike_dti(other)
elif isinstance(other, datetime):
elif isinstance(other, (datetime, np.datetime64)):
other = Timestamp(other)
if other is libts.NaT:
result = self._nat_new(box=False)
Expand All @@ -772,7 +773,8 @@ def _sub_datelike(self, other):
"timezones or no timezones")
else:
i8 = self.asi8
result = i8 - other.value
result = checked_add_with_arr(i8, -other.value,
arr_mask=self._isnan)
result = self._maybe_mask_results(result,
fill_value=libts.iNaT)
else:
Expand Down
3 changes: 2 additions & 1 deletion pandas/core/indexes/timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -381,7 +381,8 @@ def _add_datelike(self, other):
else:
other = Timestamp(other)
i8 = self.asi8
result = checked_add_with_arr(i8, other.value)
result = checked_add_with_arr(i8, other.value,
arr_mask=self._isnan)
result = self._maybe_mask_results(result, fill_value=iNaT)
return DatetimeIndex(result, name=self.name, copy=False)

Expand Down
34 changes: 34 additions & 0 deletions pandas/tests/indexes/datetimes/test_datetime.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,40 @@ def test_ufunc_coercions(self):
tm.assert_index_equal(result, exp)
assert result.freq == 'D'

def test_datetimeindex_sub_timestamp_overflow(self):
dtimax = pd.to_datetime(['now', pd.Timestamp.max])
dtimin = pd.to_datetime(['now', pd.Timestamp.min])

tsneg = Timestamp('1950-01-01')
ts_neg_variants = [tsneg,
tsneg.to_pydatetime(),
tsneg.to_datetime64().astype('datetime64[ns]'),
tsneg.to_datetime64().astype('datetime64[D]')]

tspos = Timestamp('1980-01-01')
ts_pos_variants = [tspos,
tspos.to_pydatetime(),
tspos.to_datetime64().astype('datetime64[ns]'),
tspos.to_datetime64().astype('datetime64[D]')]

for variant in ts_neg_variants:
with pytest.raises(OverflowError):
dtimax - variant

expected = pd.Timestamp.max.value - tspos.value
for variant in ts_pos_variants:
res = dtimax - variant
assert res[1].value == expected

expected = pd.Timestamp.min.value - tsneg.value
for variant in ts_neg_variants:
res = dtimin - variant
assert res[1].value == expected

for variant in ts_pos_variants:
with pytest.raises(OverflowError):
dtimin - variant

def test_week_of_month_frequency(self):
# GH 5348: "ValueError: Could not evaluate WOM-1SUN" shouldn't raise
d1 = date(2002, 9, 1)
Expand Down
20 changes: 20 additions & 0 deletions pandas/tests/indexes/timedeltas/test_ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -1282,3 +1282,23 @@ def test_add_overflow(self):
result = (to_timedelta([pd.NaT, '5 days', '1 hours']) +
to_timedelta(['7 seconds', pd.NaT, '4 hours']))
tm.assert_index_equal(result, exp)

def test_timedeltaindex_add_timestamp_nat_masking(self):
# GH17991 checking for overflow-masking with NaT
tdinat = pd.to_timedelta(['24658 days 11:15:00', 'NaT'])

tsneg = Timestamp('1950-01-01')
ts_neg_variants = [tsneg,
tsneg.to_pydatetime(),
tsneg.to_datetime64().astype('datetime64[ns]'),
tsneg.to_datetime64().astype('datetime64[D]')]

tspos = Timestamp('1980-01-01')
ts_pos_variants = [tspos,
tspos.to_pydatetime(),
tspos.to_datetime64().astype('datetime64[ns]'),
tspos.to_datetime64().astype('datetime64[D]')]

for variant in ts_neg_variants + ts_pos_variants:
res = tdinat + variant
assert res[1] is pd.NaT

0 comments on commit 8adc68f

Please sign in to comment.