From 33aa6948ab2d1806e95eeda32e0a14a55b391b2e Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Wed, 24 Dec 2025 22:22:36 +0800 Subject: [PATCH 1/4] avoid unnecessary temp QueryExecution for nested command execution --- .../spark/sql/execution/QueryExecution.scala | 34 ++++++++++++++---- .../execution/datasources/DataSource.scala | 3 +- .../v2/WriteToDataSourceV2Exec.scala | 9 +++-- .../DataSourceV2DataFrameSuite.scala | 35 +++++++++++++++++++ 4 files changed, 69 insertions(+), 12 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index 3e0aef962e719..64b75f8d17959 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -53,6 +53,7 @@ import org.apache.spark.sql.scripting.SqlScriptingExecution import org.apache.spark.sql.streaming.OutputMode import org.apache.spark.util.{LazyTry, Utils} import org.apache.spark.util.ArrayImplicits._ +import org.apache.spark.sql.execution.QueryExecution.{runCommand => refreshPhaseEnabledmode} /** * The primary workflow for executing relational queries using Spark. Designed to allow easy @@ -178,13 +179,8 @@ class QueryExecution( // with the rest of processing of the root plan being just outputting command results, // for eagerly executed commands we mark this place as beginning of execution. tracker.setReadyForExecution() - val qe = new QueryExecution(sparkSession, p, mode = mode, - shuffleCleanupMode = shuffleCleanupMode, refreshPhaseEnabled = refreshPhaseEnabled) - val result = QueryExecution.withInternalError(s"Eagerly executed $name failed.") { - SQLExecution.withNewExecutionId(qe, Some(name)) { - qe.executedPlan.executeCollect() - } - } + val (qe, result) = QueryExecution.runCommand( + sparkSession, p, name, refreshPhaseEnabled, mode, Some(shuffleCleanupMode)) CommandResult( qe.analyzed.output, qe.commandExecuted, @@ -763,4 +759,28 @@ object QueryExecution { case _ => false } } + + def runCommand( + sparkSession: SparkSession, + command: LogicalPlan, + name: String, + refreshPhaseEnabled: Boolean = true, + mode: CommandExecutionMode.Value = CommandExecutionMode.SKIP, + shuffleCleanupModeOpt: Option[ShuffleCleanupMode] = None) + : (QueryExecution, Array[InternalRow]) = { + val shuffleCleanupMode = shuffleCleanupModeOpt.getOrElse( + determineShuffleCleanupMode(sparkSession.sessionState.conf)) + val qe = new QueryExecution( + sparkSession, + command, + mode = mode, + shuffleCleanupMode = shuffleCleanupMode, + refreshPhaseEnabled = refreshPhaseEnabled) + val result = QueryExecution.withInternalError(s"Executed $name failed.") { + SQLExecution.withNewExecutionId(qe, Some(name)) { + qe.executedPlan.executeCollect() + } + } + (qe, result) + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index da3204eb221b7..be45704fdfc56 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -531,8 +531,7 @@ case class DataSource( disallowWritingIntervals( outputColumns.toStructType.asNullable, format.toString, forbidAnsiIntervals = false) val cmd = planForWritingFileFormat(format, mode, data) - val qe = sessionState(sparkSession).executePlan(cmd) - qe.assertCommandExecuted() + QueryExecution.runCommand(sparkSession, cmd, "file source write") // Replace the schema with that of the DataFrame we just wrote out to avoid re-inferring copy(userSpecifiedSchema = Some(outputColumns.toStructType.asNullable)).resolveRelation() case _ => throw SparkException.internalError( diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala index 3e4a2f792a1c8..472d65dfb2b1c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.connector.metric.CustomMetric import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, DeltaWrite, DeltaWriter, MergeSummaryImpl, PhysicalWriteInfoImpl, Write, WriterCommitMessage, WriteSummary} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} -import org.apache.spark.sql.execution.{QueryExecution, SparkPlan, SQLExecution, UnaryExecNode} +import org.apache.spark.sql.execution.{CommandExecutionMode, QueryExecution, SparkPlan, SQLExecution, UnaryExecNode} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLMetric, SQLMetrics} import org.apache.spark.sql.types.StructType @@ -737,8 +737,11 @@ private[v2] trait V2CreateTableAsSelectBaseExec extends LeafV2CommandExec { } else { AppendData.byPosition(relation, query, writeOptions) } - val qe = QueryExecution.create(session, writeCommand, refreshPhaseEnabled) - qe.assertCommandExecuted() + QueryExecution.runCommand( + session, + writeCommand, + "inner data writing for CTAS/RTAS", + refreshPhaseEnabled) DataSourceV2Utils.commitStagedChanges(sparkContext, table, metrics) Nil })(catchBlock = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala index c3164b3428f95..9bd59f23e2fff 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala @@ -2053,4 +2053,39 @@ class DataSourceV2DataFrameSuite case _ => fail(s"can't pin $ident in $catalogName") } } + + test("SPARK-50844: CTAS/RTAS should trigger two query executions") { + // CTAS/RTAS triggers 2 query executions: + // 1. The outer CTAS/RTAS command execution + // 2. The inner AppendData/OverwriteByExpression execution + var executionCount = 0 + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = { + executionCount += 1 + } + override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit = {} + } + + try { + spark.listenerManager.register(listener) + val t = "testcat.ns1.ns2.tbl" + withTable(t) { + // Test CTAS (CreateTableAsSelect) + executionCount = 0 + sql(s"CREATE TABLE $t USING foo AS SELECT 1 as id, 'a' as data") + sparkContext.listenerBus.waitUntilEmpty() + assert(executionCount == 2, + s"CTAS should trigger 2 executions, got $executionCount") + + // Test RTAS (ReplaceTableAsSelect) + executionCount = 0 + sql(s"CREATE OR REPLACE TABLE $t USING foo AS SELECT 2 as id, 'b' as data") + sparkContext.listenerBus.waitUntilEmpty() + assert(executionCount == 2, + s"RTAS should trigger 2 executions, got $executionCount") + } + } finally { + spark.listenerManager.unregister(listener) + } + } } From 22629903fed182a239b23843da517bc3ff71ecea Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Thu, 25 Dec 2025 10:35:44 +0800 Subject: [PATCH 2/4] Update sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala --- .../apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala index 9bd59f23e2fff..3a327fcf98636 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2DataFrameSuite.scala @@ -2054,7 +2054,7 @@ class DataSourceV2DataFrameSuite } } - test("SPARK-50844: CTAS/RTAS should trigger two query executions") { + test("CTAS/RTAS should trigger two query executions") { // CTAS/RTAS triggers 2 query executions: // 1. The outer CTAS/RTAS command execution // 2. The inner AppendData/OverwriteByExpression execution From af20fd387acc136862bba2bb589dba73b73ad4b7 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Thu, 25 Dec 2025 16:33:48 +0800 Subject: [PATCH 3/4] Apply suggestions from code review --- .../scala/org/apache/spark/sql/execution/QueryExecution.scala | 1 - .../sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala index 64b75f8d17959..652cf5aa7c3ef 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala @@ -53,7 +53,6 @@ import org.apache.spark.sql.scripting.SqlScriptingExecution import org.apache.spark.sql.streaming.OutputMode import org.apache.spark.util.{LazyTry, Utils} import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.sql.execution.QueryExecution.{runCommand => refreshPhaseEnabledmode} /** * The primary workflow for executing relational queries using Spark. Designed to allow easy diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala index 472d65dfb2b1c..6d874fb29e98c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.connector.metric.CustomMetric import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, DataWriterFactory, DeltaWrite, DeltaWriter, MergeSummaryImpl, PhysicalWriteInfoImpl, Write, WriterCommitMessage, WriteSummary} import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} -import org.apache.spark.sql.execution.{CommandExecutionMode, QueryExecution, SparkPlan, SQLExecution, UnaryExecNode} +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan, SQLExecution, UnaryExecNode} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.execution.metric.{CustomMetrics, SQLMetric, SQLMetrics} import org.apache.spark.sql.types.StructType From 3b206696e8a2e9057cb40d3a26fc655d5fdabe5c Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Thu, 25 Dec 2025 17:18:37 +0800 Subject: [PATCH 4/4] Update DataSource.scala --- .../org/apache/spark/sql/execution/datasources/DataSource.scala | 1 + 1 file changed, 1 insertion(+) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala index be45704fdfc56..35588df11bfca 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/DataSource.scala @@ -39,6 +39,7 @@ import org.apache.spark.sql.classic.ClassicConversions.castToImpl import org.apache.spark.sql.classic.Dataset import org.apache.spark.sql.connector.catalog.TableProvider import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} +import org.apache.spark.sql.execution.QueryExecution import org.apache.spark.sql.execution.command.DataWritingCommand import org.apache.spark.sql.execution.datasources.csv.CSVFileFormat import org.apache.spark.sql.execution.datasources.jdbc.JdbcRelationProvider