Skip to content
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

ENH: DTA/TDA add datetimelike scalar with mismatched reso #48669

Merged
merged 2 commits into from
Sep 21, 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
2 changes: 2 additions & 0 deletions pandas/_libs/tslibs/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,13 +31,15 @@
"periods_per_day",
"periods_per_second",
"is_supported_unit",
"npy_unit_to_abbrev",
]

from pandas._libs.tslibs import dtypes
from pandas._libs.tslibs.conversion import localize_pydatetime
from pandas._libs.tslibs.dtypes import (
Resolution,
is_supported_unit,
npy_unit_to_abbrev,
periods_per_day,
periods_per_second,
)
Expand Down
39 changes: 31 additions & 8 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,13 @@
Resolution,
Tick,
Timestamp,
astype_overflowsafe,
delta_to_nanoseconds,
get_unit_from_dtype,
iNaT,
ints_to_pydatetime,
ints_to_pytimedelta,
npy_unit_to_abbrev,
to_offset,
)
from pandas._libs.tslibs.fields import (
Expand Down Expand Up @@ -1130,10 +1132,13 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
return DatetimeArray._simple_new(result, dtype=result.dtype)

if self._reso != other._reso:
raise NotImplementedError(
"Addition between TimedeltaArray and Timestamp with mis-matched "
"resolutions is not yet supported."
)
# Just as with Timestamp/Timedelta, we cast to the lower resolution
# so long as doing so is lossless.
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)

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

if self._reso != other._reso:
raise NotImplementedError(
f"Addition of {type(self).__name__} with TimedeltaArray with "
"mis-matched resolutions is not yet supported."
)
# Just as with Timestamp/Timedelta, we cast to the lower resolution
# so long as doing so is lossless.
if self._reso < other._reso:
other = other._as_unit(self._unit)
else:
self = self._as_unit(other._unit)

self_i8 = self.asi8
other_i8 = other.asi8
Expand Down Expand Up @@ -2028,6 +2035,22 @@ def _unit(self) -> str:
# "ExtensionDtype"; expected "Union[DatetimeTZDtype, dtype[Any]]"
return dtype_to_unit(self.dtype) # type: ignore[arg-type]

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)

if isinstance(self.dtype, np.dtype):
new_dtype = new_values.dtype
else:
tz = cast("DatetimeArray", self).tz
new_dtype = DatetimeTZDtype(tz=tz, unit=unit)

# error: Unexpected keyword argument "freq" for "_simple_new" of
# "NDArrayBacked" [call-arg]
return type(self)._simple_new(
new_values, dtype=new_dtype, freq=self.freq # type: ignore[call-arg]
)

# --------------------------------------------------------------

def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
Expand Down
36 changes: 29 additions & 7 deletions pandas/tests/arrays/test_timedeltas.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,11 +104,18 @@ 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)

msg = "with mis-matched resolutions"
with pytest.raises(NotImplementedError, match=msg):
expected = tda + ts._as_unit(tda._unit)
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(NotImplementedError, match=msg):
with pytest.raises(ValueError, match=msg):
# mismatched reso -> check that we don't give an incorrect result
ts + tda

Expand Down Expand Up @@ -179,13 +186,28 @@ def test_add_timedeltaarraylike(self, tda):
tda_nano = TimedeltaArray(tda._ndarray.astype("m8[ns]"))

msg = "mis-matched resolutions is not yet supported"
with pytest.raises(NotImplementedError, match=msg):
expected = tda * 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
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(NotImplementedError, match=msg):
with pytest.raises(ValueError, match=msg):
tda + tda_nano
with pytest.raises(NotImplementedError, match=msg):
with pytest.raises(ValueError, match=msg):
tda - tda_nano
with pytest.raises(NotImplementedError, match=msg):
with pytest.raises(ValueError, match=msg):
tda_nano - tda

result = tda_nano + tda_nano
Expand Down
1 change: 1 addition & 0 deletions pandas/tests/tslibs/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def test_namespace():
"periods_per_day",
"periods_per_second",
"is_supported_unit",
"npy_unit_to_abbrev",
]

expected = set(submodules + api)
Expand Down