Add support of DateType and TimestampType for GetTimestamp expression (…

…NVIDIA#1779) Signed-off-by: Niranjan Artal <[email protected]>
nartal1 · Feb 23, 2021 · d6d171f · d6d171f
1 parent 3bcecd7
commit d6d171f
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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")
+  )
 }