From d6d171f10245f547ab50c756fa51065536d9872c Mon Sep 17 00:00:00 2001
From: Niranjan Artal <50492963+nartal1@users.noreply.github.com>
Date: Mon, 22 Feb 2021 17:51:10 -0800
Subject: [PATCH] Add support of DateType and TimestampType for GetTimestamp
 expression (#1779)

Signed-off-by: Niranjan Artal <nartal@nvidia.com>
---
 docs/supported_ops.md                         |  8 ++---
 .../expressions/rapids/Timestamp.scala        |  4 ++-
 .../spark/rapids/ParseDateTimeSuite.scala     | 35 +++++++++++++++++++
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/docs/supported_ops.md b/docs/supported_ops.md
index 78ceb73e177..f18b65ffb8c 100644
--- a/docs/supported_ops.md
+++ b/docs/supported_ops.md
@@ -5873,8 +5873,8 @@ Accelerator support is described below.
 <td> </td>
 <td> </td>
 <td> </td>
-<td> </td>
-<td> </td>
+<td>S</td>
+<td>S*</td>
 <td>S</td>
 <td> </td>
 <td> </td>
@@ -5937,8 +5937,8 @@ Accelerator support is described below.
 <td> </td>
 <td> </td>
 <td> </td>
-<td> </td>
-<td> </td>
+<td><b>NS</b></td>
+<td><b>NS</b></td>
 <td><b>NS</b></td>
 <td> </td>
 <td> </td>
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/catalyst/expressions/rapids/Timestamp.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/catalyst/expressions/rapids/Timestamp.scala
index da6833334d8..960a8511727 100644
--- a/sql-plugin/src/main/scala/org/apache/spark/sql/catalyst/expressions/rapids/Timestamp.scala
+++ b/sql-plugin/src/main/scala/org/apache/spark/sql/catalyst/expressions/rapids/Timestamp.scala
@@ -28,7 +28,9 @@ object TimeStamp {
     GpuOverrides.expr[GetTimestamp](
       "Gets timestamps from strings using given pattern.",
       ExprChecks.binaryProjectNotLambda(TypeSig.TIMESTAMP, TypeSig.TIMESTAMP,
-        ("timeExp", TypeSig.STRING, TypeSig.STRING),
+        ("timeExp",
+            TypeSig.STRING + TypeSig.DATE + TypeSig.TIMESTAMP,
+            TypeSig.STRING + TypeSig.DATE + TypeSig.TIMESTAMP),
         ("format", TypeSig.lit(TypeEnum.STRING)
             .withPsNote(TypeEnum.STRING, "A limited number of formats are supported"),
             TypeSig.STRING)),
diff --git a/tests/src/test/scala/com/nvidia/spark/rapids/ParseDateTimeSuite.scala b/tests/src/test/scala/com/nvidia/spark/rapids/ParseDateTimeSuite.scala
index c4e06b1411e..0979ba4db2a 100644
--- a/tests/src/test/scala/com/nvidia/spark/rapids/ParseDateTimeSuite.scala
+++ b/tests/src/test/scala/com/nvidia/spark/rapids/ParseDateTimeSuite.scala
@@ -16,6 +16,8 @@
 
 package com.nvidia.spark.rapids
 
+import java.sql.{Date, Timestamp}
+
 import org.apache.spark.{SparkConf, SparkException}
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.functions.{col, to_date, to_timestamp, unix_timestamp}
@@ -35,6 +37,18 @@ class ParseDateTimeSuite extends SparkQueryCompareTestSuite {
     df => df.withColumn("c1", to_date(col("c0"), "dd/MM/yyyy"))
   }
 
+  testSparkResultsAreEqual("to_date parse date",
+      dates,
+      conf = new SparkConf().set(SQLConf.LEGACY_TIME_PARSER_POLICY.key, "CORRECTED")) {
+    df => df.withColumn("c1", to_date(col("c0"), "yyyy-MM-dd"))
+  }
+
+  testSparkResultsAreEqual("to_date parse timestamp",
+      timestamps,
+      conf = new SparkConf().set(SQLConf.LEGACY_TIME_PARSER_POLICY.key, "CORRECTED")) {
+    df => df.withColumn("c1", to_date(col("c0"), "yyyy-MM-dd"))
+  }
+
   testSparkResultsAreEqual("to_timestamp yyyy-MM-dd",
       timestampsAsStrings,
       conf = new SparkConf().set(SQLConf.LEGACY_TIME_PARSER_POLICY.key, "CORRECTED")) {
@@ -132,6 +146,16 @@ class ParseDateTimeSuite extends SparkQueryCompareTestSuite {
     assert(cpuNowSeconds <= gpuNowSeconds)
   }
 
+  private def dates(spark: SparkSession) = {
+    import spark.implicits._
+    dateValues.toDF("c0")
+  }
+
+  private def timestamps(spark: SparkSession) = {
+    import spark.implicits._
+    tsValues.toDF("c0")
+  }
+
   private def timestampsAsStrings(spark: SparkSession) = {
     import spark.implicits._
     timestampValues.toDF("c0")
@@ -180,5 +204,16 @@ class ParseDateTimeSuite extends SparkQueryCompareTestSuite {
     "\n1999-12-31",
     "1999/12/31"
   )
+
+  private val dateValues = Seq(
+    Date.valueOf("2020-07-24"),
+    Date.valueOf("2020-07-25"),
+    Date.valueOf("1999-12-31"))
+
+  private val tsValues = Seq(
+    Timestamp.valueOf("2015-07-24 10:00:00.3"),
+    Timestamp.valueOf("2015-07-25 02:02:02.2"),
+    Timestamp.valueOf("1999-12-31 11:59:59.999")
+  )
 }