anishshri-db commented on code in PR #45051:
URL: https://github.com/apache/spark/pull/45051#discussion_r1522188039
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/TransformWithStateSuite.scala:
##########
@@ -53,8 +57,134 @@ class RunningCountStatefulProcessor extends
StatefulProcessor[String, String, (S
Iterator((key, count.toString))
}
}
+}
+
+// Class to verify stateful processor usage with adding processing time timers
+class RunningCountStatefulProcessorWithProcTimeTimer extends
RunningCountStatefulProcessor {
+ private def handleProcessingTimeBasedTimers(
+ key: String,
+ expiryTimestampMs: Long): Iterator[(String, String)] = {
+ _countState.clear()
+ Iterator((key, "-1"))
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[String],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String)] = {
+
+ if (expiredTimerInfo.isValid()) {
+ handleProcessingTimeBasedTimers(key,
expiredTimerInfo.getExpiryTimeInMs())
+ } else {
+ val currCount = _countState.getOption().getOrElse(0L)
+ if (currCount == 0 && (key == "a" || key == "c")) {
+
_processorHandle.registerTimer(timerValues.getCurrentProcessingTimeInMs()
+ + 5000)
+ }
+
+ val count = currCount + 1
+ if (count == 3) {
+ _countState.clear()
+ Iterator.empty
+ } else {
+ _countState.update(count)
+ Iterator((key, count.toString))
+ }
+ }
+ }
+}
- override def close(): Unit = {}
+// Class to verify stateful processor usage with adding/deleting processing
time timers
+class RunningCountStatefulProcessorWithAddRemoveProcTimeTimer
+ extends RunningCountStatefulProcessor {
+ @transient private var _timerState: ValueState[Long] = _
+
+ override def init(
+ handle: StatefulProcessorHandle,
+ outputMode: OutputMode,
+ timeoutMode: TimeoutMode) : Unit = {
+ super.init(handle, outputMode, timeoutMode)
+ _timerState = _processorHandle.getValueState[Long]("timerState")
+ }
+
+ private def handleProcessingTimeBasedTimers(
+ key: String,
+ expiryTimestampMs: Long): Iterator[(String, String)] = {
+ _timerState.clear()
+ Iterator((key, "-1"))
+ }
+
+ override def handleInputRows(
+ key: String,
+ inputRows: Iterator[String],
+ timerValues: TimerValues,
+ expiredTimerInfo: ExpiredTimerInfo): Iterator[(String, String)] = {
+ if (expiredTimerInfo.isValid()) {
+ handleProcessingTimeBasedTimers(key,
expiredTimerInfo.getExpiryTimeInMs())
+ } else {
+ val currCount = _countState.getOption().getOrElse(0L)
+ val count = currCount + inputRows.size
+ _countState.update(count)
+ if (key == "a") {
+ var nextTimerTs: Long = 0L
+ if (currCount == 0) {
+ nextTimerTs = timerValues.getCurrentProcessingTimeInMs() + 5000
+ _processorHandle.registerTimer(nextTimerTs)
+ _timerState.update(nextTimerTs)
+ } else if (currCount == 1) {
+ _processorHandle.deleteTimer(_timerState.get())
Review Comment:
Yes - added it in a different test
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
