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

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

jose-torres 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_r254395007
 
 

 ##########
 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:
   I’m not sure I understand this. Don’t we throw away this val immediately when we leave the foreach scope?

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