upper_bound to lower_bound in timeadd to fix a bug

Signed-off-by: Haoyang Li <[email protected]>
NVIDIA · Feb 2, 2024 · f483608 · f483608
1 parent 6fd2fb7
commit f483608
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/src/main/cpp/src/timezones.cu b/src/main/cpp/src/timezones.cu
@@ -80,7 +80,7 @@ struct convert_timestamp_tz_functor {
       (to_utc ? tz_instants : utc_instants).data<int64_t>() + tz_transitions.element_offset(0),
       static_cast<size_t>(list_size));
 
-    auto const it = thrust::upper_bound(
+    auto const it = thrust::lower_bound(
       thrust::seq, transition_times.begin(), transition_times.end(), epoch_seconds);
     auto const idx         = static_cast<size_type>(thrust::distance(transition_times.begin(), it));
     auto const list_offset = tz_transitions.element_offset(idx - 1);
@@ -156,7 +156,7 @@ struct time_add_functor {
       cuda::std::chrono::duration_cast<cudf::duration_s>(timestamp.time_since_epoch()).count());
 
     // Binary search on local to find the correct offset to convert local to utc
-    auto const local_it = thrust::upper_bound(
+    auto const local_it = thrust::lower_bound(
       thrust::seq, transition_times_tz.begin(), transition_times_tz.end(), result_epoch_seconds);
     auto local_idx =
       static_cast<size_type>(thrust::distance(transition_times_tz.begin(), local_it));

diff --git a/src/test/java/com/nvidia/spark/rapids/jni/TimeZoneTest.java b/src/test/java/com/nvidia/spark/rapids/jni/TimeZoneTest.java
@@ -243,7 +243,8 @@ void timeAddCCTest() {
       0L,
       -2177481944610644L,
       -2177481944610644L,
-      -44623990055559136L);
+      -44623990055559136L,
+      -44623987543559136L);
       ColumnVector duration = ColumnVector.durationMicroSecondsFromBoxedLongs(
         56087020233685111L,
         1000000L,
@@ -254,6 +255,7 @@ void timeAddCCTest() {
         173001810506226873L,
         86399999999L,
         86400000000L,
+        43890193694846543L,
         43890193694846543L
       );
         ColumnVector expected = ColumnVector.timestampMicroSecondsFromBoxedLongs(
@@ -266,7 +268,8 @@ void timeAddCCTest() {
           173001810506226873L,
           -2177395544610645L,
           -2177395201610644L,
-          -733799617712593L);
+          -733799617712593L,
+          -733797105712593L);
         ColumnVector actual = GpuTimeZoneDB.timeAdd(input, duration,
           ZoneId.of("Asia/Shanghai"))) {
       assertColumnsAreEqual(expected, actual);