Skip to content

Commit 87d8b5e

Browse files
authored
REF: implement shift_bday (#34761)
1 parent 82bb42d commit 87d8b5e

File tree

1 file changed

+59
-19
lines changed

1 file changed

+59
-19
lines changed

pandas/_libs/tslibs/offsets.pyx

+59-19
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ from dateutil.easter import easter
1818

1919
import numpy as np
2020
cimport numpy as cnp
21-
from numpy cimport int64_t
21+
from numpy cimport int64_t, ndarray
2222
cnp.import_array()
2323

2424
# TODO: formalize having _libs.properties "above" tslibs in the dependency structure
@@ -1380,24 +1380,7 @@ cdef class BusinessDay(BusinessMixin):
13801380
@apply_index_wraps
13811381
def apply_index(self, dtindex):
13821382
i8other = dtindex.asi8
1383-
time = (i8other % DAY_NANOS).view("timedelta64[ns]")
1384-
1385-
# to_period rolls forward to next BDay; track and
1386-
# reduce n where it does when rolling forward
1387-
asper = dtindex.to_period("B")
1388-
1389-
if self.n > 0:
1390-
shifted = (dtindex.to_perioddelta("B") - time).asi8 != 0
1391-
1392-
roll = np.where(shifted, self.n - 1, self.n)
1393-
shifted = asper._addsub_int_array(roll, operator.add)
1394-
else:
1395-
# Integer addition is deprecated, so we use _time_shift directly
1396-
roll = self.n
1397-
shifted = asper._time_shift(roll)
1398-
1399-
result = shifted.to_timestamp() + time
1400-
return result
1383+
return shift_bdays(i8other, self.n)
14011384

14021385
def is_on_offset(self, dt) -> bool:
14031386
if self.normalize and not _is_normalized(dt):
@@ -3995,6 +3978,63 @@ def shift_months(const int64_t[:] dtindex, int months, object day_opt=None):
39953978
return np.asarray(out)
39963979

39973980

3981+
cdef ndarray[int64_t] shift_bdays(const int64_t[:] i8other, int periods):
3982+
"""
3983+
Implementation of BusinessDay.apply_offset.
3984+
3985+
Parameters
3986+
----------
3987+
i8other : const int64_t[:]
3988+
periods : int
3989+
3990+
Returns
3991+
-------
3992+
ndarray[int64_t]
3993+
"""
3994+
cdef:
3995+
Py_ssize_t i, n = len(i8other)
3996+
int64_t[:] result = np.empty(n, dtype="i8")
3997+
int64_t val, res
3998+
int wday, nadj, days
3999+
npy_datetimestruct dts
4000+
4001+
for i in range(n):
4002+
val = i8other[i]
4003+
if val == NPY_NAT:
4004+
result[i] = NPY_NAT
4005+
else:
4006+
# The rest of this is effectively a copy of BusinessDay.apply
4007+
nadj = periods
4008+
weeks = nadj // 5
4009+
dt64_to_dtstruct(val, &dts)
4010+
wday = dayofweek(dts.year, dts.month, dts.day)
4011+
4012+
if nadj <= 0 and wday > 4:
4013+
# roll forward
4014+
nadj += 1
4015+
4016+
nadj -= 5 * weeks
4017+
4018+
# nadj is always >= 0 at this point
4019+
if nadj == 0 and wday > 4:
4020+
# roll back
4021+
days = 4 - wday
4022+
elif wday > 4:
4023+
# roll forward
4024+
days = (7 - wday) + (nadj - 1)
4025+
elif wday + nadj <= 4:
4026+
# shift by n days without leaving the current week
4027+
days = nadj
4028+
else:
4029+
# shift by nadj days plus 2 to get past the weekend
4030+
days = nadj + 2
4031+
4032+
res = val + (7 * weeks + days) * DAY_NANOS
4033+
result[i] = res
4034+
4035+
return result.base
4036+
4037+
39984038
def shift_month(stamp: datetime, months: int,
39994039
day_opt: object=None) -> datetime:
40004040
"""

0 commit comments

Comments
 (0)