Merge branch 'main' into codex/11361-json-duplicate-policy

Author: wjxiz1992

ci/Jenkinsfile.premerge

@@ -315,7 +315,7 @@ git --no-pager diff --name-only HEAD \$BASE -- ${PREMERGE_DOCKERFILE} || true"""
 void runPremergeBuild() {
     common.prepareArtNetrc(this)
     withEnv([
-        "ARROW_THRIFT_MIRROR_URL=https://${env.ARTIFACTORY_NAME}/artifactory/apache-remote",
+        "ARROW_THRIFT_MIRROR_URL=https://${env.ARTIFACTORY_NAME}/artifactory/sw-spark-maven/apache-remote",
         "CMAKE_NETRC=OPTIONAL",
     ]) {
         sh 'source build/env.sh && ${sclCMD} "ci/premerge-build.sh"'

src/main/cpp/src/GpuTimeZoneDBJni.cpp

@@ -140,7 +140,8 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertOr
   jlong reader_tz_info_table,
   jint reader_tz_initial_offset,
   jint reader_tz_raw_offset,
-  jintArray reader_dst_rule)
+  jintArray reader_dst_rule,
+  jboolean writer_reader_rules_differ)
 {
   JNI_NULL_CHECK(env, input_handle, "input column is null", 0);
 
@@ -160,7 +161,38 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertOr
     auto const reader = spark_rapids_jni::orc_tz_side{
       reader_tz_info_tab, reader_tz_initial_offset, reader_tz_raw_offset, reader_dst};
     return cudf::jni::release_as_jlong(spark_rapids_jni::convert_orc_writer_reader_timezones(
-      *input, static_cast<int64_t>(writer_tz_offset_at_orc_2015_base_us), writer, reader));
+      *input,
+      static_cast<int64_t>(writer_tz_offset_at_orc_2015_base_us),
+      writer,
+      reader,
+      cudf::get_default_stream(),
+      cudf::get_current_device_resource_ref(),
+      writer_reader_rules_differ));
+  }
+  JNI_CATCH(env, 0);
+}
+
+JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertOrcFromUtcWithRules(
+  JNIEnv* env,
+  jclass,
+  jlong input_handle,
+  jlong reader_tz_info_table,
+  jint reader_tz_initial_offset,
+  jint reader_tz_raw_offset,
+  jintArray reader_dst_rule)
+{
+  JNI_NULL_CHECK(env, input_handle, "input column is null", 0);
+
+  JNI_TRY
+  {
+    cudf::jni::auto_set_device(env);
+    auto const input              = reinterpret_cast<cudf::column_view const*>(input_handle);
+    auto const reader_tz_info_tab = reinterpret_cast<cudf::table_view const*>(reader_tz_info_table);
+    auto const reader_dst         = parse_dst_rule(env, reader_dst_rule);
+    cudf::jni::check_java_exception(env);
+    auto const reader = spark_rapids_jni::orc_tz_side{
+      reader_tz_info_tab, reader_tz_initial_offset, reader_tz_raw_offset, reader_dst};
+    return cudf::jni::release_as_jlong(spark_rapids_jni::convert_orc_from_utc(*input, reader));
   }
   JNI_CATCH(env, 0);
 }

src/main/cpp/src/timezones.cu

@@ -582,11 +582,13 @@ __device__ static cudf::timestamp_us convert_timestamp_between_timezones(
   cudf::timestamp_us ts,
   orc_base_offset_info writer_2015_year_base_offset,
   tz_side_info const& writer,
-  tz_side_info const& reader)
+  tz_side_info const& reader,
+  bool writer_reader_rules_differ)
 {
   int64_t const decoded_us = static_cast<int64_t>(
     cuda::std::chrono::duration_cast<cudf::duration_us>(ts.time_since_epoch()).count());
   int64_t const adjusted_us = apply_orc_base_offset(decoded_us, writer_2015_year_base_offset);
+  if (!writer_reader_rules_differ) { return cudf::timestamp_us{cudf::duration_us{adjusted_us}}; }
 
   // Floor-divide to get epoch millis (handles negative timestamps correctly)
   int64_t const epoch_millis =
@@ -665,29 +667,32 @@ CUDF_KERNEL void __launch_bounds__(CONVERT_TZ_BLOCK_SIZE)
                            cudf::size_type input_offset,
                            orc_base_offset_info writer_2015_year_base_offset,
                            orc_tz_side_kernel_args writer_args,
-                           orc_tz_side_kernel_args reader_args)
+                           orc_tz_side_kernel_args reader_args,
+                           bool writer_reader_rules_differ)
 {
   // Shared memory layout: writer transitions, writer offsets, reader transitions, reader offsets
   extern __shared__ char smem[];
 
-  char* ptr = smem;
-  int64_t const *wt_begin, *wt_end, *rt_begin, *rt_end;
-  int32_t const *wo_begin, *ro_begin;
-  stage_side_transitions(writer_args.trans,
-                         writer_args.offsets,
-                         writer_args.trans_count,
-                         ptr,
-                         wt_begin,
-                         wt_end,
-                         wo_begin);
-  stage_side_transitions(reader_args.trans,
-                         reader_args.offsets,
-                         reader_args.trans_count,
-                         ptr,
-                         rt_begin,
-                         rt_end,
-                         ro_begin);
-  __syncthreads();
+  char* ptr               = smem;
+  int64_t const *wt_begin = nullptr, *wt_end = nullptr, *rt_begin = nullptr, *rt_end = nullptr;
+  int32_t const *wo_begin = nullptr, *ro_begin = nullptr;
+  if (writer_reader_rules_differ) {
+    stage_side_transitions(writer_args.trans,
+                           writer_args.offsets,
+                           writer_args.trans_count,
+                           ptr,
+                           wt_begin,
+                           wt_end,
+                           wo_begin);
+    stage_side_transitions(reader_args.trans,
+                           reader_args.offsets,
+                           reader_args.trans_count,
+                           ptr,
+                           rt_begin,
+                           rt_end,
+                           ro_begin);
+    __syncthreads();
+  }
 
   cudf::size_type idx = blockIdx.x * blockDim.x + threadIdx.x;
   if (idx < num_rows) {
@@ -706,8 +711,8 @@ CUDF_KERNEL void __launch_bounds__(CONVERT_TZ_BLOCK_SIZE)
                               reader_args.raw_offset,
                               reader_args.dst,
                               reader_args.is_fixed};
-    output[idx] =
-      convert_timestamp_between_timezones(input[idx], writer_2015_year_base_offset, writer, reader);
+    output[idx] = convert_timestamp_between_timezones(
+      input[idx], writer_2015_year_base_offset, writer, reader, writer_reader_rules_differ);
   }
 }
 
@@ -716,7 +721,8 @@ std::unique_ptr<column> convert_timezones(cudf::column_view const& input,
                                           spark_rapids_jni::orc_tz_side writer,
                                           spark_rapids_jni::orc_tz_side reader,
                                           rmm::cuda_stream_view stream,
-                                          rmm::device_async_resource_ref mr)
+                                          rmm::device_async_resource_ref mr,
+                                          bool writer_reader_rules_differ)
 {
   SRJ_FUNC_RANGE();
 
@@ -747,13 +753,15 @@ std::unique_ptr<column> convert_timezones(cudf::column_view const& input,
   int32_t reader_trans_count = reader.tz_info_table ? reader.tz_info_table->column(0).size() : 0;
 
   size_t smem_bytes = 0;
-  if (writer_trans_count > 0 && writer_trans_count <= MAX_SMEM_TRANSITIONS) {
-    smem_bytes += writer_trans_count * (sizeof(int64_t) + sizeof(int32_t));
-  }
-  if (reader_trans_count > 0 && reader_trans_count <= MAX_SMEM_TRANSITIONS) {
-    // Alignment padding between writer offsets (int32_t) and reader transitions (int64_t)
-    smem_bytes = align_up(smem_bytes, alignof(int64_t));
-    smem_bytes += reader_trans_count * (sizeof(int64_t) + sizeof(int32_t));
+  if (writer_reader_rules_differ) {
+    if (writer_trans_count > 0 && writer_trans_count <= MAX_SMEM_TRANSITIONS) {
+      smem_bytes += writer_trans_count * (sizeof(int64_t) + sizeof(int32_t));
+    }
+    if (reader_trans_count > 0 && reader_trans_count <= MAX_SMEM_TRANSITIONS) {
+      // Alignment padding between writer offsets (int32_t) and reader transitions (int64_t)
+      smem_bytes = align_up(smem_bytes, alignof(int64_t));
+      smem_bytes += reader_trans_count * (sizeof(int64_t) + sizeof(int32_t));
+    }
   }
 
   int32_t num_blocks = cudf::util::div_rounding_up_safe(input.size(), CONVERT_TZ_BLOCK_SIZE);
@@ -789,12 +797,124 @@ std::unique_ptr<column> convert_timezones(cudf::column_view const& input,
                input.offset(),
                writer_2015_year_base_offset,
                writer_args,
-               reader_args);
+               reader_args,
+               writer_reader_rules_differ);
   CUDF_CHECK_CUDA(stream.value());
 
   return results;
 }
 
+__device__ static int64_t wrapping_subtract(int64_t lhs, int64_t rhs)
+{
+  return static_cast<int64_t>(static_cast<uint64_t>(lhs) - static_cast<uint64_t>(rhs));
+}
+
+template <typename T>
+__device__ static T convert_orc_from_utc_value(T value, tz_side_info const& reader);
+
+template <>
+__device__ cudf::timestamp_us convert_orc_from_utc_value(cudf::timestamp_us value,
+                                                         tz_side_info const& reader)
+{
+  auto const value_us = value.time_since_epoch().count();
+  auto offset_ms      = reader.raw_offset;
+  if (!reader.is_fixed) {
+    auto const value_ms = spark_rapids_jni::integer_utils::floor_div(value_us, MICROS_PER_MILLI);
+    auto const offset_lookup_ms = wrapping_subtract(value_ms, reader.raw_offset);
+    offset_ms                   = get_transition_index(offset_lookup_ms, reader);
+  }
+  auto const result_us =
+    wrapping_subtract(value_us, static_cast<int64_t>(offset_ms) * MICROS_PER_MILLI);
+  return cudf::timestamp_us{cudf::duration_us{result_us}};
+}
+
+template <typename T>
+CUDF_KERNEL void __launch_bounds__(CONVERT_TZ_BLOCK_SIZE)
+  convert_orc_from_utc_kernel(T const* __restrict__ input,
+                              cudf::bitmask_type const* __restrict__ null_mask,
+                              T* __restrict__ output,
+                              cudf::size_type num_rows,
+                              cudf::size_type input_offset,
+                              orc_tz_side_kernel_args reader_args)
+{
+  extern __shared__ char smem[];
+  char* ptr               = smem;
+  int64_t const *rt_begin = nullptr, *rt_end = nullptr;
+  int32_t const* ro_begin = nullptr;
+  if (!reader_args.is_fixed) {
+    stage_side_transitions(reader_args.trans,
+                           reader_args.offsets,
+                           reader_args.trans_count,
+                           ptr,
+                           rt_begin,
+                           rt_end,
+                           ro_begin);
+    __syncthreads();
+  }
+
+  cudf::size_type idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < num_rows) {
+    if (null_mask && !cudf::bit_is_set(null_mask, idx + input_offset)) { return; }
+    tz_side_info const reader{rt_begin,
+                              rt_end,
+                              ro_begin,
+                              reader_args.initial_offset,
+                              reader_args.raw_offset,
+                              reader_args.dst,
+                              reader_args.is_fixed};
+    output[idx] = convert_orc_from_utc_value(input[idx], reader);
+  }
+}
+
+template <typename T>
+std::unique_ptr<column> convert_orc_from_utc_typed(cudf::column_view const& input,
+                                                   spark_rapids_jni::orc_tz_side reader,
+                                                   rmm::cuda_stream_view stream,
+                                                   rmm::device_async_resource_ref mr)
+{
+  auto results = cudf::make_fixed_width_column(input.type(),
+                                               input.size(),
+                                               cudf::copy_bitmask(input, stream, mr),
+                                               input.null_count(),
+                                               stream,
+                                               mr);
+  if (input.size() == 0) { return results; }
+
+  int64_t const* reader_trans_ptr =
+    reader.tz_info_table ? reader.tz_info_table->column(0).begin<int64_t>() : nullptr;
+  int32_t const* reader_offsets_ptr =
+    reader.tz_info_table ? reader.tz_info_table->column(1).begin<int32_t>() : nullptr;
+  int32_t reader_trans_count = reader.tz_info_table ? reader.tz_info_table->column(0).size() : 0;
+  auto const is_reader_fixed = reader_trans_count == 0 && !reader.dst.has_dst;
+  size_t smem_bytes          = 0;
+  if (reader_trans_count > 0 && reader_trans_count <= MAX_SMEM_TRANSITIONS) {
+    smem_bytes = reader_trans_count * (sizeof(int64_t) + sizeof(int32_t));
+  }
+
+  auto const reader_args   = orc_tz_side_kernel_args{reader_trans_ptr,
+                                                   reader_offsets_ptr,
+                                                   reader_trans_count,
+                                                   reader.initial_offset,
+                                                   reader.raw_offset,
+                                                   reader.dst,
+                                                   is_reader_fixed};
+  int32_t num_blocks       = cudf::util::div_rounding_up_safe(input.size(), CONVERT_TZ_BLOCK_SIZE);
+  auto const launch_config = cuda::make_config(cuda::grid_dims(num_blocks),
+                                               cuda::block_dims<CONVERT_TZ_BLOCK_SIZE>(),
+                                               cuda::dynamic_shared_memory<char[]>(smem_bytes));
+  cuda::launch(stream.value(),
+               launch_config,
+               convert_orc_from_utc_kernel<T>,
+               input.begin<T>(),
+               input.null_mask(),
+               results->mutable_view().begin<T>(),
+               input.size(),
+               input.offset(),
+               reader_args);
+  CUDF_CHECK_CUDA(stream.value());
+  return results;
+}
+
 // =================== ORC timezones end ===================
 
 }  // namespace
@@ -875,9 +995,22 @@ std::unique_ptr<cudf::column> convert_orc_writer_reader_timezones(
   orc_tz_side writer,
   orc_tz_side reader,
   rmm::cuda_stream_view stream,
-  rmm::device_async_resource_ref mr)
+  rmm::device_async_resource_ref mr,
+  bool writer_reader_rules_differ)
+{
+  return convert_timezones(
+    input, writer_2015_year_base_offset_us, writer, reader, stream, mr, writer_reader_rules_differ);
+}
+
+std::unique_ptr<cudf::column> convert_orc_from_utc(cudf::column_view const& input,
+                                                   orc_tz_side reader,
+                                                   rmm::cuda_stream_view stream,
+                                                   rmm::device_async_resource_ref mr)
 {
-  return convert_timezones(input, writer_2015_year_base_offset_us, writer, reader, stream, mr);
+  validate_timezone_table(reader.tz_info_table);
+  CUDF_EXPECTS(input.type().id() == cudf::type_id::TIMESTAMP_MICROSECONDS,
+               "ORC convertFromUtc input must be TIMESTAMP_MICROSECONDS");
+  return convert_orc_from_utc_typed<cudf::timestamp_us>(input, reader, stream, mr);
 }
 
 std::unique_ptr<cudf::column> convert_orc_writer_reader_timezones(

src/main/cpp/src/timezones.hpp

@@ -158,6 +158,8 @@ struct orc_tz_side {
  *        the same frame as Apache ORC. Pass 0 for no adjustment.
  * @param writer writer timezone transition data, offsets, and DST rule.
  * @param reader reader timezone transition data, offsets, and DST rule.
+ * @param writer_reader_rules_differ whether Apache ORC would call
+ *        SerializationUtils.convertBetweenTimezones for this timezone pair.
  * @param stream CUDA stream.
  * @param mr Device memory resource.
  * @return timestamps rebased between writer and reader timezones.
@@ -168,6 +170,23 @@ struct orc_tz_side {
   orc_tz_side writer,
   orc_tz_side reader,
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(),
+  bool writer_reader_rules_differ   = true);
+
+/**
+ * @brief Apply Apache ORC SerializationUtils.convertFromUtc semantics.
+ *
+ * @param input TIMESTAMP_MICROSECONDS input column.
+ * @param reader reader timezone transition data, offsets, and DST rule.
+ * @param stream CUDA stream.
+ * @param mr Device memory resource.
+ * @return converted column with the same type as input.
+ * @throws cudf::logic_error If input is not TIMESTAMP_MICROSECONDS.
+ */
+[[nodiscard]] std::unique_ptr<cudf::column> convert_orc_from_utc(
+  cudf::column_view const& input,
+  orc_tz_side reader,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 /**