Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support from_utc_timestamp on the GPU for non-UTC timezones (non-DST) #9810

Merged
merged 12 commits into from
Nov 28, 2023
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Add plugin initialization code and change config flag for non UTC tim…
…ezone support

Signed-off-by: Navin Kumar <navink@nvidia.com>
  • Loading branch information
NVnavkumar committed Nov 21, 2023
commit 93f5660b833853da981b651f9c72112c2afa6073
4 changes: 2 additions & 2 deletions integration_tests/src/main/python/date_time_test.py
Original file line number Diff line number Diff line change
@@ -281,9 +281,9 @@ def test_from_utc_timestamp_unsupported_timezone_fallback(data_gen, time_zone):
@pytest.mark.parametrize('time_zone', ["UTC", "Asia/Shanghai", "EST", "MST", "VST"], ids=idfn)
@pytest.mark.parametrize('data_gen', [timestamp_gen], ids=idfn)
def test_from_utc_timestamp_supported_timezones(data_gen, time_zone):
# Remove spark.rapids.test.CPU.timezone configuration when GPU kernel is ready to really test on GPU
# TODO: Remove spark.rapids.sql.nonUTC.enabled configuration
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there an issue for this?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Created #9881 and added to timezone epic

assert_gpu_and_cpu_are_equal_collect(
lambda spark: unary_op_df(spark, data_gen).select(f.from_utc_timestamp(f.col('a'), time_zone)), conf = {"spark.rapids.test.CPU.timezone": "true"})
lambda spark: unary_op_df(spark, data_gen).select(f.from_utc_timestamp(f.col('a'), time_zone)), conf = {"spark.rapids.sql.nonUTC.enabled": "true"})
winningsix marked this conversation as resolved.
Show resolved Hide resolved


@allow_non_gpu('ProjectExec')
Original file line number Diff line number Diff line change
@@ -24,6 +24,7 @@ import scala.sys.process._
import scala.util.Try
import ai.rapids.cudf.{Cuda, CudaException, CudaFatalException, CudfException, MemoryCleaner}
import com.nvidia.spark.rapids.filecache.{FileCache, FileCacheLocalityManager, FileCacheLocalityMsg}
import com.nvidia.spark.rapids.jni.GpuTimeZoneDB
import com.nvidia.spark.rapids.python.PythonWorkerSemaphore
import org.apache.commons.lang3.exception.ExceptionUtils
import org.apache.spark.{ExceptionFailure, SparkConf, SparkContext, TaskFailedReason}
@@ -378,6 +379,9 @@ class RapidsExecutorPlugin extends ExecutorPlugin with Logging {
s"Driver timezone is $driverTimezone and executor timezone is " +
s"$executorTimezone. Set executor timezone to $driverTimezone.")
}
if (conf.nonUTCTimeZoneEnabled) {
GpuTimeZoneDB.cacheDatabase()
}
}

GpuCoreDumpHandler.executorInit(conf, pluginContext)
@@ -500,6 +504,9 @@ class RapidsExecutorPlugin extends ExecutorPlugin with Logging {
}

override def shutdown(): Unit = {
if (conf.nonUTCTimeZoneEnabled) {
GpuTimeZoneDB.shutdown()
}
GpuSemaphore.shutdown()
PythonWorkerSemaphore.shutdown()
GpuDeviceManager.shutdown()
Original file line number Diff line number Diff line change
@@ -2044,6 +2044,13 @@ object RapidsConf {
"The gpu to disk spill bounce buffer must have a positive size")
.createWithDefault(128L * 1024 * 1024)

val NON_UTC_TIME_ZONE_ENABLED =
NVnavkumar marked this conversation as resolved.
Show resolved Hide resolved
conf("spark.rapids.sql.nonUTC.enabled")
.doc("An option to enable/disable non-UTC time zone support.")
.internal()
.booleanConf
.createWithDefault(false)

val SPLIT_UNTIL_SIZE_OVERRIDE = conf("spark.rapids.sql.test.overrides.splitUntilSize")
.doc("Only for tests: override the value of GpuDeviceManager.splitUntilSize")
.internal()
@@ -2056,12 +2063,6 @@ object RapidsConf {
.booleanConf
.createOptional

val TEST_USE_TIMEZONE_CPU_BACKEND = conf("spark.rapids.test.CPU.timezone")
.doc("Only for tests: verify for timezone related functions")
.internal()
.booleanConf
.createOptional

private def printSectionHeader(category: String): Unit =
println(s"\n### $category")

@@ -2754,6 +2755,8 @@ class RapidsConf(conf: Map[String, String]) extends Logging {

lazy val splitUntilSizeOverride: Option[Long] = get(SPLIT_UNTIL_SIZE_OVERRIDE)

lazy val nonUTCTimeZoneEnabled: Boolean = get(NON_UTC_TIME_ZONE_ENABLED)

private val optimizerDefaults = Map(
// this is not accurate because CPU projections do have a cost due to appending values
// to each row that is produced, but this needs to be a really small number because
Original file line number Diff line number Diff line change
@@ -23,7 +23,6 @@ import ai.rapids.cudf.{BinaryOp, CaptureGroups, ColumnVector, ColumnView, DType,
import com.nvidia.spark.rapids.{BinaryExprMeta, BoolUtils, DataFromReplacementRule, DateUtils, GpuBinaryExpression, GpuBinaryExpressionArgsAnyScalar, GpuCast, GpuColumnVector, GpuExpression, GpuScalar, GpuUnaryExpression, RapidsConf, RapidsMeta}
import com.nvidia.spark.rapids.Arm._
import com.nvidia.spark.rapids.GpuOverrides.{extractStringLit, getTimeParserPolicy}
import com.nvidia.spark.rapids.RapidsConf.TEST_USE_TIMEZONE_CPU_BACKEND
import com.nvidia.spark.rapids.RapidsPluginImplicits._
import com.nvidia.spark.rapids.shims.ShimBinaryExpression

@@ -1046,7 +1045,7 @@ class FromUTCTimestampExprMeta(
extends BinaryExprMeta[FromUTCTimestamp](expr, conf, parent, rule) {

private[this] var timezoneId: ZoneId = null
private[this] val isOnCPU: Boolean = conf.get(TEST_USE_TIMEZONE_CPU_BACKEND).getOrElse(false)
private[this] val isOnCPU: Boolean = conf.nonUTCTimeZoneEnabled

override def tagExprForGpu(): Unit = {
extractStringLit(expr.right) match {