Skip to content

Commit a9bfd55

Browse files
Chong Gaores-life
Chong Gao
authored and
res-life
committed
Fix ORC timezone conversion semantics
Preserve Apache ORC same-rule behavior and add ORC-compatible conversion from local timestamp values. Refs NVIDIA/cudf-spark#15449. Signed-off-by: Chong Gao <chongg@nvidia.com>
1 parent 52598d9 commit a9bfd55

4 files changed

Lines changed: 241 additions & 12 deletions

File tree

‎src/main/cpp/src/GpuTimeZoneDBJni.cpp‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,8 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertOr
140140
jlong reader_tz_info_table,
141141
jint reader_tz_initial_offset,
142142
jint reader_tz_raw_offset,
143-
jintArray reader_dst_rule)
143+
jintArray reader_dst_rule,
144+
jboolean writer_reader_rules_differ)
144145
{
145146
JNI_NULL_CHECK(env, input_handle, "input column is null", 0);
146147

@@ -160,7 +161,38 @@ JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertOr
160161
auto const reader = spark_rapids_jni::orc_tz_side{
161162
reader_tz_info_tab, reader_tz_initial_offset, reader_tz_raw_offset, reader_dst};
162163
return cudf::jni::release_as_jlong(spark_rapids_jni::convert_orc_writer_reader_timezones(
163-
*input, static_cast<int64_t>(writer_tz_offset_at_orc_2015_base_us), writer, reader));
164+
*input,
165+
static_cast<int64_t>(writer_tz_offset_at_orc_2015_base_us),
166+
writer,
167+
reader,
168+
cudf::get_default_stream(),
169+
cudf::get_current_device_resource_ref(),
170+
writer_reader_rules_differ));
171+
}
172+
JNI_CATCH(env, 0);
173+
}
174+
175+
JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertOrcFromUtcWithRules(
176+
JNIEnv* env,
177+
jclass,
178+
jlong input_handle,
179+
jlong reader_tz_info_table,
180+
jint reader_tz_initial_offset,
181+
jint reader_tz_raw_offset,
182+
jintArray reader_dst_rule)
183+
{
184+
JNI_NULL_CHECK(env, input_handle, "input column is null", 0);
185+
186+
JNI_TRY
187+
{
188+
cudf::jni::auto_set_device(env);
189+
auto const input = reinterpret_cast<cudf::column_view const*>(input_handle);
190+
auto const reader_tz_info_tab = reinterpret_cast<cudf::table_view const*>(reader_tz_info_table);
191+
auto const reader_dst = parse_dst_rule(env, reader_dst_rule);
192+
cudf::jni::check_java_exception(env);
193+
auto const reader = spark_rapids_jni::orc_tz_side{
194+
reader_tz_info_tab, reader_tz_initial_offset, reader_tz_raw_offset, reader_dst};
195+
return cudf::jni::release_as_jlong(spark_rapids_jni::convert_orc_from_utc(*input, reader));
164196
}
165197
JNI_CATCH(env, 0);
166198
}

‎src/main/cpp/src/timezones.cu‎

Lines changed: 129 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -582,11 +582,13 @@ __device__ static cudf::timestamp_us convert_timestamp_between_timezones(
582582
cudf::timestamp_us ts,
583583
orc_base_offset_info writer_2015_year_base_offset,
584584
tz_side_info const& writer,
585-
tz_side_info const& reader)
585+
tz_side_info const& reader,
586+
bool writer_reader_rules_differ)
586587
{
587588
int64_t const decoded_us = static_cast<int64_t>(
588589
cuda::std::chrono::duration_cast<cudf::duration_us>(ts.time_since_epoch()).count());
589590
int64_t const adjusted_us = apply_orc_base_offset(decoded_us, writer_2015_year_base_offset);
591+
if (!writer_reader_rules_differ) { return cudf::timestamp_us{cudf::duration_us{adjusted_us}}; }
590592

591593
// Floor-divide to get epoch millis (handles negative timestamps correctly)
592594
int64_t const epoch_millis =
@@ -665,7 +667,8 @@ CUDF_KERNEL void __launch_bounds__(CONVERT_TZ_BLOCK_SIZE)
665667
cudf::size_type input_offset,
666668
orc_base_offset_info writer_2015_year_base_offset,
667669
orc_tz_side_kernel_args writer_args,
668-
orc_tz_side_kernel_args reader_args)
670+
orc_tz_side_kernel_args reader_args,
671+
bool writer_reader_rules_differ)
669672
{
670673
// Shared memory layout: writer transitions, writer offsets, reader transitions, reader offsets
671674
extern __shared__ char smem[];
@@ -706,8 +709,8 @@ CUDF_KERNEL void __launch_bounds__(CONVERT_TZ_BLOCK_SIZE)
706709
reader_args.raw_offset,
707710
reader_args.dst,
708711
reader_args.is_fixed};
709-
output[idx] =
710-
convert_timestamp_between_timezones(input[idx], writer_2015_year_base_offset, writer, reader);
712+
output[idx] = convert_timestamp_between_timezones(
713+
input[idx], writer_2015_year_base_offset, writer, reader, writer_reader_rules_differ);
711714
}
712715
}
713716

@@ -716,7 +719,8 @@ std::unique_ptr<column> convert_timezones(cudf::column_view const& input,
716719
spark_rapids_jni::orc_tz_side writer,
717720
spark_rapids_jni::orc_tz_side reader,
718721
rmm::cuda_stream_view stream,
719-
rmm::device_async_resource_ref mr)
722+
rmm::device_async_resource_ref mr,
723+
bool writer_reader_rules_differ)
720724
{
721725
SRJ_FUNC_RANGE();
722726

@@ -789,12 +793,116 @@ std::unique_ptr<column> convert_timezones(cudf::column_view const& input,
789793
input.offset(),
790794
writer_2015_year_base_offset,
791795
writer_args,
792-
reader_args);
796+
reader_args,
797+
writer_reader_rules_differ);
793798
CUDF_CHECK_CUDA(stream.value());
794799

795800
return results;
796801
}
797802

803+
__device__ static int64_t wrapping_subtract(int64_t lhs, int64_t rhs)
804+
{
805+
return static_cast<int64_t>(static_cast<uint64_t>(lhs) - static_cast<uint64_t>(rhs));
806+
}
807+
808+
template <typename T>
809+
__device__ static T convert_orc_from_utc_value(T value, tz_side_info const& reader);
810+
811+
template <>
812+
__device__ cudf::timestamp_us convert_orc_from_utc_value(cudf::timestamp_us value,
813+
tz_side_info const& reader)
814+
{
815+
auto const value_us = value.time_since_epoch().count();
816+
auto const value_ms = spark_rapids_jni::integer_utils::floor_div(value_us, MICROS_PER_MILLI);
817+
auto const offset_lookup_ms = wrapping_subtract(value_ms, reader.raw_offset);
818+
auto const offset_ms = get_transition_index(offset_lookup_ms, reader);
819+
auto const result_us =
820+
wrapping_subtract(value_us, static_cast<int64_t>(offset_ms) * MICROS_PER_MILLI);
821+
return cudf::timestamp_us{cudf::duration_us{result_us}};
822+
}
823+
824+
template <typename T>
825+
CUDF_KERNEL void __launch_bounds__(CONVERT_TZ_BLOCK_SIZE)
826+
convert_orc_from_utc_kernel(T const* __restrict__ input,
827+
cudf::bitmask_type const* __restrict__ null_mask,
828+
T* __restrict__ output,
829+
cudf::size_type num_rows,
830+
cudf::size_type input_offset,
831+
orc_tz_side_kernel_args reader_args)
832+
{
833+
extern __shared__ char smem[];
834+
char* ptr = smem;
835+
int64_t const *rt_begin, *rt_end;
836+
int32_t const* ro_begin;
837+
stage_side_transitions(reader_args.trans,
838+
reader_args.offsets,
839+
reader_args.trans_count,
840+
ptr,
841+
rt_begin,
842+
rt_end,
843+
ro_begin);
844+
__syncthreads();
845+
846+
cudf::size_type idx = blockIdx.x * blockDim.x + threadIdx.x;
847+
if (idx < num_rows) {
848+
if (null_mask && !cudf::bit_is_set(null_mask, idx + input_offset)) { return; }
849+
tz_side_info const reader{rt_begin,
850+
rt_end,
851+
ro_begin,
852+
reader_args.initial_offset,
853+
reader_args.raw_offset,
854+
reader_args.dst};
855+
output[idx] = convert_orc_from_utc_value(input[idx], reader);
856+
}
857+
}
858+
859+
template <typename T>
860+
std::unique_ptr<column> convert_orc_from_utc_typed(cudf::column_view const& input,
861+
spark_rapids_jni::orc_tz_side reader,
862+
rmm::cuda_stream_view stream,
863+
rmm::device_async_resource_ref mr)
864+
{
865+
auto results = cudf::make_fixed_width_column(input.type(),
866+
input.size(),
867+
cudf::copy_bitmask(input, stream, mr),
868+
input.null_count(),
869+
stream,
870+
mr);
871+
if (input.size() == 0) { return results; }
872+
873+
int64_t const* reader_trans_ptr =
874+
reader.tz_info_table ? reader.tz_info_table->column(0).begin<int64_t>() : nullptr;
875+
int32_t const* reader_offsets_ptr =
876+
reader.tz_info_table ? reader.tz_info_table->column(1).begin<int32_t>() : nullptr;
877+
int32_t reader_trans_count = reader.tz_info_table ? reader.tz_info_table->column(0).size() : 0;
878+
size_t smem_bytes = 0;
879+
if (reader_trans_count > 0 && reader_trans_count <= MAX_SMEM_TRANSITIONS) {
880+
smem_bytes = reader_trans_count * (sizeof(int64_t) + sizeof(int32_t));
881+
}
882+
883+
auto const reader_args = orc_tz_side_kernel_args{reader_trans_ptr,
884+
reader_offsets_ptr,
885+
reader_trans_count,
886+
reader.initial_offset,
887+
reader.raw_offset,
888+
reader.dst};
889+
int32_t num_blocks = cudf::util::div_rounding_up_safe(input.size(), CONVERT_TZ_BLOCK_SIZE);
890+
auto const launch_config = cuda::make_config(cuda::grid_dims(num_blocks),
891+
cuda::block_dims<CONVERT_TZ_BLOCK_SIZE>(),
892+
cuda::dynamic_shared_memory<char[]>(smem_bytes));
893+
cuda::launch(stream.value(),
894+
launch_config,
895+
convert_orc_from_utc_kernel<T>,
896+
input.begin<T>(),
897+
input.null_mask(),
898+
results->mutable_view().begin<T>(),
899+
input.size(),
900+
input.offset(),
901+
reader_args);
902+
CUDF_CHECK_CUDA(stream.value());
903+
return results;
904+
}
905+
798906
// =================== ORC timezones end ===================
799907

800908
} // namespace
@@ -875,9 +983,22 @@ std::unique_ptr<cudf::column> convert_orc_writer_reader_timezones(
875983
orc_tz_side writer,
876984
orc_tz_side reader,
877985
rmm::cuda_stream_view stream,
878-
rmm::device_async_resource_ref mr)
986+
rmm::device_async_resource_ref mr,
987+
bool writer_reader_rules_differ)
988+
{
989+
return convert_timezones(
990+
input, writer_2015_year_base_offset_us, writer, reader, stream, mr, writer_reader_rules_differ);
991+
}
992+
993+
std::unique_ptr<cudf::column> convert_orc_from_utc(cudf::column_view const& input,
994+
orc_tz_side reader,
995+
rmm::cuda_stream_view stream,
996+
rmm::device_async_resource_ref mr)
879997
{
880-
return convert_timezones(input, writer_2015_year_base_offset_us, writer, reader, stream, mr);
998+
validate_timezone_table(reader.tz_info_table);
999+
CUDF_EXPECTS(input.type().id() == cudf::type_id::TIMESTAMP_MICROSECONDS,
1000+
"ORC convertFromUtc input must be TIMESTAMP_MICROSECONDS");
1001+
return convert_orc_from_utc_typed<cudf::timestamp_us>(input, reader, stream, mr);
8811002
}
8821003

8831004
std::unique_ptr<cudf::column> convert_orc_writer_reader_timezones(

‎src/main/cpp/src/timezones.hpp‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,8 @@ struct orc_tz_side {
158158
* the same frame as Apache ORC. Pass 0 for no adjustment.
159159
* @param writer writer timezone transition data, offsets, and DST rule.
160160
* @param reader reader timezone transition data, offsets, and DST rule.
161+
* @param writer_reader_rules_differ whether Apache ORC would call
162+
* SerializationUtils.convertBetweenTimezones for this timezone pair.
161163
* @param stream CUDA stream.
162164
* @param mr Device memory resource.
163165
* @return timestamps rebased between writer and reader timezones.
@@ -168,6 +170,22 @@ struct orc_tz_side {
168170
orc_tz_side writer,
169171
orc_tz_side reader,
170172
rmm::cuda_stream_view stream = cudf::get_default_stream(),
173+
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(),
174+
bool writer_reader_rules_differ = true);
175+
176+
/**
177+
* @brief Apply Apache ORC SerializationUtils.convertFromUtc semantics.
178+
*
179+
* @param input TIMESTAMP_MICROSECONDS input column.
180+
* @param reader reader timezone transition data, offsets, and DST rule.
181+
* @param stream CUDA stream.
182+
* @param mr Device memory resource.
183+
* @return converted column with the same type as input.
184+
*/
185+
[[nodiscard]] std::unique_ptr<cudf::column> convert_orc_from_utc(
186+
cudf::column_view const& input,
187+
orc_tz_side reader,
188+
rmm::cuda_stream_view stream = cudf::get_default_stream(),
171189
rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
172190

173191
/**

‎src/main/java/com/nvidia/spark/rapids/jni/GpuTimeZoneDB.java‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -581,10 +581,13 @@ public static final class OrcTimezoneContext implements AutoCloseable {
581581
private final int readerInitialOffset;
582582
private final int readerRawOffset;
583583
private final int[] readerDstRule;
584+
private final long readerFirstTransitionUs;
585+
private final boolean writerReaderRulesDiffer;
584586
private boolean closed;
585587

586588
private OrcTimezoneContext(Table writerTzInfoTable, Table readerTzInfoTable,
587-
String writerTimezone, OrcTimezoneInfo writerTzInfo, OrcTimezoneInfo readerTzInfo) {
589+
String writerTimezone, String readerTimezone,
590+
OrcTimezoneInfo writerTzInfo, OrcTimezoneInfo readerTzInfo) {
588591
this.writerTzInfoTable = writerTzInfoTable;
589592
this.readerTzInfoTable = readerTzInfoTable;
590593
this.writerTzOffsetAtOrc2015BaseUs = TimeUnit.MILLISECONDS.toMicros(
@@ -595,6 +598,18 @@ private OrcTimezoneContext(Table writerTzInfoTable, Table readerTzInfoTable,
595598
this.readerInitialOffset = readerTzInfo.initialOffset;
596599
this.readerRawOffset = readerTzInfo.rawOffset;
597600
this.readerDstRule = dstRuleToArray(readerTzInfo.dstRule);
601+
this.readerFirstTransitionUs =
602+
readerTzInfo.transitions == null || readerTzInfo.transitions.length == 0
603+
? Long.MIN_VALUE
604+
: TimeUnit.MILLISECONDS.toMicros(readerTzInfo.transitions[0]);
605+
TimeZone writerTz = TimeZone.getTimeZone(getZoneId(writerTimezone).getId());
606+
TimeZone readerTz = TimeZone.getTimeZone(getZoneId(readerTimezone).getId());
607+
this.writerReaderRulesDiffer = !writerTz.hasSameRules(readerTz);
608+
}
609+
610+
public long getReaderFirstTransitionUs() {
611+
ensureOpen();
612+
return readerFirstTransitionUs;
598613
}
599614

600615
private void ensureOpen() {
@@ -634,7 +649,7 @@ public static OrcTimezoneContext buildOrcTimezoneContext(
634649
writerTzInfoTable = getTableForUtilTZ(writerTzInfo);
635650
readerTzInfoTable = getTableForUtilTZ(readerTzInfo);
636651
return new OrcTimezoneContext(writerTzInfoTable, readerTzInfoTable,
637-
writerTimezone, writerTzInfo, readerTzInfo);
652+
writerTimezone, readerTimezone, writerTzInfo, readerTzInfo);
638653
} catch (RuntimeException | Error e) {
639654
try {
640655
Arms.closeAll(writerTzInfoTable, readerTzInfoTable);
@@ -667,9 +682,44 @@ public static ColumnVector convertOrcTimezones(
667682
context.readerTzInfoTable != null ? context.readerTzInfoTable.getNativeView() : 0L,
668683
context.readerInitialOffset,
669684
context.readerRawOffset,
685+
context.readerDstRule,
686+
context.writerReaderRulesDiffer));
687+
}
688+
689+
/**
690+
* Apply Apache ORC's {@code SerializationUtils.convertFromUtc} semantics using a pre-built
691+
* ORC timezone context. The input must be TIMESTAMP_MICROSECONDS.
692+
*
693+
* @param input values to convert
694+
* @param context timezone metadata whose reader side identifies the target timezone
695+
* @return converted values with the same type as {@code input}
696+
*/
697+
public static ColumnVector convertOrcFromUtc(
698+
ColumnView input, OrcTimezoneContext context) {
699+
context.ensureOpen();
700+
return new ColumnVector(convertOrcFromUtcWithRules(
701+
input.getNativeView(),
702+
context.readerTzInfoTable != null ? context.readerTzInfoTable.getNativeView() : 0L,
703+
context.readerInitialOffset,
704+
context.readerRawOffset,
670705
context.readerDstRule));
671706
}
672707

708+
/**
709+
* Apply Apache ORC's {@code SerializationUtils.convertFromUtc} semantics.
710+
*
711+
* @param input TIMESTAMP_MICROSECONDS values
712+
* @param readerTimezone target timezone
713+
* @return converted values with the same type as {@code input}
714+
*/
715+
public static ColumnVector convertOrcFromUtc(
716+
ColumnView input, String readerTimezone) {
717+
try (OrcTimezoneContext context =
718+
buildOrcTimezoneContext(readerTimezone, readerTimezone)) {
719+
return convertOrcFromUtc(input, context);
720+
}
721+
}
722+
673723
private static int[] dstRuleToArray(OrcDstRuleExtractor.DstRule rule) {
674724
if (rule == null) {
675725
return null;
@@ -737,5 +787,13 @@ private static native long convertOrcTimezonesWithRules(
737787
long readerTzInfoTable,
738788
int readerTzInitialOffset,
739789
int readerTzRawOffset,
790+
int[] readerDstRule,
791+
boolean writerReaderRulesDiffer);
792+
793+
private static native long convertOrcFromUtcWithRules(
794+
long input,
795+
long readerTzInfoTable,
796+
int readerTzInitialOffset,
797+
int readerTzRawOffset,
740798
int[] readerDstRule);
741799
}

0 commit comments

Comments
 (0)