diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/PipelineExecutionMetadata.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/PipelineExecutionMetadata.scala new file mode 100644 index 0000000000000..7cca5fd2f60ce --- /dev/null +++ b/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) + } + } +} diff --git a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala b/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala index c665c344c53b4..5cf5c35cad4b4 100644 --- a/sql/pipelines/src/main/scala/org/apache/spark/sql/pipelines/graph/FlowExecution.scala +++ b/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) } } } diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/PipelineExecutionMetadataSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/PipelineExecutionMetadataSuite.scala new file mode 100644 index 0000000000000..d051a3f241764 --- /dev/null +++ b/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) + } + } +} diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SinkExecutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SinkExecutionSuite.scala index 9e6010a611e97..233da69687018 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/SinkExecutionSuite.scala +++ b/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, diff --git a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala index 9525a61dedd53..4f018915b4c8b 100644 --- a/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala +++ b/sql/pipelines/src/test/scala/org/apache/spark/sql/pipelines/graph/TriggeredGraphExecutionSuite.scala @@ -17,13 +17,23 @@ package org.apache.spark.sql.pipelines.graph +import java.util.UUID +import java.util.concurrent.ConcurrentLinkedQueue + +import scala.jdk.CollectionConverters._ + import org.scalatest.time.{Seconds, Span} +import org.apache.spark.SparkContext +import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent, SparkListenerJobStart} import org.apache.spark.sql.{functions, Row} import org.apache.spark.sql.catalyst.TableIdentifier import org.apache.spark.sql.classic.{DataFrame, Dataset} import org.apache.spark.sql.connector.catalog.{CatalogV2Util, Identifier, TableCatalog} +import org.apache.spark.sql.execution.SQLExecution import org.apache.spark.sql.execution.streaming.runtime.MemoryStream +import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionStart +import org.apache.spark.sql.pipelines.PipelineExecutionMetadata._ import org.apache.spark.sql.pipelines.common.{FlowStatus, RunState} import org.apache.spark.sql.pipelines.graph.TriggeredGraphExecution.StreamState import org.apache.spark.sql.pipelines.logging.{EventLevel, FlowProgress} @@ -33,6 +43,90 @@ import org.apache.spark.sql.types.{IntegerType, StringType, StructType} class TriggeredGraphExecutionSuite extends ExecutionTest with SharedSparkSession { + test("batch 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 pipelineDef = new TestGraphRegistrationContext(spark) { + registerMaterializedView("attributed_batch", query = dfFlowFunc(Seq(1, 2).toDF("value"))) + } + val graph = pipelineDef.toDataflowGraph + val flowIdentifier = fullyQualifiedIdentifier("attributed_batch") + val callerProperty = "spark.sql.pipelines.test.caller" + + def runPipeline(callerValue: String, callerTag: String): String = { + spark.sparkContext.setLocalProperty(callerProperty, callerValue) + spark.sparkContext.addJobTag(callerTag) + try { + val updateContext = TestPipelineUpdateContext(spark, graph, storageRoot) + updateContext.pipelineExecution.runPipeline() + updateContext.pipelineExecution.awaitCompletion() + updateContext.pipelineExecution.graphExecution.get + .flowExecutions(flowIdentifier) + .executionId + } finally { + spark.sparkContext.setLocalProperty(callerProperty, null) + spark.sparkContext.removeJobTag(callerTag) + } + } + + val firstExecutionId = runPipeline("first", "first-caller-tag") + val secondExecutionId = runPipeline("second", "second-caller-tag") + val expectedCallerState = Map( + firstExecutionId -> ("first", "first-caller-tag"), + secondExecutionId -> ("second", "second-caller-tag")) + assert(expectedCallerState.size == 2) + spark.sparkContext.listenerBus.waitUntilEmpty() + + val attributedJobs = jobStarts.asScala.filter { event => + event.properties.getProperty(FLOW_IDENTIFIER_PROPERTY) == flowIdentifier.quotedString + }.toSeq + assert(attributedJobs.nonEmpty) + + val executionIds = attributedJobs.map( + _.properties.getProperty(FLOW_EXECUTION_ID_PROPERTY)).toSet + assert(executionIds == expectedCallerState.keySet) + executionIds.foreach(UUID.fromString) + + attributedJobs.foreach { event => + val executionId = event.properties.getProperty(FLOW_EXECUTION_ID_PROPERTY) + val (callerValue, callerTag) = expectedCallerState(executionId) + val tags = event.properties + .getProperty(SparkContext.SPARK_JOB_TAGS) + .split(SparkContext.SPARK_JOB_TAGS_SEP) + .toSet + assert(tags.contains(flowExecutionIdTag(executionId))) + assert(tags.contains(callerTag)) + assert((tags intersect Set("first-caller-tag", "second-caller-tag")) == Set(callerTag)) + assert(event.properties.getProperty(callerProperty) == callerValue) + + val sqlExecutionId = event.properties.getProperty(SQLExecution.EXECUTION_ID_KEY).toLong + val sqlExecutionStart = sqlExecutionStarts.asScala.find( + _.executionId == sqlExecutionId).get + assert(sqlExecutionStart.jobTags.contains(flowExecutionIdTag(executionId))) + assert(sqlExecutionStart.jobTags.contains(callerTag)) + assert( + (sqlExecutionStart.jobTags intersect Set("first-caller-tag", "second-caller-tag")) == + Set(callerTag)) + } + } finally { + spark.sparkContext.removeSparkListener(listener) + } + } + /** Returns a Dataset of Longs from the table with the given identifier. */ private def getTable(identifier: TableIdentifier): Dataset[Long] = { val session = spark