From b09775a86c9d50c4bee5cf8f7736c45bcb242c4c Mon Sep 17 00:00:00 2001 From: Raza Jafri Date: Fri, 17 Jul 2020 13:04:10 -0700 Subject: [PATCH] [REVIEW] Timesub tests (#365) * timesub tests * addressed review comments * TimeSub should be able to subtract negative interval Co-authored-by: Raza Jafri --- .../src/main/python/date_time_test.py | 14 ++++++++++++-- .../nvidia/spark/rapids/TimeOperatorsSuite.scala | 11 ----------- .../spark/sql/rapids/datetimeExpressions.scala | 6 ++++-- 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/integration_tests/src/main/python/date_time_test.py b/integration_tests/src/main/python/date_time_test.py index 3b881531fb1..5373628c7bb 100644 --- a/integration_tests/src/main/python/date_time_test.py +++ b/integration_tests/src/main/python/date_time_test.py @@ -13,14 +13,24 @@ # limitations under the License. import pytest - from asserts import assert_gpu_and_cpu_are_equal_collect from data_gen import * -from datetime import date, datetime, timedelta, timezone +from datetime import date, datetime, timezone from marks import incompat from pyspark.sql.types import * import pyspark.sql.functions as f +# We only support literal intervals for TimeSub +vals = [(-584, 1563), (1943, 1101), (2693, 2167), (2729, 0), (44, 1534), (2635, 3319), + (1885, -2828), (0, 2463), (932, 2286), (0, 0)] +@pytest.mark.parametrize('data_gen', vals, ids=idfn) +def test_timesub(data_gen): + days, seconds = data_gen + assert_gpu_and_cpu_are_equal_collect( + # We are starting at year 0015 to make sure we don't go before year 0001 while doing TimeSub + lambda spark: unary_op_df(spark, TimestampGen(start=datetime(15, 1, 1, tzinfo=timezone.utc)), seed=1) + .selectExpr("a - (interval {} days {} seconds)".format(days, seconds))) + @pytest.mark.parametrize('data_gen', date_gens, ids=idfn) def test_datediff(data_gen): assert_gpu_and_cpu_are_equal_collect( diff --git a/integration_tests/src/test/scala/com/nvidia/spark/rapids/TimeOperatorsSuite.scala b/integration_tests/src/test/scala/com/nvidia/spark/rapids/TimeOperatorsSuite.scala index 26d1a1277e7..949a78945ad 100644 --- a/integration_tests/src/test/scala/com/nvidia/spark/rapids/TimeOperatorsSuite.scala +++ b/integration_tests/src/test/scala/com/nvidia/spark/rapids/TimeOperatorsSuite.scala @@ -32,15 +32,4 @@ class TimeOperatorsSuite extends SparkQueryCompareTestSuite { frame => frame.select(from_unixtime(col("dates"),"dd/LL/yy HH:mm:ss.SSSSSS")) } - testSparkResultsAreEqual("Test timesub - 4000 seconds", epochDf) { - frame => frame.selectExpr("cast(dates as timestamp) - (interval 40000 seconds)") - } - - testSparkResultsAreEqual("Test timesub - 4 day", epochDf) { - frame => frame.selectExpr("cast(dates as timestamp) - (interval 4 days)") - } - - testSparkResultsAreEqual("Test timesub - 4 day 1000 seconds", epochDf) { - frame => frame.selectExpr("cast(dates as timestamp) - (interval 4 days 1000 seconds)") - } } diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala index 51a17d7715f..3edc4570d35 100644 --- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala +++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala @@ -152,8 +152,8 @@ case class GpuTimeSub( if (intvl.months != 0) { throw new UnsupportedOperationException("Months aren't supported at the moment") } - val usToSub = intvl.days * 24 * 60 * 60 * 1000 * 1000L + intvl.microseconds - if (usToSub > 0) { + val usToSub = intvl.days.toLong * 24 * 60 * 60 * 1000 * 1000 + intvl.microseconds + if (usToSub != 0) { withResource(Scalar.fromLong(usToSub)) { us_s => withResource(l.getBase.castTo(DType.INT64)) { us => withResource(us.sub(us_s)) {longResult => @@ -161,6 +161,8 @@ case class GpuTimeSub( } } } + } else { + l.incRefCount() } case _ => throw new UnsupportedOperationException("GpuTimeSub takes column and interval as an " +