[SPARK-51981][SS] Add JobTags to queryStartedEvent

### What changes were proposed in this pull request?
Adding a new jobTags parameter for QueryStartedEvent so that it can be connected to the actual spark connect command that triggered this streaming. Also besides adding the parameter, a fix has been applied to the timestamp because previously json reads the wrong argument

### Why are the changes needed?
Without this, there is no way to tell where does this streaming originate from.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Unit test is added

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #50780 from gjxdxh/lingkai-kong_data/SPARK-51981.

Authored-by: Lingkai Kong <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: Lingkai Kong

python/pyspark/sql/streaming/listener.py

@@ -16,7 +16,7 @@
 #
 import uuid
 import json
-from typing import Any, Dict, List, Optional, TYPE_CHECKING
+from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING
 from abc import ABC, abstractmethod
 
 from pyspark.sql import Row
@@ -178,29 +178,44 @@ class QueryStartedEvent:
     """
 
     def __init__(
-        self, id: uuid.UUID, runId: uuid.UUID, name: Optional[str], timestamp: str
+        self,
+        id: uuid.UUID,
+        runId: uuid.UUID,
+        name: Optional[str],
+        timestamp: str,
+        jobTags: Set[str],
     ) -> None:
         self._id: uuid.UUID = id
         self._runId: uuid.UUID = runId
         self._name: Optional[str] = name
         self._timestamp: str = timestamp
+        self._jobTags: Set[str] = jobTags
 
     @classmethod
     def fromJObject(cls, jevent: "JavaObject") -> "QueryStartedEvent":
+        job_tags = set()
+        java_iterator = jevent.jobTags().iterator()
+        while java_iterator.hasNext():
+            job_tags.add(java_iterator.next().toString())
+
         return cls(
             id=uuid.UUID(jevent.id().toString()),
             runId=uuid.UUID(jevent.runId().toString()),
             name=jevent.name(),
             timestamp=jevent.timestamp(),
+            jobTags=job_tags,
         )
 
     @classmethod
     def fromJson(cls, j: Dict[str, Any]) -> "QueryStartedEvent":
+        # Json4s will convert jobTags to a list, so we need to convert it back to a set.
+        job_tags = j["jobTags"] if "jobTags" in j else []
         return cls(
             id=uuid.UUID(j["id"]),
             runId=uuid.UUID(j["runId"]),
             name=j["name"],
             timestamp=j["timestamp"],
+            jobTags=set(job_tags),
         )
 
     @property
@@ -233,6 +248,13 @@ def timestamp(self) -> str:
         """
         return self._timestamp
 
+    @property
+    def jobTags(self) -> Set[str]:
+        """
+        The job tags of the query.
+        """
+        return self._jobTags
+
 
 class QueryProgressEvent:
     """

python/pyspark/sql/tests/streaming/test_streaming_listener.py

@@ -45,6 +45,7 @@ def check_start_event(self, event):
             datetime.strptime(event.timestamp, "%Y-%m-%dT%H:%M:%S.%fZ")
         except ValueError:
             self.fail("'%s' is not in ISO 8601 format.")
+        self.assertTrue(isinstance(event.jobTags, set))
 
     def check_progress_event(self, event, is_stateful):
         """Check QueryProgressEvent"""
@@ -287,7 +288,7 @@ def get_number_of_public_methods(clz):
             get_number_of_public_methods(
                 "org.apache.spark.sql.streaming.StreamingQueryListener$QueryStartedEvent"
             ),
-            15,
+            16,
             msg,
         )
         self.assertEqual(
@@ -451,20 +452,38 @@ def verify(test_listener):
         verify(TestListenerV2())
 
     def test_query_started_event_fromJson(self):
-        start_event = """
+        start_event_old = """
             {
                 "id" : "78923ec2-8f4d-4266-876e-1f50cf3c283b",
                 "runId" : "55a95d45-e932-4e08-9caa-0a8ecd9391e8",
                 "name" : null,
                 "timestamp" : "2023-06-09T18:13:29.741Z"
             }
         """
-        start_event = QueryStartedEvent.fromJson(json.loads(start_event))
-        self.check_start_event(start_event)
-        self.assertEqual(start_event.id, uuid.UUID("78923ec2-8f4d-4266-876e-1f50cf3c283b"))
-        self.assertEqual(start_event.runId, uuid.UUID("55a95d45-e932-4e08-9caa-0a8ecd9391e8"))
-        self.assertIsNone(start_event.name)
-        self.assertEqual(start_event.timestamp, "2023-06-09T18:13:29.741Z")
+        start_event_old = QueryStartedEvent.fromJson(json.loads(start_event_old))
+        self.check_start_event(start_event_old)
+        self.assertEqual(start_event_old.id, uuid.UUID("78923ec2-8f4d-4266-876e-1f50cf3c283b"))
+        self.assertEqual(start_event_old.runId, uuid.UUID("55a95d45-e932-4e08-9caa-0a8ecd9391e8"))
+        self.assertIsNone(start_event_old.name)
+        self.assertEqual(start_event_old.timestamp, "2023-06-09T18:13:29.741Z")
+        self.assertEqual(start_event_old.jobTags, set())
+
+        start_event_new = """
+            {
+                "id" : "78923ec2-8f4d-4266-876e-1f50cf3c283b",
+                "runId" : "55a95d45-e932-4e08-9caa-0a8ecd9391e8",
+                "name" : null,
+                "timestamp" : "2023-06-09T18:13:29.741Z",
+                "jobTags": ["jobTag1", "jobTag2"]
+            }
+        """
+        start_event_new = QueryStartedEvent.fromJson(json.loads(start_event_new))
+        self.check_start_event(start_event_new)
+        self.assertEqual(start_event_new.id, uuid.UUID("78923ec2-8f4d-4266-876e-1f50cf3c283b"))
+        self.assertEqual(start_event_new.runId, uuid.UUID("55a95d45-e932-4e08-9caa-0a8ecd9391e8"))
+        self.assertIsNone(start_event_new.name)
+        self.assertEqual(start_event_new.timestamp, "2023-06-09T18:13:29.741Z")
+        self.assertEqual(start_event_new.jobTags, set(["jobTag1", "jobTag2"]))
 
     def test_query_terminated_event_fromJson(self):
         terminated_json = """

sql/api/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryListener.scala

@@ -19,10 +19,12 @@ package org.apache.spark.sql.streaming
 
 import java.util.UUID
 
+import scala.jdk.CollectionConverters._
+
 import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper}
 import com.fasterxml.jackson.databind.node.TreeTraversingParser
 import com.fasterxml.jackson.module.scala.{ClassTagExtensions, DefaultScalaModule}
-import org.json4s.{JObject, JString}
+import org.json4s.{JArray, JObject, JString}
 import org.json4s.JsonAST.JValue
 import org.json4s.JsonDSL.{jobject2assoc, pair2Assoc}
 import org.json4s.jackson.JsonMethods.{compact, render}
@@ -123,6 +125,14 @@ object StreamingQueryListener extends Serializable {
     private val tree = mapper.readTree(json)
     def getString(name: String): String = tree.get(name).asText()
     def getUUID(name: String): UUID = UUID.fromString(getString(name))
+    def getStringArray(name: String): List[String] = {
+      val node = tree.get(name)
+      if (node.isArray()) {
+        node.elements().asScala.map(_.asText()).toList
+      } else {
+        List()
+      }
+    }
     def getProgress(name: String): StreamingQueryProgress = {
       val parser = new TreeTraversingParser(tree.get(name), mapper)
       parser.readValueAs(classOf[StreamingQueryProgress])
@@ -146,24 +156,32 @@ object StreamingQueryListener extends Serializable {
    *   User-specified name of the query, null if not specified.
    * @param timestamp
    *   The timestamp to start a query.
+   * @param jobTags
+   *   The job tags that have been assigned to all the jobs started by this thread
    * @since 2.1.0
    */
   @Evolving
   class QueryStartedEvent private[sql] (
       val id: UUID,
       val runId: UUID,
       val name: String,
-      val timestamp: String)
+      val timestamp: String,
+      val jobTags: Set[String])
       extends Event
       with Serializable {
 
+    def this(id: UUID, runId: UUID, name: String, timestamp: String) = {
+      this(id, runId, name, timestamp, Set())
+    }
+
     def json: String = compact(render(jsonValue))
 
     private def jsonValue: JValue = {
       ("id" -> JString(id.toString)) ~
         ("runId" -> JString(runId.toString)) ~
         ("name" -> JString(name)) ~
-        ("timestamp" -> JString(timestamp))
+        ("timestamp" -> JString(timestamp)) ~
+        ("jobTags" -> JArray(jobTags.toList.map(JString)))
     }
   }
 
@@ -175,7 +193,8 @@ object StreamingQueryListener extends Serializable {
         parser.getUUID("id"),
         parser.getUUID("runId"),
         parser.getString("name"),
-        parser.getString("name"))
+        parser.getString("timestamp"),
+        parser.getStringArray("jobTags").toSet)
     }
   }
 

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

@@ -288,7 +288,14 @@ abstract class StreamExecution(
       // `postEvent` does not throw non fatal exception.
       val startTimestamp = triggerClock.getTimeMillis()
       postEvent(
-        new QueryStartedEvent(id, runId, name, progressReporter.formatTimestamp(startTimestamp)))
+        new QueryStartedEvent(
+          id,
+          runId,
+          name,
+          progressReporter.formatTimestamp(startTimestamp),
+          sparkSession.sparkContext.getJobTags()
+        )
+      )
 
       // Unblock starting thread
       startLatch.countDown()

sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingQueryListenerSuite.scala

@@ -262,12 +262,28 @@ class StreamingQueryListenerSuite extends StreamTest with BeforeAndAfter {
       assert(newEvent.id === event.id)
       assert(newEvent.runId === event.runId)
       assert(newEvent.name === event.name)
+      assert(newEvent.timestamp === event.timestamp)
+      assert(newEvent.jobTags === event.jobTags)
     }
 
     testSerialization(
-      new QueryStartedEvent(UUID.randomUUID, UUID.randomUUID, "name", "2016-12-05T20:54:20.827Z"))
+      new QueryStartedEvent(
+        UUID.randomUUID,
+        UUID.randomUUID,
+        "name",
+        "2016-12-05T20:54:20.827Z",
+        Set()
+      )
+    )
     testSerialization(
-      new QueryStartedEvent(UUID.randomUUID, UUID.randomUUID, null, "2016-12-05T20:54:20.827Z"))
+      new QueryStartedEvent(
+        UUID.randomUUID,
+        UUID.randomUUID,
+        null,
+        "2016-12-05T20:54:20.827Z",
+        Set("tag1", "tag2")
+      )
+    )
   }
 
   test("QueryProgressEvent serialization") {
@@ -349,6 +365,32 @@ class StreamingQueryListenerSuite extends StreamTest with BeforeAndAfter {
     }
   }
 
+  test("QueryStartedEvent has the right jobTags set") {
+    val session = spark.newSession()
+    val collector = new EventCollectorV2
+    val jobTag1 = "test-jobTag1"
+    val jobTag2 = "test-jobTag2"
+
+    def runQuery(session: SparkSession): Unit = {
+      collector.reset()
+      session.sparkContext.addJobTag(jobTag1)
+      session.sparkContext.addJobTag(jobTag2)
+      val mem = MemoryStream[Int](implicitly[Encoder[Int]], session.sqlContext)
+      testStream(mem.toDS())(
+        AddData(mem, 1, 2, 3),
+        CheckAnswer(1, 2, 3)
+      )
+      session.sparkContext.listenerBus.waitUntilEmpty()
+      session.sparkContext.clearJobTags()
+    }
+
+    withListenerAdded(collector, session) {
+      runQuery(session)
+      assert(collector.startEvent !== null)
+      assert(collector.startEvent.jobTags === Set(jobTag1, jobTag2))
+    }
+  }
+
   test("listener only posts events from queries started in the related sessions") {
     val session1 = spark.newSession()
     val session2 = spark.newSession()