[SPARK-51402][SQL][TESTS] Test TimeType in UDF

sql/core/src/test/scala/org/apache/spark/sql/UDFSuite.scala

@@ -862,7 +862,7 @@ class UDFSuite extends QueryTest with SharedSparkSession {
       .select(myUdf1(Column("col"))),
     Row(ArrayBuffer(100)))
 
-     val myUdf2 = udf((a: immutable.ArraySeq[Int]) =>
+    val myUdf2 = udf((a: immutable.ArraySeq[Int]) =>
       immutable.ArraySeq.unsafeWrapArray[Int]((a :+ 5 :+ 6).toArray))
     checkAnswer(Seq(Array(1, 2, 3))
       .toDF("col")
@@ -1197,6 +1197,35 @@ class UDFSuite extends QueryTest with SharedSparkSession {
       Row(Row(null)))
   }
 
+  test("SPARK-51402: Test TimeType in UDF") {
+    // Mocks
+    val mockTimeStr = "00:00:00.000000"
+    val input = Seq(java.time.LocalTime.parse(mockTimeStr)).toDF("currentTime")
+    // Regular case
+    val plusHour = udf((l: java.time.LocalTime) => l.plusHours(1))
+    val result = input.select(plusHour($"currentTime").as("newTime"))
+    checkAnswer(result, Row(java.time.LocalTime.parse("01:00:00.000000")) :: Nil)
+    assert(result.schema === new StructType().add("newTime", TimeType()))
+    // UDF produces `null`
+    val nullFunc = udf((_: java.time.LocalTime) => null.asInstanceOf[java.time.LocalTime])
+    val nullResult = input.select(nullFunc($"currentTime").as("nullTime"))
+    checkAnswer(nullResult, Row(null) :: Nil)
+    assert(nullResult.schema === new StructType().add("nullTime", TimeType()))
+    // Input parameter of UDF is null
+    val nullInput = Seq(null.asInstanceOf[java.time.LocalTime]).toDF("nullTime")
+    val constDuration = udf((_: java.time.LocalTime) =>
+      java.time.LocalTime.parse(mockTimeStr))
+    val constResult = nullInput.select(constDuration($"nullTime").as("zeroHour"))
+    checkAnswer(constResult, Row(java.time.LocalTime.parse(mockTimeStr)) :: Nil)
+    assert(constResult.schema === new StructType().add("zeroHour", TimeType()))
+    // Error in the conversion of UDF result to the internal representation of time
+    val invalidFunc = udf((l: java.time.LocalTime) => "Zero".toLong)
+    val e = intercept[SparkException] {
+      input.select(invalidFunc($"currentTime")).collect()
+    }
+    assert(e.getCause.isInstanceOf[java.lang.NumberFormatException])
+  }
+
   test("char/varchar as UDF return type") {
     Seq(CharType(5), VarcharType(5)).foreach { dt =>
       val f = udf(