Skip to content

API: Change DTA/TDA add/sub to match numpy #48748

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 6 commits into from
Sep 28, 2022
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
16 changes: 7 additions & 9 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -1134,13 +1134,12 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
return DatetimeArray._simple_new(result, dtype=result.dtype)

if self._reso != other._reso:
# Just as with Timestamp/Timedelta, we cast to the lower resolution
# so long as doing so is lossless.
# Just as with Timestamp/Timedelta, we cast to the higher resolution
if self._reso < other._reso:
other = other._as_unit(self._unit, round_ok=False)
else:
unit = npy_unit_to_abbrev(other._reso)
self = self._as_unit(unit)
else:
other = other._as_unit(self._unit)

i8 = self.asi8
result = checked_add_with_arr(i8, other.value, arr_mask=self._isnan)
Expand Down Expand Up @@ -1296,12 +1295,11 @@ def _add_timedelta_arraylike(
self = cast("DatetimeArray | TimedeltaArray", self)

if self._reso != other._reso:
# Just as with Timestamp/Timedelta, we cast to the lower resolution
# so long as doing so is lossless.
# Just as with Timestamp/Timedelta, we cast to the higher resolution
if self._reso < other._reso:
other = other._as_unit(self._unit)
else:
self = self._as_unit(other._unit)
else:
other = other._as_unit(self._unit)

self_i8 = self.asi8
other_i8 = other.asi8
Expand Down Expand Up @@ -2039,7 +2037,7 @@ def _unit(self) -> str:

def _as_unit(self: TimelikeOpsT, unit: str) -> TimelikeOpsT:
dtype = np.dtype(f"{self.dtype.kind}8[{unit}]")
new_values = astype_overflowsafe(self._ndarray, dtype, round_ok=False)
new_values = astype_overflowsafe(self._ndarray, dtype, round_ok=True)

if isinstance(self.dtype, np.dtype):
new_dtype = new_values.dtype
Expand Down
11 changes: 11 additions & 0 deletions pandas/tests/arrays/test_datetimes.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,17 @@ def test_compare_mismatched_resolutions(self, comparison_op):
np_res = op(left._ndarray, right._ndarray)
tm.assert_numpy_array_equal(np_res[1:], ~expected[1:])

def test_add_mismatched_reso_doesnt_downcast(self):
# https://github.com/pandas-dev/pandas/pull/48748#issuecomment-1260181008
td = pd.Timedelta(microseconds=1)
dti = pd.date_range("2016-01-01", periods=3) - td
dta = dti._data._as_unit("us")

res = dta + td._as_unit("us")
# even though the result is an even number of days
# (so we _could_ downcast to unit="s"), we do not.
assert res._unit == "us"


class TestDatetimeArrayComparisons:
# TODO: merge this into tests/arithmetic/test_datetime64 once it is
Expand Down
33 changes: 4 additions & 29 deletions pandas/tests/arrays/test_timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,22 +104,13 @@ def test_add_pdnat(self, tda):
def test_add_datetimelike_scalar(self, tda, tz_naive_fixture):
ts = pd.Timestamp("2016-01-01", tz=tz_naive_fixture)

expected = tda + ts._as_unit(tda._unit)
expected = tda._as_unit("ns") + ts
res = tda + ts
tm.assert_extension_array_equal(res, expected)
res = ts + tda
tm.assert_extension_array_equal(res, expected)

ts += Timedelta(1) # so we can't cast losslessly
msg = "Cannot losslessly convert units"
with pytest.raises(ValueError, match=msg):
# mismatched reso -> check that we don't give an incorrect result
tda + ts
with pytest.raises(ValueError, match=msg):
# mismatched reso -> check that we don't give an incorrect result
ts + tda

ts = ts._as_unit(tda._unit)
ts += Timedelta(1) # case where we can't cast losslessly

exp_values = tda._ndarray + ts.asm8
expected = (
Expand Down Expand Up @@ -185,35 +176,19 @@ def test_add_timedeltaarraylike(self, tda):
# TODO(2.0): just do `tda_nano = tda.astype("m8[ns]")`
tda_nano = TimedeltaArray(tda._ndarray.astype("m8[ns]"))

msg = "mis-matched resolutions is not yet supported"
expected = tda * 2
expected = tda_nano * 2
res = tda_nano + tda
tm.assert_extension_array_equal(res, expected)
res = tda + tda_nano
tm.assert_extension_array_equal(res, expected)

expected = tda * 0
expected = tda_nano * 0
res = tda - tda_nano
tm.assert_extension_array_equal(res, expected)

res = tda_nano - tda
tm.assert_extension_array_equal(res, expected)

tda_nano[:] = np.timedelta64(1, "ns") # can't round losslessly
msg = "Cannot losslessly cast '-?1 ns' to"
with pytest.raises(ValueError, match=msg):
tda_nano + tda
with pytest.raises(ValueError, match=msg):
tda + tda_nano
with pytest.raises(ValueError, match=msg):
tda - tda_nano
with pytest.raises(ValueError, match=msg):
tda_nano - tda

result = tda_nano + tda_nano
expected = tda_nano * 2
tm.assert_extension_array_equal(result, expected)


class TestTimedeltaArray:
@pytest.mark.parametrize("dtype", [int, np.int32, np.int64, "uint32", "uint64"])
Expand Down
7 changes: 7 additions & 0 deletions pandas/tests/scalar/timestamp/test_timestamp.py
Original file line number Diff line number Diff line change
Expand Up @@ -1030,6 +1030,13 @@ def test_sub_timedeltalike_mismatched_reso(self, ts_tz):
assert res == exp
assert res._reso == max(ts._reso, other._reso)

def test_addition_doesnt_downcast_reso(self):
# https://github.com/pandas-dev/pandas/pull/48748#pullrequestreview-1122635413
ts = Timestamp(year=2022, month=1, day=1, microsecond=999999)._as_unit("us")
td = Timedelta(microseconds=1)._as_unit("us")
res = ts + td
assert res._reso == ts._reso

def test_sub_timedelta64_mismatched_reso(self, ts_tz):
ts = ts_tz

Expand Down