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

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

arunmahadevan 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_r254421452
 
 

 ##########
 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 {
 
 Review comment:
   We use this to finally update the `statefulOperatorToWatermark` below (its passed to `updateWaterMarkMap`). 
   
   So basically we collect all the event time (via `EventTimeWaterMarkExec`) inputs to a stateful operator. The (input) watermark of that stateful operator is then the minimum of all the input watermarks (the event times minus the lag) coming into that node.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on 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