Skip to content

Commit a6c0ae4

Browse files
authored
PERF: datetimelike addition (#56373)
* PERF: datetimelike addition * re-add import * remove no-longer-used * mypy fixup * troubleshoot 32bit builds
1 parent d95a7a7 commit a6c0ae4

File tree

10 files changed

+59
-202
lines changed

10 files changed

+59
-202
lines changed

asv_bench/benchmarks/arithmetic.py

-37
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
date_range,
1313
to_timedelta,
1414
)
15-
from pandas.core.algorithms import checked_add_with_arr
1615

1716
from .pandas_vb_common import numeric_dtypes
1817

@@ -389,42 +388,6 @@ def time_add_timedeltas(self, df):
389388
df["timedelta"] + df["timedelta"]
390389

391390

392-
class AddOverflowScalar:
393-
params = [1, -1, 0]
394-
param_names = ["scalar"]
395-
396-
def setup(self, scalar):
397-
N = 10**6
398-
self.arr = np.arange(N)
399-
400-
def time_add_overflow_scalar(self, scalar):
401-
checked_add_with_arr(self.arr, scalar)
402-
403-
404-
class AddOverflowArray:
405-
def setup(self):
406-
N = 10**6
407-
self.arr = np.arange(N)
408-
self.arr_rev = np.arange(-N, 0)
409-
self.arr_mixed = np.array([1, -1]).repeat(N / 2)
410-
self.arr_nan_1 = np.random.choice([True, False], size=N)
411-
self.arr_nan_2 = np.random.choice([True, False], size=N)
412-
413-
def time_add_overflow_arr_rev(self):
414-
checked_add_with_arr(self.arr, self.arr_rev)
415-
416-
def time_add_overflow_arr_mask_nan(self):
417-
checked_add_with_arr(self.arr, self.arr_mixed, arr_mask=self.arr_nan_1)
418-
419-
def time_add_overflow_b_mask_nan(self):
420-
checked_add_with_arr(self.arr, self.arr_mixed, b_mask=self.arr_nan_1)
421-
422-
def time_add_overflow_both_arg_nan(self):
423-
checked_add_with_arr(
424-
self.arr, self.arr_mixed, arr_mask=self.arr_nan_1, b_mask=self.arr_nan_2
425-
)
426-
427-
428391
hcal = pd.tseries.holiday.USFederalHolidayCalendar()
429392
# These offsets currently raise a NotImplementedError with .apply_index()
430393
non_apply = [

pandas/_libs/tslibs/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
"npy_unit_to_abbrev",
3535
"get_supported_reso",
3636
"guess_datetime_format",
37+
"add_overflowsafe",
3738
]
3839

3940
from pandas._libs.tslibs import dtypes # pylint: disable=import-self
@@ -55,6 +56,7 @@
5556
from pandas._libs.tslibs.np_datetime import (
5657
OutOfBoundsDatetime,
5758
OutOfBoundsTimedelta,
59+
add_overflowsafe,
5860
astype_overflowsafe,
5961
is_unitless,
6062
py_get_unit_from_dtype as get_unit_from_dtype,

pandas/_libs/tslibs/np_datetime.pxd

+2
Original file line numberDiff line numberDiff line change
@@ -118,3 +118,5 @@ cdef int64_t convert_reso(
118118
NPY_DATETIMEUNIT to_reso,
119119
bint round_ok,
120120
) except? -1
121+
122+
cpdef cnp.ndarray add_overflowsafe(cnp.ndarray left, cnp.ndarray right)

pandas/_libs/tslibs/np_datetime.pyi

+4
Original file line numberDiff line numberDiff line change
@@ -19,3 +19,7 @@ def is_unitless(dtype: np.dtype) -> bool: ...
1919
def compare_mismatched_resolutions(
2020
left: np.ndarray, right: np.ndarray, op
2121
) -> npt.NDArray[np.bool_]: ...
22+
def add_overflowsafe(
23+
left: npt.NDArray[np.int64],
24+
right: npt.NDArray[np.int64],
25+
) -> npt.NDArray[np.int64]: ...

pandas/_libs/tslibs/np_datetime.pyx

+41
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
cimport cython
12
from cpython.datetime cimport (
23
PyDateTime_CheckExact,
34
PyDateTime_DATE_GET_HOUR,
@@ -678,3 +679,43 @@ cdef int64_t _convert_reso_with_dtstruct(
678679
raise OutOfBoundsDatetime from err
679680

680681
return result
682+
683+
684+
@cython.overflowcheck(True)
685+
cpdef cnp.ndarray add_overflowsafe(cnp.ndarray left, cnp.ndarray right):
686+
"""
687+
Overflow-safe addition for datetime64/timedelta64 dtypes.
688+
689+
`right` may either be zero-dim or of the same shape as `left`.
690+
"""
691+
cdef:
692+
Py_ssize_t N = left.size
693+
int64_t lval, rval, res_value
694+
ndarray iresult = cnp.PyArray_EMPTY(
695+
left.ndim, left.shape, cnp.NPY_INT64, 0
696+
)
697+
cnp.broadcast mi = cnp.PyArray_MultiIterNew3(iresult, left, right)
698+
699+
# Note: doing this try/except outside the loop improves performance over
700+
# doing it inside the loop.
701+
try:
702+
for i in range(N):
703+
# Analogous to: lval = lvalues[i]
704+
lval = (<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 1))[0]
705+
706+
# Analogous to: rval = rvalues[i]
707+
rval = (<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 2))[0]
708+
709+
if lval == NPY_DATETIME_NAT or rval == NPY_DATETIME_NAT:
710+
res_value = NPY_DATETIME_NAT
711+
else:
712+
res_value = lval + rval
713+
714+
# Analogous to: result[i] = res_value
715+
(<int64_t*>cnp.PyArray_MultiIter_DATA(mi, 0))[0] = res_value
716+
717+
cnp.PyArray_MultiIter_NEXT(mi)
718+
except OverflowError as err:
719+
raise OverflowError("Overflow in int64 addition") from err
720+
721+
return iresult

pandas/core/algorithms.py

-92
Original file line numberDiff line numberDiff line change
@@ -1119,98 +1119,6 @@ def rank(
11191119
return ranks
11201120

11211121

1122-
def checked_add_with_arr(
1123-
arr: npt.NDArray[np.int64],
1124-
b: int | npt.NDArray[np.int64],
1125-
arr_mask: npt.NDArray[np.bool_] | None = None,
1126-
b_mask: npt.NDArray[np.bool_] | None = None,
1127-
) -> npt.NDArray[np.int64]:
1128-
"""
1129-
Perform array addition that checks for underflow and overflow.
1130-
1131-
Performs the addition of an int64 array and an int64 integer (or array)
1132-
but checks that they do not result in overflow first. For elements that
1133-
are indicated to be NaN, whether or not there is overflow for that element
1134-
is automatically ignored.
1135-
1136-
Parameters
1137-
----------
1138-
arr : np.ndarray[int64] addend.
1139-
b : array or scalar addend.
1140-
arr_mask : np.ndarray[bool] or None, default None
1141-
array indicating which elements to exclude from checking
1142-
b_mask : np.ndarray[bool] or None, default None
1143-
array or scalar indicating which element(s) to exclude from checking
1144-
1145-
Returns
1146-
-------
1147-
sum : An array for elements x + b for each element x in arr if b is
1148-
a scalar or an array for elements x + y for each element pair
1149-
(x, y) in (arr, b).
1150-
1151-
Raises
1152-
------
1153-
OverflowError if any x + y exceeds the maximum or minimum int64 value.
1154-
"""
1155-
# For performance reasons, we broadcast 'b' to the new array 'b2'
1156-
# so that it has the same size as 'arr'.
1157-
b2 = np.broadcast_to(b, arr.shape)
1158-
if b_mask is not None:
1159-
# We do the same broadcasting for b_mask as well.
1160-
b2_mask = np.broadcast_to(b_mask, arr.shape)
1161-
else:
1162-
b2_mask = None
1163-
1164-
# For elements that are NaN, regardless of their value, we should
1165-
# ignore whether they overflow or not when doing the checked add.
1166-
if arr_mask is not None and b2_mask is not None:
1167-
not_nan = np.logical_not(arr_mask | b2_mask)
1168-
elif arr_mask is not None:
1169-
not_nan = np.logical_not(arr_mask)
1170-
elif b_mask is not None:
1171-
# error: Argument 1 to "__call__" of "_UFunc_Nin1_Nout1" has
1172-
# incompatible type "Optional[ndarray[Any, dtype[bool_]]]";
1173-
# expected "Union[_SupportsArray[dtype[Any]], _NestedSequence
1174-
# [_SupportsArray[dtype[Any]]], bool, int, float, complex, str
1175-
# , bytes, _NestedSequence[Union[bool, int, float, complex, str
1176-
# , bytes]]]"
1177-
not_nan = np.logical_not(b2_mask) # type: ignore[arg-type]
1178-
else:
1179-
not_nan = np.empty(arr.shape, dtype=bool)
1180-
not_nan.fill(True)
1181-
1182-
# gh-14324: For each element in 'arr' and its corresponding element
1183-
# in 'b2', we check the sign of the element in 'b2'. If it is positive,
1184-
# we then check whether its sum with the element in 'arr' exceeds
1185-
# np.iinfo(np.int64).max. If so, we have an overflow error. If it
1186-
# it is negative, we then check whether its sum with the element in
1187-
# 'arr' exceeds np.iinfo(np.int64).min. If so, we have an overflow
1188-
# error as well.
1189-
i8max = lib.i8max
1190-
i8min = iNaT
1191-
1192-
mask1 = b2 > 0
1193-
mask2 = b2 < 0
1194-
1195-
if not mask1.any():
1196-
to_raise = ((i8min - b2 > arr) & not_nan).any()
1197-
elif not mask2.any():
1198-
to_raise = ((i8max - b2 < arr) & not_nan).any()
1199-
else:
1200-
to_raise = ((i8max - b2[mask1] < arr[mask1]) & not_nan[mask1]).any() or (
1201-
(i8min - b2[mask2] > arr[mask2]) & not_nan[mask2]
1202-
).any()
1203-
1204-
if to_raise:
1205-
raise OverflowError("Overflow in int64 addition")
1206-
1207-
result = arr + b
1208-
if arr_mask is not None or b2_mask is not None:
1209-
np.putmask(result, ~not_nan, iNaT)
1210-
1211-
return result
1212-
1213-
12141122
# ---- #
12151123
# take #
12161124
# ---- #

pandas/core/arrays/datetimelike.py

+6-14
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
Tick,
3636
Timedelta,
3737
Timestamp,
38+
add_overflowsafe,
3839
astype_overflowsafe,
3940
get_unit_from_dtype,
4041
iNaT,
@@ -112,7 +113,6 @@
112113
ops,
113114
)
114115
from pandas.core.algorithms import (
115-
checked_add_with_arr,
116116
isin,
117117
map_array,
118118
unique1d,
@@ -1038,7 +1038,7 @@ def _get_i8_values_and_mask(
10381038
self, other
10391039
) -> tuple[int | npt.NDArray[np.int64], None | npt.NDArray[np.bool_]]:
10401040
"""
1041-
Get the int64 values and b_mask to pass to checked_add_with_arr.
1041+
Get the int64 values and b_mask to pass to add_overflowsafe.
10421042
"""
10431043
if isinstance(other, Period):
10441044
i8values = other.ordinal
@@ -1094,9 +1094,7 @@ def _add_datetimelike_scalar(self, other) -> DatetimeArray:
10941094
self = cast("TimedeltaArray", self)
10951095

10961096
other_i8, o_mask = self._get_i8_values_and_mask(other)
1097-
result = checked_add_with_arr(
1098-
self.asi8, other_i8, arr_mask=self._isnan, b_mask=o_mask
1099-
)
1097+
result = add_overflowsafe(self.asi8, np.asarray(other_i8, dtype="i8"))
11001098
res_values = result.view(f"M8[{self.unit}]")
11011099

11021100
dtype = tz_to_dtype(tz=other.tz, unit=self.unit)
@@ -1159,9 +1157,7 @@ def _sub_datetimelike(self, other: Timestamp | DatetimeArray) -> TimedeltaArray:
11591157
raise type(err)(new_message) from err
11601158

11611159
other_i8, o_mask = self._get_i8_values_and_mask(other)
1162-
res_values = checked_add_with_arr(
1163-
self.asi8, -other_i8, arr_mask=self._isnan, b_mask=o_mask
1164-
)
1160+
res_values = add_overflowsafe(self.asi8, np.asarray(-other_i8, dtype="i8"))
11651161
res_m8 = res_values.view(f"timedelta64[{self.unit}]")
11661162

11671163
new_freq = self._get_arithmetic_result_freq(other)
@@ -1227,9 +1223,7 @@ def _add_timedeltalike(self, other: Timedelta | TimedeltaArray):
12271223
self = cast("DatetimeArray | TimedeltaArray", self)
12281224

12291225
other_i8, o_mask = self._get_i8_values_and_mask(other)
1230-
new_values = checked_add_with_arr(
1231-
self.asi8, other_i8, arr_mask=self._isnan, b_mask=o_mask
1232-
)
1226+
new_values = add_overflowsafe(self.asi8, np.asarray(other_i8, dtype="i8"))
12331227
res_values = new_values.view(self._ndarray.dtype)
12341228

12351229
new_freq = self._get_arithmetic_result_freq(other)
@@ -1297,9 +1291,7 @@ def _sub_periodlike(self, other: Period | PeriodArray) -> npt.NDArray[np.object_
12971291
self._check_compatible_with(other)
12981292

12991293
other_i8, o_mask = self._get_i8_values_and_mask(other)
1300-
new_i8_data = checked_add_with_arr(
1301-
self.asi8, -other_i8, arr_mask=self._isnan, b_mask=o_mask
1302-
)
1294+
new_i8_data = add_overflowsafe(self.asi8, np.asarray(-other_i8, dtype="i8"))
13031295
new_data = np.array([self.freq.base * x for x in new_i8_data])
13041296

13051297
if o_mask is None:

pandas/core/arrays/period.py

+3-8
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
NaT,
2626
NaTType,
2727
Timedelta,
28+
add_overflowsafe,
2829
astype_overflowsafe,
2930
dt64arr_to_periodarr as c_dt64arr_to_periodarr,
3031
get_unit_from_dtype,
@@ -72,7 +73,6 @@
7273
)
7374
from pandas.core.dtypes.missing import isna
7475

75-
import pandas.core.algorithms as algos
7676
from pandas.core.arrays import datetimelike as dtl
7777
import pandas.core.common as com
7878

@@ -855,7 +855,7 @@ def _addsub_int_array_or_scalar(
855855
assert op in [operator.add, operator.sub]
856856
if op is operator.sub:
857857
other = -other
858-
res_values = algos.checked_add_with_arr(self.asi8, other, arr_mask=self._isnan)
858+
res_values = add_overflowsafe(self.asi8, np.asarray(other, dtype="i8"))
859859
return type(self)(res_values, dtype=self.dtype)
860860

861861
def _add_offset(self, other: BaseOffset):
@@ -920,12 +920,7 @@ def _add_timedelta_arraylike(
920920
"not an integer multiple of the PeriodArray's freq."
921921
) from err
922922

923-
b_mask = np.isnat(delta)
924-
925-
res_values = algos.checked_add_with_arr(
926-
self.asi8, delta.view("i8"), arr_mask=self._isnan, b_mask=b_mask
927-
)
928-
np.putmask(res_values, self._isnan | b_mask, iNaT)
923+
res_values = add_overflowsafe(self.asi8, np.asarray(delta.view("i8")))
929924
return type(self)(res_values, dtype=self.dtype)
930925

931926
def _check_timedeltalike_freq_compat(self, other):

pandas/tests/test_algos.py

-51
Original file line numberDiff line numberDiff line change
@@ -1837,57 +1837,6 @@ def test_pct_max_many_rows(self):
18371837
assert result == 1
18381838

18391839

1840-
def test_int64_add_overflow():
1841-
# see gh-14068
1842-
msg = "Overflow in int64 addition"
1843-
m = np.iinfo(np.int64).max
1844-
n = np.iinfo(np.int64).min
1845-
1846-
with pytest.raises(OverflowError, match=msg):
1847-
algos.checked_add_with_arr(np.array([m, m]), m)
1848-
with pytest.raises(OverflowError, match=msg):
1849-
algos.checked_add_with_arr(np.array([m, m]), np.array([m, m]))
1850-
with pytest.raises(OverflowError, match=msg):
1851-
algos.checked_add_with_arr(np.array([n, n]), n)
1852-
with pytest.raises(OverflowError, match=msg):
1853-
algos.checked_add_with_arr(np.array([n, n]), np.array([n, n]))
1854-
with pytest.raises(OverflowError, match=msg):
1855-
algos.checked_add_with_arr(np.array([m, n]), np.array([n, n]))
1856-
with pytest.raises(OverflowError, match=msg):
1857-
algos.checked_add_with_arr(
1858-
np.array([m, m]), np.array([m, m]), arr_mask=np.array([False, True])
1859-
)
1860-
with pytest.raises(OverflowError, match=msg):
1861-
algos.checked_add_with_arr(
1862-
np.array([m, m]), np.array([m, m]), b_mask=np.array([False, True])
1863-
)
1864-
with pytest.raises(OverflowError, match=msg):
1865-
algos.checked_add_with_arr(
1866-
np.array([m, m]),
1867-
np.array([m, m]),
1868-
arr_mask=np.array([False, True]),
1869-
b_mask=np.array([False, True]),
1870-
)
1871-
with pytest.raises(OverflowError, match=msg):
1872-
algos.checked_add_with_arr(np.array([m, m]), np.array([np.nan, m]))
1873-
1874-
# Check that the nan boolean arrays override whether or not
1875-
# the addition overflows. We don't check the result but just
1876-
# the fact that an OverflowError is not raised.
1877-
algos.checked_add_with_arr(
1878-
np.array([m, m]), np.array([m, m]), arr_mask=np.array([True, True])
1879-
)
1880-
algos.checked_add_with_arr(
1881-
np.array([m, m]), np.array([m, m]), b_mask=np.array([True, True])
1882-
)
1883-
algos.checked_add_with_arr(
1884-
np.array([m, m]),
1885-
np.array([m, m]),
1886-
arr_mask=np.array([True, False]),
1887-
b_mask=np.array([False, True]),
1888-
)
1889-
1890-
18911840
class TestMode:
18921841
def test_no_mode(self):
18931842
exp = Series([], dtype=np.float64, index=Index([], dtype=int))

0 commit comments

Comments
 (0)