[GitHub] [spark] HeartSaVioR commented on a change in pull request #23576: [SPARK-26655] [SS] Support multiple aggregates in append mode

[GitBox <gi...@apache.org>]

HeartSaVioR commented on a change in pull request #23576: [SPARK-26655] [SS] Support multiple aggregates in append mode
URL: https://github.com/apache/spark/pull/23576#discussion_r287939969
 
 

 ##########
 File path: sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/WatermarkTracker.scala
 ##########
 @@ -80,34 +80,84 @@ case object MaxWatermark extends MultipleWatermarkPolicy {
 /** Tracks the watermark value of a streaming query based on a given `policy` */
 case class WatermarkTracker(policy: MultipleWatermarkPolicy) extends Logging {
   private val operatorToWatermarkMap = mutable.HashMap[Int, Long]()
+  private val statefulOperatorToWatermark = mutable.HashMap[Long, Long]()
+  private val statefulOperatorToEventTimeMap = mutable.HashMap[Long, mutable.HashMap[Int, Long]]()
+
   private var globalWatermarkMs: Long = 0
 
+  private def updateWaterMarkMap(eventTimeExecs: Seq[EventTimeWatermarkExec],
+                                 map: mutable.HashMap[Int, Long]): Unit = {
+    eventTimeExecs.zipWithIndex.foreach {
+      case (e, index) if e.eventTimeStats.value.count > 0 =>
+        logDebug(s"Observed event time stats $index: ${e.eventTimeStats.value}")
+        val newWatermarkMs = e.eventTimeStats.value.max - e.delayMs
+        val prevWatermarkMs = map.get(index)
+        if (prevWatermarkMs.isEmpty || newWatermarkMs > prevWatermarkMs.get) {
+          map.put(index, newWatermarkMs)
+        }
+
+      // Populate 0 if we haven't seen any data yet for this watermark node.
+      case (_, index) =>
+        if (!map.isDefinedAt(index)) {
+          map.put(index, 0)
+        }
+    }
+  }
+
   def setWatermark(newWatermarkMs: Long): Unit = synchronized {
     globalWatermarkMs = newWatermarkMs
   }
 
+  def setOperatorWatermarks(operatorWatermarks: Map[Long, Long]): Unit = synchronized {
+    statefulOperatorToWatermark ++= operatorWatermarks
+  }
+
   def updateWatermark(executedPlan: SparkPlan): Unit = synchronized {
     val watermarkOperators = executedPlan.collect {
       case e: EventTimeWatermarkExec => e
     }
     if (watermarkOperators.isEmpty) return
 
-    watermarkOperators.zipWithIndex.foreach {
-      case (e, index) if e.eventTimeStats.value.count > 0 =>
-        logDebug(s"Observed event time stats $index: ${e.eventTimeStats.value}")
-        val newWatermarkMs = e.eventTimeStats.value.max - e.delayMs
-        val prevWatermarkMs = operatorToWatermarkMap.get(index)
-        if (prevWatermarkMs.isEmpty || newWatermarkMs > prevWatermarkMs.get) {
-          operatorToWatermarkMap.put(index, newWatermarkMs)
-        }
+    updateWaterMarkMap(watermarkOperators, operatorToWatermarkMap)
 
-      // Populate 0 if we haven't seen any data yet for this watermark node.
-      case (_, index) =>
-        if (!operatorToWatermarkMap.isDefinedAt(index)) {
-          operatorToWatermarkMap.put(index, 0)
-        }
+    // compute the per stateful operator watermark
+    val statefulOperators = executedPlan.collect {
+      case s: StatefulOperator => s
     }
 
+    statefulOperators.foreach(statefulOperator => {
+      // find the first event time child node(s)
+      val eventTimeExecs = statefulOperator match {
+        case op: UnaryExecNode =>
+          op.collectFirst {
+            case e: EventTimeWatermarkExec => e
+          }.map(Seq(_)).getOrElse(Seq())
+        case op: BinaryExecNode =>
+          val left = op.left.collectFirst {
+            case e: EventTimeWatermarkExec => e
+          }.map(Seq(_)).getOrElse(Seq())
+          val right = op.right.collectFirst {
+            case e: EventTimeWatermarkExec => e
+          }.map(Seq(_)).getOrElse(Seq())
+          left ++ right
+      }
+
+      // compute watermark for the stateful operator node
+      statefulOperator.stateInfo.foreach(state => {
+        if (eventTimeExecs.nonEmpty) {
+          updateWaterMarkMap(eventTimeExecs,
+            statefulOperatorToEventTimeMap.getOrElseUpdate(state.operatorId,
+              new mutable.HashMap[Int, Long]()))
+          val newWatermarkMs = statefulOperatorToEventTimeMap(state.operatorId).values.toSeq.min
 
 Review comment:
   Just to record here, I agree this should be just `min`.

----------------------------------------------------------------
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.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org