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.

Sign up for GitHub

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

Jump to bottom

Support unix_timestamp and to_unix_timestamp with non-UTC timezones (non-DST) #9816

Merged
merged 18 commits into from
Dec 13, 2023
Merged

Support unix_timestamp and to_unix_timestamp with non-UTC timezones (non-DST) #9816

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
3faacd8
Support current_date with timezone
winningsix Nov 15, 2023
9ed56ea
Remove type based timezone checker
winningsix Nov 17, 2023
c7722f9
Add timezone checker for expressions
winningsix Nov 17, 2023
a37d1ce
Add check for cast
winningsix Nov 17, 2023
e26e299
code style fix
winningsix Nov 17, 2023
3696466
Fix premerge fails
winningsix Nov 20, 2023
ce90809
Fix
winningsix Nov 20, 2023
3566f1c
Refine comments
winningsix Nov 21, 2023
749cfbf
Support unix_timestamp with non-DST timezones
winningsix Nov 21, 2023
f239624
Merge branch 'branch-24.02' into unix_timestamp
winningsix Dec 6, 2023
465f0f3
Fix issue after merge
winningsix Dec 7, 2023
877c9eb
Fix
winningsix Dec 8, 2023
83fee5b
Merge branch 'branch-24.02' into unix_timestamp
winningsix Dec 8, 2023
fa439da
Fix merge issue
winningsix Dec 8, 2023
8ed6ff0
Update test cases
winningsix Dec 8, 2023
3b72923
Refine
winningsix Dec 8, 2023
003cf3a
Enable fallback case, address comments
winningsix Dec 11, 2023
0752ffc
Merge branch 'branch-24.02' into unix_timestamp
winningsix Dec 13, 2023
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
28 changes: 18 additions & 10 deletions integration_tests/src/main/python/date_time_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,11 +244,16 @@ def test_dayofyear(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark : unary_op_df(spark, data_gen).select(f.dayofyear(f.col('a'))))


non_utc_unix_time_allow = ['ProjectExec'] if not is_supported_time_zone() else []


@pytest.mark.parametrize('data_gen', date_n_time_gens, ids=idfn)
@allow_non_gpu(*non_utc_allow)
@allow_non_gpu(*non_utc_unix_time_allow)
def test_unix_timestamp(data_gen):
assert_gpu_and_cpu_are_equal_collect(
lambda spark : unary_op_df(spark, data_gen).select(f.unix_timestamp(f.col('a'))))
Copy link
Collaborator

Choose a reason for hiding this comment

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

Add cases for not supported Time Zone.

e.g.:

@pytest.mark.skipif(not is_supported_time_zone(), reason="not all time zones are supported now, refer to https://github.com/NVIDIA/spark-rapids/issues/6839, please update after all time zones are supported")
def test_from_unixtime(data_gen, date_format):

@pytest.mark.skipif(is_supported_time_zone(), reason="not all time zones are supported now, refer to https://github.com/NVIDIA/spark-rapids/issues/6839, please update after all time zones are supported")
def test_from_unixtime_fall_back(data_gen, date_format):

And test TZs locally:
export TZ=Iran
export TZ= America/Los_Angeles

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

updated

lambda spark : unary_op_df(spark, data_gen).select(f.unix_timestamp(f.col('a'))),
{"spark.rapids.sql.nonUTC.enabled": "true"})


@allow_non_gpu('ProjectExec')
Expand Down Expand Up @@ -381,11 +386,12 @@ def fun(spark):

@pytest.mark.parametrize('ansi_enabled', [True, False], ids=['ANSI_ON', 'ANSI_OFF'])
@pytest.mark.parametrize('data_gen', date_n_time_gens, ids=idfn)
@allow_non_gpu(*non_utc_allow)
@allow_non_gpu(*non_utc_unix_time_allow)
def test_unix_timestamp(data_gen, ansi_enabled):
assert_gpu_and_cpu_are_equal_collect(
lambda spark : unary_op_df(spark, data_gen).select(f.unix_timestamp(f.col("a"))),
{'spark.sql.ansi.enabled': ansi_enabled})
{'spark.sql.ansi.enabled': ansi_enabled, "spark.rapids.sql.nonUTC.enabled": "true"})


str_date_and_format_gen = [pytest.param(StringGen('[0-9]{4}/[01][0-9]'),'yyyy/MM', marks=pytest.mark.xfail(reason="cudf does no checks")),
(StringGen('[0-9]{4}/[01][12]/[0-2][1-8]'),'yyyy/MM/dd'),
Expand All @@ -398,12 +404,13 @@ def invalid_date_string_df(spark):
return spark.createDataFrame([['invalid_date_string']], "a string")

@pytest.mark.parametrize('ansi_enabled', [True, False], ids=['ANSI_ON', 'ANSI_OFF'])
@pytest.mark.parametrize('non_utc_timezone_enabled', [True, False], ids=['NOT_UTC_ON', 'NOT_UTC_OFF'])
@pytest.mark.parametrize('data_gen,date_form', str_date_and_format_gen, ids=idfn)
@allow_non_gpu(*non_utc_allow)
def test_string_to_unix_timestamp(data_gen, date_form, ansi_enabled):
@allow_non_gpu('ProjectExec')
Copy link
Collaborator

Choose a reason for hiding this comment

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

I don't want to have a hard coded ProjectExec that does not know if we should fall back to the CPU or not!.

def test_string_to_unix_timestamp(data_gen, date_form, ansi_enabled, non_utc_timezone_enabled):
Copy link
Collaborator

Choose a reason for hiding this comment

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

This parameter is not used and nonUTC.enabled is going away.

assert_gpu_and_cpu_are_equal_collect(
lambda spark : unary_op_df(spark, data_gen, seed=1).selectExpr("to_unix_timestamp(a, '{}')".format(date_form)),
{'spark.sql.ansi.enabled': ansi_enabled})
{'spark.sql.ansi.enabled': ansi_enabled, "spark.rapids.sql.nonUTC.enabled": "true"})

def test_string_to_unix_timestamp_ansi_exception():
assert_gpu_and_cpu_error(
Expand All @@ -412,12 +419,13 @@ def test_string_to_unix_timestamp_ansi_exception():
conf=ansi_enabled_conf)

@pytest.mark.parametrize('ansi_enabled', [True, False], ids=['ANSI_ON', 'ANSI_OFF'])
@pytest.mark.parametrize('non_utc_timezone_enabled', [True, False], ids=['NOT_UTC_ON', 'NOT_UTC_OFF'])
@pytest.mark.parametrize('data_gen,date_form', str_date_and_format_gen, ids=idfn)
@allow_non_gpu(*non_utc_allow)
def test_string_unix_timestamp(data_gen, date_form, ansi_enabled):
@allow_non_gpu('ProjectExec')
Copy link
Collaborator

Choose a reason for hiding this comment

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

Here too.

def test_string_unix_timestamp(data_gen, date_form, ansi_enabled, non_utc_timezone_enabled):
assert_gpu_and_cpu_are_equal_collect(
lambda spark : unary_op_df(spark, data_gen, seed=1).select(f.unix_timestamp(f.col('a'), date_form)),
{'spark.sql.ansi.enabled': ansi_enabled})
{'spark.sql.ansi.enabled': ansi_enabled, "spark.rapids.sql.nonUTC.enabled": non_utc_timezone_enabled})

def test_string_unix_timestamp_ansi_exception():
assert_gpu_and_cpu_error(
Expand Down
25 changes: 23 additions & 2 deletions sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import scala.util.control.NonFatal

import ai.rapids.cudf.DType
import com.nvidia.spark.rapids.RapidsConf.{SUPPRESS_PLANNING_FAILURE, TEST_CONF}
import com.nvidia.spark.rapids.jni.GpuTimeZoneDB
import com.nvidia.spark.rapids.shims._
import org.apache.hadoop.fs.Path

Expand Down Expand Up @@ -625,6 +626,10 @@ object GpuOverrides extends Logging {
timezoneId.normalized() == UTC_TIMEZONE_ID
}

def isUTCTimezone(timezoneIdStr: String): Boolean = {
isUTCTimezone(GpuTimeZoneDB.getZoneId(timezoneIdStr))
}

def isUTCTimezone(): Boolean = {
val zoneId = DateTimeUtils.getZoneId(SQLConf.get.sessionLocalTimeZone)
isUTCTimezone(zoneId.normalized())
Expand Down Expand Up @@ -1695,8 +1700,16 @@ object GpuOverrides extends Logging {
.withPsNote(TypeEnum.STRING, "A limited number of formats are supported"),
TypeSig.STRING)),
(a, conf, p, r) => new UnixTimeExprMeta[ToUnixTimestamp](a, conf, p, r) {
// String type is not supported yet for non-UTC timezone.
override def isTimeZoneSupported: Boolean = a.timeZoneId.forall { zoneID =>
a.left.dataType match {
case _: StringType => GpuOverrides.isUTCTimezone(zoneID)
case _ => true
}
}

override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
GpuToUnixTimestamp(lhs, rhs, sparkFormat, strfFormat)
GpuToUnixTimestamp(lhs, rhs, sparkFormat, strfFormat, a.timeZoneId)
}
}),
expr[UnixTimestamp](
Expand All @@ -1709,8 +1722,16 @@ object GpuOverrides extends Logging {
.withPsNote(TypeEnum.STRING, "A limited number of formats are supported"),
TypeSig.STRING)),
(a, conf, p, r) => new UnixTimeExprMeta[UnixTimestamp](a, conf, p, r) {
// String type is not supported yet for non-UTC timezone.
override def isTimeZoneSupported: Boolean = a.timeZoneId.forall { zoneID =>
a.left.dataType match {
case _: StringType => GpuOverrides.isUTCTimezone(zoneID)
case _ => true
}
}

override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
GpuUnixTimestamp(lhs, rhs, sparkFormat, strfFormat)
GpuUnixTimestamp(lhs, rhs, sparkFormat, strfFormat, a.timeZoneId)
}
}),
expr[Hour](
Expand Down
54 changes: 35 additions & 19 deletions sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -824,25 +824,41 @@ abstract class GpuToTimestamp
val failOnError: Boolean = SQLConf.get.ansiEnabled

override def doColumnar(lhs: GpuColumnVector, rhs: GpuScalar): ColumnVector = {
val tmp = if (lhs.dataType == StringType) {
// rhs is ignored we already parsed the format
if (getTimeParserPolicy == LegacyTimeParserPolicy) {
parseStringAsTimestampWithLegacyParserPolicy(
lhs,
sparkFormat,
strfFormat,
DType.TIMESTAMP_MICROSECONDS,
(col, strfFormat) => col.asTimestampMicroseconds(strfFormat))
} else {
parseStringAsTimestamp(
lhs,
sparkFormat,
strfFormat,
DType.TIMESTAMP_MICROSECONDS,
failOnError)
}
} else { // Timestamp or DateType
lhs.getBase.asTimestampMicroseconds()
val tmp = lhs.dataType match {
case _: StringType =>
// rhs is ignored we already parsed the format
if (getTimeParserPolicy == LegacyTimeParserPolicy) {
parseStringAsTimestampWithLegacyParserPolicy(
lhs,
sparkFormat,
strfFormat,
DType.TIMESTAMP_MICROSECONDS,
(col, strfFormat) => col.asTimestampMicroseconds(strfFormat))
} else {
parseStringAsTimestamp(
lhs,
sparkFormat,
strfFormat,
DType.TIMESTAMP_MICROSECONDS,
failOnError)
}
case _: DateType =>
timeZoneId match {
Copy link
Collaborator

Choose a reason for hiding this comment

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

Use zoneId defined in TimeZoneAwareExpression

@transient lazy val zoneId: ZoneId = DateTimeUtils.getZoneId(timeZoneId.get)

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Still timeZoneId to match non case but use it later.

case Some(_) =>
if (GpuOverrides.isUTCTimezone(zoneId)) {
lhs.getBase.asTimestampMicroseconds()
} else {
assert(GpuTimeZoneDB.isSupportedTimeZone(zoneId))
withResource(lhs.getBase.asTimestampMicroseconds) { tsInMs =>
GpuTimeZoneDB.fromTimestampToUtcTimestamp(tsInMs, zoneId)
}
}
case None => lhs.getBase.asTimestampMicroseconds()
}
case _ =>
// Consistent with Spark's behavior which ignores timeZone for other types like timestamp
// and timestampNtp.
lhs.getBase.asTimestampMicroseconds()
}
// Return Timestamp value if dataType it is expecting is of TimestampType
if (dataType.equals(TimestampType)) {
Expand Down