-
Notifications
You must be signed in to change notification settings - Fork 64
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Custom kernel for converting timestamps b/n UTC and non-UTC non-DST t…
…imezones (#1553) * Semi-working kernel for timestamp timezone conversion Signed-off-by: Navin Kumar <navink@nvidia.com> * Updated gtest with transition list * Refactor tests to use transitions as fixture * Add more items in the column to test each transition * Updated unit gtests for timezone kernel * Implementation of GpuTimeZoneDB with matching interface with CPU POC. * Add minimal convert from UTC test Signed-off-by: Navin Kumar <navink@nvidia.com> * Fix wrong offset bug in creating transition DB and update tests to sync with real timezone DB Signed-off-by: Navin Kumar <navink@nvidia.com> * Cleanup and sync test with CPP version. Signed-off-by: Navin Kumar <navink@nvidia.com> * Fix bug that happens when we pass a timestamp on the transition time exactly, switch to upper bound. Update tests for edge case. Signed-off-by: Navin Kumar <navink@nvidia.com> * Update timezone handling for convert to UTC and update tests Signed-off-by: Navin Kumar <navink@nvidia.com> * Internalize the daemon thread running to cache the timezone db Signed-off-by: Navin Kumar <navink@nvidia.com> * Fix null pointer exception by creating the instance automatically Signed-off-by: Navin Kumar <navink@nvidia.com> * Fix the visibility of these methods. Signed-off-by: Navin Kumar <navink@nvidia.com> * Add comment to note the type of the column vector stored in the database. Signed-off-by: Navin Kumar <navink@nvidia.com> * Remove the TIMESTAMP_DAYS code here. Signed-off-by: Navin Kumar <navink@nvidia.com> * Update this. I think the subtracting one second now doesn't make sense, because the transition would still happen on that exact time. Signed-off-by: Navin Kumar <navink@nvidia.com> * Update tests to handle around the instant of transition. Signed-off-by: Navin Kumar <navink@nvidia.com> * Add comments and add auto const Signed-off-by: Navin Kumar <navink@nvidia.com> * Run clang format and cleanup Signed-off-by: Navin Kumar <navink@nvidia.com> * Address some feedback: 1) Using auto const 2) removing counting iterator in thrust::transform 3) using thrust::distance 4) removing chrono type aliasing Signed-off-by: Navin Kumar <navink@nvidia.com> * Convert device method to functor to clean up calls to thrust::transform Signed-off-by: Navin Kumar <navink@nvidia.com> * Port Spark zoneId normalization to GpuTimeZoneDB Signed-off-by: Navin Kumar <navink@nvidia.com> * These need to be able to run independently of whether the timezone DB is loaded on the GPU Signed-off-by: Navin Kumar <navink@nvidia.com> * Fix logic of isSupportedTimeZone. Also lazy load GPU timezone database if needed (to handle config flag in plugin) Signed-off-by: Navin Kumar <navink@nvidia.com> * Prevent a double close of fixed transitions HCV Signed-off-by: Navin Kumar <navink@nvidia.com> * Address some feedback with some cleanup Signed-off-by: Navin Kumar <navink@nvidia.com> * Refactor into template function Signed-off-by: Navin Kumar <navink@nvidia.com> * Address more feedback by adding some aliases and cleanup Signed-off-by: Navin Kumar <navink@nvidia.com> * Header reordering by near-to-far * Address some Java feedback Signed-off-by: Navin Kumar <navink@nvidia.com> * Fix comment Signed-off-by: Navin Kumar <navink@nvidia.com> * fix formatting of JNI CPP file Signed-off-by: Navin Kumar <navink@nvidia.com> * Addressing more feedback, refactor again. Signed-off-by: Navin Kumar <navink@nvidia.com> * Pre-commit clang-format Signed-off-by: Navin Kumar <navink@nvidia.com> --------- Signed-off-by: Navin Kumar <navink@nvidia.com>
- Loading branch information
1 parent
edc0868
commit 0fa5796
Showing
8 changed files
with
1,103 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,53 @@ | ||
/* Copyright (c) 2023, NVIDIA CORPORATION. | ||
* | ||
* Licensed under the Apache License, Version 2.0 (the "License"); | ||
* you may not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
|
||
#include "cudf_jni_apis.hpp" | ||
#include "timezones.hpp" | ||
|
||
extern "C" { | ||
|
||
JNIEXPORT jlong JNICALL Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertTimestampColumnToUTC( | ||
JNIEnv* env, jclass, jlong input_handle, jlong transitions_handle, jint tz_index) | ||
{ | ||
JNI_NULL_CHECK(env, input_handle, "column is null", 0); | ||
JNI_NULL_CHECK(env, transitions_handle, "column is null", 0); | ||
try { | ||
cudf::jni::auto_set_device(env); | ||
auto const input = reinterpret_cast<cudf::column_view const*>(input_handle); | ||
auto const transitions = reinterpret_cast<cudf::table_view const*>(transitions_handle); | ||
auto const index = static_cast<cudf::size_type>(tz_index); | ||
return cudf::jni::ptr_as_jlong( | ||
spark_rapids_jni::convert_timestamp_to_utc(*input, *transitions, index).release()); | ||
} | ||
CATCH_STD(env, 0); | ||
} | ||
|
||
JNIEXPORT jlong JNICALL | ||
Java_com_nvidia_spark_rapids_jni_GpuTimeZoneDB_convertUTCTimestampColumnToTimeZone( | ||
JNIEnv* env, jclass, jlong input_handle, jlong transitions_handle, jint tz_index) | ||
{ | ||
JNI_NULL_CHECK(env, input_handle, "column is null", 0); | ||
JNI_NULL_CHECK(env, transitions_handle, "column is null", 0); | ||
try { | ||
cudf::jni::auto_set_device(env); | ||
auto const input = reinterpret_cast<cudf::column_view const*>(input_handle); | ||
auto const transitions = reinterpret_cast<cudf::table_view const*>(transitions_handle); | ||
auto const index = static_cast<cudf::size_type>(tz_index); | ||
return cudf::jni::ptr_as_jlong( | ||
spark_rapids_jni::convert_utc_timestamp_to_timezone(*input, *transitions, index).release()); | ||
} | ||
CATCH_STD(env, 0); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,166 @@ | ||
/* | ||
* Copyright (c) 2023, NVIDIA CORPORATION. | ||
* | ||
* Licensed under the Apache License, Version 2.0 (the "License"); | ||
* you may not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
|
||
#include "timezones.hpp" | ||
|
||
#include <cudf/column/column.hpp> | ||
#include <cudf/column/column_device_view.cuh> | ||
#include <cudf/column/column_factories.hpp> | ||
#include <cudf/detail/null_mask.hpp> | ||
#include <cudf/lists/list_device_view.cuh> | ||
#include <cudf/lists/lists_column_device_view.cuh> | ||
#include <cudf/table/table.hpp> | ||
#include <cudf/types.hpp> | ||
|
||
#include <thrust/binary_search.h> | ||
|
||
#include <rmm/cuda_stream_view.hpp> | ||
#include <rmm/exec_policy.hpp> | ||
|
||
using column = cudf::column; | ||
using column_device_view = cudf::column_device_view; | ||
using column_view = cudf::column_view; | ||
using lists_column_device_view = cudf::detail::lists_column_device_view; | ||
using size_type = cudf::size_type; | ||
using struct_view = cudf::struct_view; | ||
using table_view = cudf::table_view; | ||
|
||
namespace { | ||
|
||
// This device functor uses a binary search to find the instant of the transition | ||
// to find the right offset to do the transition. | ||
// To transition to UTC: do a binary search on the tzInstant child column and subtract | ||
// the offset. | ||
// To transition from UTC: do a binary search on the utcInstant child column and add | ||
// the offset. | ||
template <typename timestamp_type> | ||
struct convert_timestamp_tz_functor { | ||
using duration_type = typename timestamp_type::duration; | ||
|
||
// The list column of transitions to figure out the correct offset | ||
// to adjust the timestamp. The type of the values in this column is | ||
// LIST<STRUCT<utcInstant: int64, tzInstant: int64, utcOffset: int32>>. | ||
lists_column_device_view const transitions; | ||
// the index of the specified zone id in the transitions table | ||
size_type const tz_index; | ||
// whether we are converting to UTC or converting to the timezone | ||
bool const to_utc; | ||
|
||
/** | ||
* @brief Convert the timestamp value to either UTC or a specified timezone | ||
* @param timestamp input timestamp | ||
* | ||
*/ | ||
__device__ timestamp_type operator()(timestamp_type const& timestamp) const | ||
{ | ||
auto const utc_instants = transitions.child().child(0); | ||
auto const tz_instants = transitions.child().child(1); | ||
auto const utc_offsets = transitions.child().child(2); | ||
|
||
auto const epoch_seconds = static_cast<int64_t>( | ||
cuda::std::chrono::duration_cast<cudf::duration_s>(timestamp.time_since_epoch()).count()); | ||
auto const tz_transitions = cudf::list_device_view{transitions, tz_index}; | ||
auto const list_size = tz_transitions.size(); | ||
|
||
auto const transition_times = cudf::device_span<int64_t const>( | ||
(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( | ||
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); | ||
auto const utc_offset = cuda::std::chrono::duration_cast<duration_type>( | ||
cudf::duration_s{static_cast<int64_t>(utc_offsets.element<int32_t>(list_offset))}); | ||
return to_utc ? timestamp - utc_offset : timestamp + utc_offset; | ||
} | ||
}; | ||
|
||
template <typename timestamp_type> | ||
auto convert_timestamp_tz(column_view const& input, | ||
table_view const& transitions, | ||
size_type tz_index, | ||
bool to_utc, | ||
rmm::cuda_stream_view stream, | ||
rmm::mr::device_memory_resource* mr) | ||
{ | ||
// get the fixed transitions | ||
auto const ft_cdv_ptr = column_device_view::create(transitions.column(0), stream); | ||
auto const fixed_transitions = lists_column_device_view{*ft_cdv_ptr}; | ||
|
||
auto results = cudf::make_timestamp_column(input.type(), | ||
input.size(), | ||
cudf::detail::copy_bitmask(input, stream, mr), | ||
input.null_count(), | ||
stream, | ||
mr); | ||
|
||
thrust::transform( | ||
rmm::exec_policy(stream), | ||
input.begin<timestamp_type>(), | ||
input.end<timestamp_type>(), | ||
results->mutable_view().begin<timestamp_type>(), | ||
convert_timestamp_tz_functor<timestamp_type>{fixed_transitions, tz_index, to_utc}); | ||
|
||
return results; | ||
} | ||
|
||
} // namespace | ||
|
||
namespace spark_rapids_jni { | ||
|
||
std::unique_ptr<column> convert_timestamp(column_view const& input, | ||
table_view const& transitions, | ||
size_type tz_index, | ||
bool to_utc, | ||
rmm::cuda_stream_view stream, | ||
rmm::mr::device_memory_resource* mr) | ||
{ | ||
auto const type = input.type().id(); | ||
|
||
switch (type) { | ||
case cudf::type_id::TIMESTAMP_SECONDS: | ||
return convert_timestamp_tz<cudf::timestamp_s>( | ||
input, transitions, tz_index, to_utc, stream, mr); | ||
case cudf::type_id::TIMESTAMP_MILLISECONDS: | ||
return convert_timestamp_tz<cudf::timestamp_ms>( | ||
input, transitions, tz_index, to_utc, stream, mr); | ||
case cudf::type_id::TIMESTAMP_MICROSECONDS: | ||
return convert_timestamp_tz<cudf::timestamp_us>( | ||
input, transitions, tz_index, to_utc, stream, mr); | ||
default: CUDF_FAIL("Unsupported timestamp unit for timezone conversion"); | ||
} | ||
} | ||
|
||
std::unique_ptr<column> convert_timestamp_to_utc(column_view const& input, | ||
table_view const& transitions, | ||
size_type tz_index, | ||
rmm::cuda_stream_view stream, | ||
rmm::mr::device_memory_resource* mr) | ||
{ | ||
return convert_timestamp(input, transitions, tz_index, true, stream, mr); | ||
} | ||
|
||
std::unique_ptr<column> convert_utc_timestamp_to_timezone(column_view const& input, | ||
table_view const& transitions, | ||
size_type tz_index, | ||
rmm::cuda_stream_view stream, | ||
rmm::mr::device_memory_resource* mr) | ||
{ | ||
return convert_timestamp(input, transitions, tz_index, false, stream, mr); | ||
} | ||
|
||
} // namespace spark_rapids_jni |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,69 @@ | ||
/* | ||
* Copyright (c) 2023, NVIDIA CORPORATION. | ||
* | ||
* Licensed under the Apache License, Version 2.0 (the "License"); | ||
* you may not use this file except in compliance with the License. | ||
* You may obtain a copy of the License at | ||
* | ||
* http://www.apache.org/licenses/LICENSE-2.0 | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, | ||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
* See the License for the specific language governing permissions and | ||
* limitations under the License. | ||
*/ | ||
|
||
#include <cudf/column/column_view.hpp> | ||
#include <cudf/table/table.hpp> | ||
#include <cudf/utilities/default_stream.hpp> | ||
|
||
#include <rmm/cuda_stream_view.hpp> | ||
|
||
#include <cstddef> | ||
|
||
namespace spark_rapids_jni { | ||
|
||
/** | ||
* @brief Convert input column timestamps in current timezone to UTC | ||
* | ||
* The transition rules are in enclosed in a table, and the index corresponding to the | ||
* current timezone is given. | ||
* | ||
* This method is the inverse of convert_utc_timestamp_to_timezone. | ||
* | ||
* @param input the column of input timestamps in the current timezone | ||
* @param transitions the table of transitions for all timezones | ||
* @param tz_index the index of the row in `transitions` corresponding to the current timezone | ||
* @param stream CUDA stream used for device memory operations and kernel launches. | ||
* @param mr Device memory resource used to allocate the returned timestamp column's memory | ||
*/ | ||
std::unique_ptr<cudf::column> convert_timestamp_to_utc( | ||
cudf::column_view const& input, | ||
cudf::table_view const& transitions, | ||
cudf::size_type tz_index, | ||
rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); | ||
|
||
/** | ||
* @brief Convert input column timestamps in UTC to specified timezone | ||
* | ||
* The transition rules are in enclosed in a table, and the index corresponding to the | ||
* specific timezone is given. | ||
* | ||
* This method is the inverse of convert_timestamp_to_utc. | ||
* | ||
* @param input the column of input timestamps in UTC | ||
* @param transitions the table of transitions for all timezones | ||
* @param tz_index the index of the row in `transitions` corresponding to the specific timezone | ||
* @param stream CUDA stream used for device memory operations and kernel launches. | ||
* @param mr Device memory resource used to allocate the returned timestamp column's memory | ||
*/ | ||
std::unique_ptr<cudf::column> convert_utc_timestamp_to_timezone( | ||
cudf::column_view const& input, | ||
cudf::table_view const& transitions, | ||
cudf::size_type tz_index, | ||
rmm::cuda_stream_view stream = cudf::get_default_stream(), | ||
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); | ||
|
||
} // namespace spark_rapids_jni |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.