[SPARK-59395][SDP] Expose deterministic flow execution attribution

sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/PipelineExecutionMetadata.scala

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.pipelines
+
+import org.apache.spark.SparkContext
+import org.apache.spark.annotation.{DeveloperApi, Since}
+
+/**
+ * Metadata used to correlate Spark Declarative Pipelines flow executions with Spark jobs.
+ *
+ * @since 4.4.0
+ */
+@DeveloperApi
+@Since("4.4.0")
+object PipelineExecutionMetadata {
+
+  /** Spark local property containing the canonical identifier of the executing flow. */
+  val FLOW_IDENTIFIER_PROPERTY: String = "spark.sql.pipelines.flow.identifier"
+
+  /** Spark local property containing the unique identifier of the flow execution attempt. */
+  val FLOW_EXECUTION_ID_PROPERTY: String = "spark.sql.pipelines.flow.executionId"
+
+  /** Returns the Spark job tag associated with a flow execution attempt. */
+  def flowExecutionIdTag(executionId: String): String =
+    s"$FLOW_EXECUTION_ID_PROPERTY:$executionId"
+
+  private[pipelines] def withFlowExecutionMetadata[T](
+      sc: SparkContext,
+      flowIdentifier: String,
+      executionId: String)(body: => T): T = {
+    val previousFlowIdentifier = sc.getLocalProperty(FLOW_IDENTIFIER_PROPERTY)
+    val previousExecutionId = sc.getLocalProperty(FLOW_EXECUTION_ID_PROPERTY)
+    val executionTag = flowExecutionIdTag(executionId)
+    val executionTagAlreadySet = sc.getJobTags().contains(executionTag)
+
+    sc.setLocalProperty(FLOW_IDENTIFIER_PROPERTY, flowIdentifier)
+    sc.setLocalProperty(FLOW_EXECUTION_ID_PROPERTY, executionId)
+    if (!executionTagAlreadySet) {
+      sc.addJobTag(executionTag)
+    }
+    try {
+      body
+    } finally {
+      if (!executionTagAlreadySet) {
+        sc.removeJobTag(executionTag)
+      }
+      sc.setLocalProperty(FLOW_IDENTIFIER_PROPERTY, previousFlowIdentifier)
+      sc.setLocalProperty(FLOW_EXECUTION_ID_PROPERTY, previousExecutionId)
+    }
+  }
+}

sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala

@@ -17,6 +17,7 @@
 
 package org.apache.spark.sql.pipelines.graph
 
+import java.util.UUID
 import java.util.concurrent.ThreadPoolExecutor
 import java.util.concurrent.atomic.AtomicBoolean
 
@@ -28,6 +29,7 @@ import org.apache.spark.sql.{Dataset, Row}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.classic.ClassicConversions._
 import org.apache.spark.sql.classic.SparkSession
+import org.apache.spark.sql.pipelines.PipelineExecutionMetadata.withFlowExecutionMetadata
 import org.apache.spark.sql.pipelines.autocdc.{
   Scd1BatchProcessor,
   Scd1ForeachBatchHandler,
@@ -37,7 +39,7 @@ import org.apache.spark.sql.pipelines.autocdc.{
 import org.apache.spark.sql.pipelines.graph.QueryOrigin.ExceptionHelpers
 import org.apache.spark.sql.pipelines.util.SparkSessionUtils
 import org.apache.spark.sql.streaming.{OutputMode, StreamingQuery, Trigger}
-import org.apache.spark.util.ThreadUtils
+import org.apache.spark.util.{ThreadUtils, Utils}
 
 /**
  * A flow's execution may complete for two reasons:
@@ -56,6 +58,9 @@ object ExecutionResult {
 /** A `FlowExecution` specifies how to execute a flow and manages its execution. */
 trait FlowExecution {
 
+  /** Unique identifier for this flow execution attempt. */
+  final val executionId: String = UUID.randomUUID().toString
+
   /** Identifier of this physical flow */
   def identifier: TableIdentifier
 
@@ -78,6 +83,14 @@ trait FlowExecution {
    */
   protected def spark: SparkSession = updateContext.spark
 
+  /** Runs a block with this flow execution's attribution metadata. */
+  protected final def withExecutionMetadata[T](body: => T): T = {
+    withFlowExecutionMetadata(
+      spark.sparkContext,
+      identifier.quotedString,
+      executionId)(body)
+  }
+
   /**
    * Origin to use when recording events for this flow.
    */
@@ -215,7 +228,9 @@ trait StreamingFlowExecution extends FlowExecution with Logging {
       log"Starting ${MDC(LogKeys.TABLE_NAME, identifier)} with " +
       log"checkpoint location ${MDC(LogKeys.CHECKPOINT_PATH, checkpointPath)}"
     )
-    val streamingQuery = SparkSessionUtils.withSqlConf(spark, sqlConf.toList: _*)(startStream())
+    val streamingQuery = SparkSessionUtils.withSqlConf(spark, sqlConf.toList: _*) {
+      withExecutionMetadata(startStream())
+    }
     _streamingQuery = Option(streamingQuery)
     Future(streamingQuery.awaitTermination())
   }
@@ -265,23 +280,32 @@ class BatchTableWrite(
     SparkSessionUtils.withSqlConf(spark, sqlConf.toList: _*) {
       updateContext.flowProgressEventLogger.recordRunning(flow = flow)
       val data = graph.reanalyzeFlow(flow, sessionCaseSensitive).df
+      val localProperties = Utils.cloneProperties(spark.sparkContext.getLocalProperties)
       Future {
-        val dataFrameWriter = data.write
-        destination.format.foreach(dataFrameWriter.format)
-
-        // In "append" mode with saveAsTable, partition/cluster columns must be specified in query
-        // because the format and options of the existing table is used, and the table could
-        // have been created with partition columns.
-        destination.clusterCols.foreach { clusterCols =>
-          dataFrameWriter.clusterBy(clusterCols.head, clusterCols.tail: _*)
-        }
-        destination.partitionCols.foreach { partitionCols =>
-          dataFrameWriter.partitionBy(partitionCols: _*)
+        val previousLocalProperties = spark.sparkContext.getLocalProperties
+        spark.sparkContext.setLocalProperties(localProperties)
+        try {
+          withExecutionMetadata {
+            val dataFrameWriter = data.write
+            destination.format.foreach(dataFrameWriter.format)
+
+            // In "append" mode with saveAsTable, partition/cluster columns must be specified in
+            // query because the format and options of the existing table is used, and the table
+            // could have been created with partition columns.
+            destination.clusterCols.foreach { clusterCols =>
+              dataFrameWriter.clusterBy(clusterCols.head, clusterCols.tail: _*)
+            }
+            destination.partitionCols.foreach { partitionCols =>
+              dataFrameWriter.partitionBy(partitionCols: _*)
+            }
+
+            dataFrameWriter
+              .mode("append")
+              .saveAsTable(destination.identifier.unquotedString)
+          }
+        } finally {
+          spark.sparkContext.setLocalProperties(previousLocalProperties)
         }
-
-        dataFrameWriter
-          .mode("append")
-          .saveAsTable(destination.identifier.unquotedString)
       }
     }
   }

sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/PipelineExecutionMetadataSuite.scala

@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.pipelines
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.pipelines.PipelineExecutionMetadata._
+import org.apache.spark.sql.test.SharedSparkSession
+
+class PipelineExecutionMetadataSuite extends QueryTest with SharedSparkSession {
+
+  test("flow execution metadata names and tag format") {
+    assert(FLOW_IDENTIFIER_PROPERTY == "spark.sql.pipelines.flow.identifier")
+    assert(FLOW_EXECUTION_ID_PROPERTY == "spark.sql.pipelines.flow.executionId")
+    assert(
+      flowExecutionIdTag("execution-id") ==
+        "spark.sql.pipelines.flow.executionId:execution-id")
+  }
+
+  test("flow execution metadata is scoped and restores previous state") {
+    val sc = spark.sparkContext
+    val previousIdentifier = "previous-flow"
+    val previousExecutionId = "previous-execution"
+    val unrelatedTag = "unrelated-tag"
+    val executionId = "execution-id"
+    val executionTag = flowExecutionIdTag(executionId)
+
+    sc.setLocalProperty(FLOW_IDENTIFIER_PROPERTY, previousIdentifier)
+    sc.setLocalProperty(FLOW_EXECUTION_ID_PROPERTY, previousExecutionId)
+    sc.addJobTag(unrelatedTag)
+    sc.addJobTag(executionTag)
+
+    try {
+      withFlowExecutionMetadata(sc, "`catalog`.`schema`.`target`", executionId) {
+        assert(sc.getLocalProperty(FLOW_IDENTIFIER_PROPERTY) ==
+          "`catalog`.`schema`.`target`")
+        assert(sc.getLocalProperty(FLOW_EXECUTION_ID_PROPERTY) == executionId)
+        assert(sc.getJobTags().contains(executionTag))
+        assert(sc.getJobTags().contains(unrelatedTag))
+      }
+
+      assert(sc.getLocalProperty(FLOW_IDENTIFIER_PROPERTY) == previousIdentifier)
+      assert(sc.getLocalProperty(FLOW_EXECUTION_ID_PROPERTY) == previousExecutionId)
+      assert(sc.getJobTags().contains(executionTag))
+      assert(sc.getJobTags().contains(unrelatedTag))
+    } finally {
+      sc.setLocalProperty(FLOW_IDENTIFIER_PROPERTY, null)
+      sc.setLocalProperty(FLOW_EXECUTION_ID_PROPERTY, null)
+      sc.removeJobTag(executionTag)
+      sc.removeJobTag(unrelatedTag)
+    }
+  }
+
+  test("flow execution metadata restores previous state after failure") {
+    val sc = spark.sparkContext
+    val unrelatedTag = "unrelated-tag"
+    val executionId = "execution-id"
+
+    sc.addJobTag(unrelatedTag)
+    try {
+      val error = intercept[SparkException] {
+        withFlowExecutionMetadata(sc, "`target`", executionId) {
+          throw new SparkException("expected failure")
+        }
+      }
+      assert(error.getMessage == "expected failure")
+      assert(sc.getLocalProperty(FLOW_IDENTIFIER_PROPERTY) == null)
+      assert(sc.getLocalProperty(FLOW_EXECUTION_ID_PROPERTY) == null)
+      assert(!sc.getJobTags().contains(flowExecutionIdTag(executionId)))
+      assert(sc.getJobTags().contains(unrelatedTag))
+    } finally {
+      sc.removeJobTag(unrelatedTag)
+    }
+  }
+}

sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SinkExecutionSuite.scala

@@ -17,16 +17,89 @@
 
 package org.apache.spark.sql.pipelines.graph
 
+import java.util.UUID
+import java.util.concurrent.ConcurrentLinkedQueue
+
+import scala.jdk.CollectionConverters._
+
 import org.apache.hadoop.fs.Path
 
+import org.apache.spark.SparkContext
+import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart}
 import org.apache.spark.sql.catalyst.TableIdentifier
 import org.apache.spark.sql.classic.DataFrame
+import org.apache.spark.sql.execution.SQLExecution
 import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream, StreamingQueryWrapper}
+import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionStart
+import org.apache.spark.sql.pipelines.PipelineExecutionMetadata._
 import org.apache.spark.sql.pipelines.utils.{ExecutionTest, TestGraphRegistrationContext}
 import org.apache.spark.sql.streaming.StreamingQuery
 import org.apache.spark.sql.test.SharedSparkSession
 
 class SinkExecutionSuite extends ExecutionTest with SharedSparkSession {
+
+  test("streaming flow execution attribution is exposed on Spark jobs") {
+    val session = spark
+    import session.implicits._
+
+    val jobStarts = new ConcurrentLinkedQueue[SparkListenerJobStart]()
+    val sqlExecutionStarts = new ConcurrentLinkedQueue[SparkListenerSQLExecutionStart]()
+    val listener = new SparkListener {
+      override def onJobStart(jobStart: SparkListenerJobStart): Unit = jobStarts.add(jobStart)
+
+      override def onOtherEvent(event: SparkListenerEvent): Unit = event match {
+        case start: SparkListenerSQLExecutionStart => sqlExecutionStarts.add(start)
+        case _ =>
+      }
+    }
+    spark.sparkContext.addSparkListener(listener)
+
+    try {
+      val ints = MemoryStream[Int]
+      ints.addData(1, 2, 3, 4)
+      val graph = createDataflowGraph(
+        ints.toDF(),
+        "attributed_sink",
+        "flow_to_attributed_sink",
+        "memory")
+      val updateContext = TestPipelineUpdateContext(spark, graph, storageRoot)
+
+      updateContext.pipelineExecution.startPipeline()
+      updateContext.pipelineExecution.awaitCompletion()
+      spark.sparkContext.listenerBus.waitUntilEmpty()
+
+      val flowIdentifier = updateContext.pipelineExecution.graphExecution.get
+        .flowExecutions.keys
+        .find(_.table == "flow_to_attributed_sink")
+        .get
+        .quotedString
+      val attributedJobs = jobStarts.asScala.filter { event =>
+        event.properties.getProperty(FLOW_IDENTIFIER_PROPERTY) == flowIdentifier
+      }.toSeq
+      assert(attributedJobs.nonEmpty)
+
+      val executionIds = attributedJobs.map(
+        _.properties.getProperty(FLOW_EXECUTION_ID_PROPERTY)).toSet
+      assert(executionIds.size == 1)
+      val executionId = executionIds.head
+      UUID.fromString(executionId)
+
+      attributedJobs.foreach { event =>
+        val tags = event.properties
+          .getProperty(SparkContext.SPARK_JOB_TAGS)
+          .split(SparkContext.SPARK_JOB_TAGS_SEP)
+          .toSet
+        assert(tags.contains(flowExecutionIdTag(executionId)))
+        val sqlExecutionId = event.properties.getProperty(SQLExecution.EXECUTION_ID_KEY).toLong
+        val sqlExecutionStart = sqlExecutionStarts.asScala.find(
+          _.executionId == sqlExecutionId).get
+        assert(sqlExecutionStart.jobTags.contains(flowExecutionIdTag(executionId)))
+      }
+    } finally {
+      spark.sparkContext.removeSparkListener(listener)
+    }
+  }
+
   def createDataflowGraph(
       inputs: DataFrame,
       sinkName: String,