@@ -116,9 +116,8 @@ timedelta-like}
116
116
"""
117
117
cdef:
118
118
int64_t[::1 ] deltas
119
- int64_t[:] idx_shifted, idx_shifted_left, idx_shifted_right
120
119
ndarray[uint8_t, cast= True ] ambiguous_array, both_nat, both_eq
121
- Py_ssize_t i, idx, pos, ntrans, n = vals.shape[0 ]
120
+ Py_ssize_t i, isl, isr, idx, pos, ntrans, n = vals.shape[0 ]
122
121
Py_ssize_t delta_idx_offset, delta_idx, pos_left, pos_right
123
122
int64_t * tdata
124
123
int64_t v, left, right, val, v_left, v_right, new_local, remaining_mins
@@ -194,21 +193,28 @@ timedelta-like}
194
193
result_a[:] = NPY_NAT
195
194
result_b[:] = NPY_NAT
196
195
197
- idx_shifted_left = (np.maximum(0 , trans.searchsorted(
198
- vals - DAY_NANOS, side = ' right' ) - 1 )).astype(np.int64)
199
-
200
- idx_shifted_right = (np.maximum(0 , trans.searchsorted(
201
- vals + DAY_NANOS, side = ' right' ) - 1 )).astype(np.int64)
202
-
203
196
for i in range (n):
204
197
val = vals[i]
205
- v_left = val - deltas[idx_shifted_left[i]]
198
+ if val == NPY_NAT:
199
+ continue
200
+
201
+ # TODO: be careful of overflow in val-DAY_NANOS
202
+ isl = bisect_right_i8(tdata, val - DAY_NANOS, ntrans) - 1
203
+ if isl < 0 :
204
+ isl = 0
205
+
206
+ v_left = val - deltas[isl]
206
207
pos_left = bisect_right_i8(tdata, v_left, ntrans) - 1
207
208
# timestamp falls to the left side of the DST transition
208
209
if v_left + deltas[pos_left] == val:
209
210
result_a[i] = v_left
210
211
211
- v_right = val - deltas[idx_shifted_right[i]]
212
+ # TODO: be careful of overflow in val+DAY_NANOS
213
+ isr = bisect_right_i8(tdata, val + DAY_NANOS, ntrans) - 1
214
+ if isr < 0 :
215
+ isr = 0
216
+
217
+ v_right = val - deltas[isr]
212
218
pos_right = bisect_right_i8(tdata, v_right, ntrans) - 1
213
219
# timestamp falls to the right side of the DST transition
214
220
if v_right + deltas[pos_right] == val:
@@ -309,7 +315,9 @@ timedelta-like}
309
315
# Subtract 1 since the beginning hour is _inclusive_ of
310
316
# nonexistent times
311
317
new_local = val - remaining_mins - 1
312
- delta_idx = trans.searchsorted(new_local, side = ' right' )
318
+
319
+ delta_idx = bisect_right_i8(tdata, new_local, ntrans)
320
+
313
321
# Shift the delta_idx by if the UTC offset of
314
322
# the target tz is greater than 0 and we're moving forward
315
323
# or vice versa
@@ -333,17 +341,22 @@ timedelta-like}
333
341
334
342
cdef inline Py_ssize_t bisect_right_i8(int64_t * data,
335
343
int64_t val, Py_ssize_t n):
344
+ # Caller is responsible for checking n > 0
345
+ # This looks very similar to local_search_right in the ndarray.searchsorted
346
+ # implementation.
336
347
cdef:
337
348
Py_ssize_t pivot, left = 0 , right = n
338
349
339
- assert n >= 1
340
-
341
350
# edge cases
342
351
if val > data[n - 1 ]:
343
352
return n
344
353
345
- if val < data[0 ]:
346
- return 0
354
+ # Caller is responsible for ensuring 'val >= data[0]'. This is
355
+ # ensured by the fact that 'data' comes from get_dst_info where data[0]
356
+ # is *always* NPY_NAT+1. If that ever changes, we will need to restore
357
+ # the following disabled check.
358
+ # if val < data[0]:
359
+ # return 0
347
360
348
361
while left < right:
349
362
pivot = left + (right - left) // 2
@@ -403,6 +416,7 @@ cpdef int64_t tz_convert_from_utc_single(int64_t val, tzinfo tz):
403
416
int64_t delta
404
417
int64_t[::1 ] deltas
405
418
ndarray[int64_t, ndim= 1 ] trans
419
+ int64_t* tdata
406
420
intp_t pos
407
421
408
422
if val == NPY_NAT:
@@ -418,7 +432,8 @@ cpdef int64_t tz_convert_from_utc_single(int64_t val, tzinfo tz):
418
432
return val + delta
419
433
else :
420
434
trans, deltas, _ = get_dst_info(tz)
421
- pos = trans.searchsorted(val, side = " right" ) - 1
435
+ tdata = < int64_t* > cnp.PyArray_DATA(trans)
436
+ pos = bisect_right_i8(tdata, val, trans.shape[0 ]) - 1
422
437
return val + deltas[pos]
423
438
424
439
@@ -462,10 +477,11 @@ cdef const int64_t[:] _tz_convert_from_utc(const int64_t[:] vals, tzinfo tz):
462
477
"""
463
478
cdef:
464
479
int64_t[::1 ] converted, deltas
465
- Py_ssize_t i, n = vals.shape[0 ]
480
+ Py_ssize_t i, ntrans = - 1 , n = vals.shape[0 ]
466
481
int64_t val, delta = 0 # avoid not-initialized-warning
467
- intp_t[:] pos
482
+ intp_t pos
468
483
ndarray[int64_t] trans
484
+ int64_t* tdata = NULL
469
485
str typ
470
486
bint use_tzlocal = False , use_fixed = False , use_utc = True
471
487
@@ -479,13 +495,14 @@ cdef const int64_t[:] _tz_convert_from_utc(const int64_t[:] vals, tzinfo tz):
479
495
use_tzlocal = True
480
496
else :
481
497
trans, deltas, typ = get_dst_info(tz)
498
+ ntrans = trans.shape[0 ]
482
499
483
500
if typ not in [" pytz" , " dateutil" ]:
484
501
# FixedOffset, we know len(deltas) == 1
485
502
delta = deltas[0 ]
486
503
use_fixed = True
487
504
else :
488
- pos = trans.searchsorted(vals, side = " right " ) - 1
505
+ tdata = < int64_t * > cnp.PyArray_DATA(trans)
489
506
490
507
converted = np.empty(n, dtype = np.int64)
491
508
@@ -502,7 +519,8 @@ cdef const int64_t[:] _tz_convert_from_utc(const int64_t[:] vals, tzinfo tz):
502
519
elif use_fixed:
503
520
converted[i] = val + delta
504
521
else :
505
- converted[i] = val + deltas[pos[i]]
522
+ pos = bisect_right_i8(tdata, val, ntrans) - 1
523
+ converted[i] = val + deltas[pos]
506
524
507
525
return converted
508
526
0 commit comments