[SPARK-48934][SS][3.4] Python datetime types converted incorrectly for setting timeout in applyInPandasWithState

### What changes were proposed in this pull request?
Fix the way applyInPandasWithState's setTimeoutTimestamp() handles argument of datetime

### Why are the changes needed?
In applyInPandasWithState(), when state.setTimeoutTimestamp() is passed in with datetime.datetime type, it doesn't function as expected. Fix it.
Also, fix another bug of reporting VALUE_NOT_POSITIVE. This issue will trigger when the converted value is 0.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
Add unit test coverage for thie scenario

### Was this patch authored or co-authored using generative AI tooling?
No.

Closes #47413 from HeartSaVioR/SPARK-48934-3.4.

Authored-by: Siying Dong <siying.dong@databricks.com>
Signed-off-by: Jungtaek Lim <kabhwan.opensource@gmail.com>

Author: siying

python/pyspark/sql/streaming/state.py

@@ -18,7 +18,7 @@
 import json
 from typing import Tuple, Optional
 
-from pyspark.sql.types import DateType, Row, StructType
+from pyspark.sql.types import Row, StructType, TimestampType
 from pyspark.sql.utils import has_numpy
 
 __all__ = ["GroupState", "GroupStateTimeout"]
@@ -195,7 +195,7 @@ def setTimeoutTimestamp(self, timestampMs: int) -> None:
             )
 
         if isinstance(timestampMs, datetime.datetime):
-            timestampMs = DateType().toInternal(timestampMs)
+            timestampMs = TimestampType().toInternal(timestampMs) / 1000
 
         if timestampMs <= 0:
             raise ValueError("Timeout timestamp must be positive")

sql/core/src/test/scala/org/apache/spark/sql/streaming/FlatMapGroupsInPandasWithStateSuite.scala

@@ -366,91 +366,101 @@ class FlatMapGroupsInPandasWithStateSuite extends StateStoreMetricsTest {
     )
   }
 
-  test("applyInPandasWithState - streaming w/ event time timeout + watermark") {
-    assume(shouldTestPandasUDFs)
+  Seq(true, false).map { ifUseDateTimeType =>
+    test("applyInPandasWithState - streaming w/ event time timeout + watermark " +
+      s"ifUseDateTimeType=$ifUseDateTimeType") {
+      assume(shouldTestPandasUDFs)
 
-    // timestamp_seconds assumes the base timezone is UTC. However, the provided function
-    // localizes it. Therefore, this test assumes the timezone is in UTC
-    withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") {
-      val pythonScript =
-        """
-          |import calendar
-          |import os
-          |import datetime
-          |import pandas as pd
-          |from pyspark.sql.types import StructType, StringType, StructField, IntegerType
-          |
-          |tpe = StructType([
-          |    StructField("key", StringType()),
-          |    StructField("maxEventTimeSec", IntegerType())])
-          |
-          |def func(key, pdf_iter, state):
-          |    assert state.getCurrentProcessingTimeMs() >= 0
-          |    assert state.getCurrentWatermarkMs() >= -1
-          |
-          |    timeout_delay_sec = 5
-          |    if state.hasTimedOut:
-          |        state.remove()
-          |        yield pd.DataFrame({'key': [key[0]], 'maxEventTimeSec': [-1]})
-          |    else:
-          |        m = state.getOption
-          |        if m is None:
-          |            max_event_time_sec = 0
-          |        else:
-          |            max_event_time_sec = m[0]
-          |
-          |        for pdf in pdf_iter:
-          |            pser = pdf.eventTime.apply(
-          |                lambda dt: (int(calendar.timegm(dt.utctimetuple()) + dt.microsecond)))
-          |            max_event_time_sec = int(max(pser.max(), max_event_time_sec))
-          |
-          |        state.update((max_event_time_sec,))
-          |        timeout_timestamp_sec = max_event_time_sec + timeout_delay_sec
-          |        state.setTimeoutTimestamp(timeout_timestamp_sec * 1000)
-          |        yield pd.DataFrame({'key': [key[0]],
-          |                            'maxEventTimeSec': [max_event_time_sec]})
-          |""".stripMargin
-      val pythonFunc = TestGroupedMapPandasUDFWithState(
-        name = "pandas_grouped_map_with_state", pythonScript = pythonScript)
+      // timestamp_seconds assumes the base timezone is UTC. However, the provided function
+      // localizes it. Therefore, this test assumes the timezone is in UTC
+      withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") {
+        val timeoutMs = if (ifUseDateTimeType) {
+          "datetime.datetime.fromtimestamp(timeout_timestamp_sec)"
+        } else {
+          "timeout_timestamp_sec * 1000"
+        }
 
-      val inputData = MemoryStream[(String, Int)]
-      val inputDataDF =
-        inputData.toDF.select($"_1".as("key"), timestamp_seconds($"_2").as("eventTime"))
-      val outputStructType = StructType(
-        Seq(
-          StructField("key", StringType),
-          StructField("maxEventTimeSec", IntegerType)))
-      val stateStructType = StructType(Seq(StructField("maxEventTimeSec", LongType)))
-      val result =
-        inputDataDF
-          .withWatermark("eventTime", "10 seconds")
-          .groupBy("key")
-          .applyInPandasWithState(
-            pythonFunc(inputDataDF("key"), inputDataDF("eventTime")).expr.asInstanceOf[PythonUDF],
-            outputStructType,
-            stateStructType,
-            "Update",
-            "EventTimeTimeout")
+        val pythonScript =
+          s"""
+             |import calendar
+             |import os
+             |import datetime
+             |import pandas as pd
+             |from pyspark.sql.types import StructType, StringType, StructField, IntegerType
+             |
+             |tpe = StructType([
+             |    StructField("key", StringType()),
+             |    StructField("maxEventTimeSec", IntegerType())])
+             |
+             |def func(key, pdf_iter, state):
+             |    assert state.getCurrentProcessingTimeMs() >= 0
+             |    assert state.getCurrentWatermarkMs() >= -1
+             |
+             |    timeout_delay_sec = 5
+             |    if state.hasTimedOut:
+             |        state.remove()
+             |        yield pd.DataFrame({'key': [key[0]], 'maxEventTimeSec': [-1]})
+             |    else:
+             |        m = state.getOption
+             |        if m is None:
+             |            max_event_time_sec = 0
+             |        else:
+             |            max_event_time_sec = m[0]
+             |
+             |        for pdf in pdf_iter:
+             |            pser = pdf.eventTime.apply(
+             |                lambda dt: (int(calendar.timegm(dt.utctimetuple()) + dt.microsecond)))
+             |            max_event_time_sec = int(max(pser.max(), max_event_time_sec))
+             |
+             |        state.update((max_event_time_sec,))
+             |        timeout_timestamp_sec = max_event_time_sec + timeout_delay_sec
+             |        state.setTimeoutTimestamp($timeoutMs)
+             |        yield pd.DataFrame({'key': [key[0]],
+             |                            'maxEventTimeSec': [max_event_time_sec]})
+             |""".stripMargin.format("")
+        val pythonFunc = TestGroupedMapPandasUDFWithState(
+          name = "pandas_grouped_map_with_state", pythonScript = pythonScript)
 
-      testStream(result, Update)(
-        StartStream(),
+        val inputData = MemoryStream[(String, Int)]
+        val inputDataDF =
+          inputData.toDF().select($"_1".as("key"), timestamp_seconds($"_2").as("eventTime"))
+        val outputStructType = StructType(
+          Seq(
+            StructField("key", StringType),
+            StructField("maxEventTimeSec", IntegerType)))
+        val stateStructType = StructType(Seq(StructField("maxEventTimeSec", LongType)))
+        val result =
+          inputDataDF
+            .withWatermark("eventTime", "10 seconds")
+            .groupBy("key")
+            .applyInPandasWithState(
+              pythonFunc(inputDataDF("key"), inputDataDF("eventTime")).expr.asInstanceOf[PythonUDF],
+              outputStructType,
+              stateStructType,
+              "Update",
+              "EventTimeTimeout")
 
-        AddData(inputData, ("a", 11), ("a", 13), ("a", 15)),
-        // Max event time = 15. Timeout timestamp for "a" = 15 + 5 = 20. Watermark = 15 - 10 = 5.
-        CheckNewAnswer(("a", 15)), // Output = max event time of a
+        testStream(result, Update)(
+          StartStream(),
 
-        AddData(inputData, ("a", 4)), // Add data older than watermark for "a"
-        CheckNewAnswer(), // No output as data should get filtered by watermark
+          AddData(inputData, ("a", 11), ("a", 13), ("a", 15)),
+          // Max event time = 15. Timeout timestamp for "a" = 15 + 5 = 20. Watermark = 15 - 10 = 5.
+          CheckNewAnswer(("a", 15)), // Output = max event time of a
 
-        AddData(inputData, ("a", 10)), // Add data newer than watermark for "a"
-        CheckNewAnswer(("a", 15)), // Max event time is still the same
-        // Timeout timestamp for "a" is still 20 as max event time for "a" is still 15.
-        // Watermark is still 5 as max event time for all data is still 15.
+          AddData(inputData, ("a", 4)), // Add data older than watermark for "a"
+          CheckNewAnswer(), // No output as data should get filtered by watermark
 
-        AddData(inputData, ("b", 31)), // Add data newer than watermark for "b", not "a"
-        // Watermark = 31 - 10 = 21, so "a" should be timed out as timeout timestamp for "a" is 20.
-        CheckNewAnswer(("a", -1), ("b", 31)) // State for "a" should timeout and emit -1
-      )
+          AddData(inputData, ("a", 10)), // Add data newer than watermark for "a"
+          CheckNewAnswer(("a", 15)), // Max event time is still the same
+          // Timeout timestamp for "a" is still 20 as max event time for "a" is still 15.
+          // Watermark is still 5 as max event time for all data is still 15.
+
+          AddData(inputData, ("b", 31)), // Add data newer than watermark for "b", not "a"
+          // Watermark = 31 - 10 = 21, so "a" should be timed out as timeout timestamp for "a" is
+          // 20.
+          CheckNewAnswer(("a", -1), ("b", 31)) // State for "a" should timeout and emit -1
+        )
+      }
     }
   }