Add ContinuousQueryInfo to make ContinuousQueryListener events serializable

Author: zsxwing

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala

@@ -34,7 +34,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap}
 import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan}
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.execution.QueryExecution
-import org.apache.spark.sql.util.ContinuousQueryListener
+import org.apache.spark.sql.util.{ContinuousQueryInfo, ContinuousQueryListener}
 import org.apache.spark.sql.util.ContinuousQueryListener._
 import org.apache.spark.util.{Clock, UninterruptibleThread, Utils}
 
@@ -167,7 +167,7 @@ class StreamExecution(
       // Mark ACTIVE and then post the event. QueryStarted event is synchronously sent to listeners,
       // so must mark this as ACTIVE first.
       state = ACTIVE
-      postEvent(new QueryStarted(this)) // Assumption: Does not throw exception.
+      postEvent(new QueryStarted(this.toInfo)) // Assumption: Does not throw exception.
 
       // Unblock starting thread
       startLatch.countDown()
@@ -206,7 +206,7 @@ class StreamExecution(
     } finally {
       state = TERMINATED
       sparkSession.streams.notifyQueryTermination(StreamExecution.this)
-      postEvent(new QueryTerminated(this))
+      postEvent(new QueryTerminated(this.toInfo))
       terminationLatch.countDown()
     }
   }
@@ -374,7 +374,7 @@ class StreamExecution(
     logInfo(s"Completed up to $availableOffsets in ${batchTime}ms")
     // Update committed offsets.
     committedOffsets ++= availableOffsets
-    postEvent(new QueryProgress(this))
+    postEvent(new QueryProgress(this.toInfo))
   }
 
   private def postEvent(event: ContinuousQueryListener.Event) {
@@ -484,6 +484,15 @@ class StreamExecution(
      """.stripMargin
   }
 
+  private def toInfo: ContinuousQueryInfo = {
+    ContinuousQueryInfo(
+      this.name,
+      this.isActive,
+      this.sourceStatuses,
+      this.sinkStatus,
+      this.exception)
+  }
+
   trait State
   case object INITIALIZED extends State
   case object ACTIVE extends State

sql/core/src/main/scala/org/apache/spark/sql/util/ContinuousQueryListener.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.util
 
 import org.apache.spark.annotation.Experimental
-import org.apache.spark.sql.ContinuousQuery
+import org.apache.spark.sql._
 import org.apache.spark.sql.util.ContinuousQueryListener._
 
 /**
@@ -65,11 +65,23 @@ object ContinuousQueryListener {
   trait Event
 
   /** Event representing the start of a query */
-  class QueryStarted private[sql](val query: ContinuousQuery) extends Event
+  class QueryStarted private[sql](val queryInfo: ContinuousQueryInfo) extends Event
 
   /** Event representing any progress updates in a query */
-  class QueryProgress private[sql](val query: ContinuousQuery) extends Event
+  class QueryProgress private[sql](val queryInfo: ContinuousQueryInfo) extends Event
 
   /** Event representing that termination of a query */
-  class QueryTerminated private[sql](val query: ContinuousQuery) extends Event
+  class QueryTerminated private[sql](val queryInfo: ContinuousQueryInfo) extends Event
 }
+
+/**
+ * :: Experimental ::
+ * Companion object of [[ContinuousQueryListener]] that defines the listener events.
+ */
+@Experimental
+case class ContinuousQueryInfo(
+    name: String,
+    isActive: Boolean,
+    sourceStatuses: Seq[SourceStatus],
+    sinkStatus: SinkStatus,
+    exception: Option[ContinuousQueryException])

sql/core/src/test/scala/org/apache/spark/sql/util/ContinuousQueryListenerSuite.scala

@@ -52,7 +52,7 @@ class ContinuousQueryListenerSuite extends StreamTest with SharedSQLContext with
         Assert("Incorrect query status in onQueryStarted") {
           val status = listener.startStatus
           assert(status != null)
-          assert(status.active == true)
+          assert(status.isActive == true)
           assert(status.sourceStatuses.size === 1)
           assert(status.sourceStatuses(0).description.contains("Memory"))
 
@@ -74,7 +74,7 @@ class ContinuousQueryListenerSuite extends StreamTest with SharedSQLContext with
             assert(listener.progressStatuses.size === 1)
             val status = listener.progressStatuses.peek()
             assert(status != null)
-            assert(status.active == true)
+            assert(status.isActive == true)
             assert(status.sourceStatuses(0).offset === Some(LongOffset(0)))
             assert(status.sinkStatus.offset === CompositeOffset.fill(LongOffset(0)))
 
@@ -88,7 +88,7 @@ class ContinuousQueryListenerSuite extends StreamTest with SharedSQLContext with
             val status = listener.terminationStatus
             assert(status != null)
 
-            assert(status.active === false) // must be inactive by the time onQueryTerm is called
+            assert(status.isActive === false) // must be inactive by the time onQueryTerm is called
             assert(status.sourceStatuses(0).offset === Some(LongOffset(0)))
             assert(status.sinkStatus.offset === CompositeOffset.fill(LongOffset(0)))
           }
@@ -165,9 +165,9 @@ class ContinuousQueryListenerSuite extends StreamTest with SharedSQLContext with
     // to catch errors in the async listener events
     @volatile private var asyncTestWaiter = new Waiter
 
-    @volatile var startStatus: QueryStatus = null
-    @volatile var terminationStatus: QueryStatus = null
-    val progressStatuses = new ConcurrentLinkedQueue[QueryStatus]
+    @volatile var startStatus: ContinuousQueryInfo = null
+    @volatile var terminationStatus: ContinuousQueryInfo = null
+    val progressStatuses = new ConcurrentLinkedQueue[ContinuousQueryInfo]
 
     def reset(): Unit = {
       startStatus = null
@@ -183,35 +183,23 @@ class ContinuousQueryListenerSuite extends StreamTest with SharedSQLContext with
 
     override def onQueryStarted(queryStarted: QueryStarted): Unit = {
       asyncTestWaiter {
-        startStatus = QueryStatus(queryStarted.query)
+        startStatus = queryStarted.queryInfo
       }
     }
 
     override def onQueryProgress(queryProgress: QueryProgress): Unit = {
       asyncTestWaiter {
         assert(startStatus != null, "onQueryProgress called before onQueryStarted")
-        progressStatuses.add(QueryStatus(queryProgress.query))
+        progressStatuses.add(queryProgress.queryInfo)
       }
     }
 
     override def onQueryTerminated(queryTerminated: QueryTerminated): Unit = {
       asyncTestWaiter {
         assert(startStatus != null, "onQueryTerminated called before onQueryStarted")
-        terminationStatus = QueryStatus(queryTerminated.query)
+        terminationStatus = queryTerminated.queryInfo
       }
       asyncTestWaiter.dismiss()
     }
   }
-
-  case class QueryStatus(
-    active: Boolean,
-    exception: Option[Exception],
-    sourceStatuses: Array[SourceStatus],
-    sinkStatus: SinkStatus)
-
-  object QueryStatus {
-    def apply(query: ContinuousQuery): QueryStatus = {
-      QueryStatus(query.isActive, query.exception, query.sourceStatuses, query.sinkStatus)
-    }
-  }
 }