Skip to content

[SPARK-53413][SQL] Shuffle cleanup for commands #52157

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
karuppayya wants to merge 5 commits into from
Closed

[SPARK-53413][SQL] Shuffle cleanup for commands #52157

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
616b5e3
SPARK-53413: Shuffle cleanup for commands
karuppayya Aug 18, 2025
602dfb1
Add logs, handle DataWritingCommandExec
karuppayya Aug 29, 2025
6cef7bd
Write test for DataWritingCommandExec
karuppayya Aug 29, 2025
6850ebf
Address review comments: Remove unused code, fix test to write csv
karuppayya Sep 1, 2025
4fa440c
Trigger CI
karuppayya Sep 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
SPARK-53413: Shuffle cleanup for commands 
Fix
  • Loading branch information
@karuppayya
karuppayya committed Sep 1, 2025
commit 616b5e3237cb794cc87e144d830607dafedc6506
2 changes: 1 addition & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ 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 = sparkSession.sessionState.executePlan(p, mode)
val qe = sparkSession.sessionState.executePlan(p, mode, shuffleCleanupMode)
Copy link
Contributor

@cloud-fan cloud-fan Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

instead of calling executePlan, can we construct QueryExecution instance directly? I feel this executePlan method is a bit useless.

Copy link
Contributor Author

@karuppayya karuppayya Aug 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems to have been introduced here. It doesnt mention the rationale behind. But like you mentioned it does nt serve any real purpose, creating new QueryExecution

val result = QueryExecution.withInternalError(s"Eagerly executed $name failed.") {
SQLExecution.withNewExecutionId(qe, Some(name)) {
qe.executedPlan.executeCollect()
Expand Down
23 changes: 16 additions & 7 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import org.apache.spark.internal.config.{SPARK_DRIVER_PREFIX, SPARK_EXECUTOR_PRE
import org.apache.spark.internal.config.Tests.IS_TESTING
import org.apache.spark.sql.classic.SparkSession
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
import org.apache.spark.sql.execution.datasources.v2.V2CommandExec
import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
import org.apache.spark.sql.execution.ui.{SparkListenerSQLExecutionEnd, SparkListenerSQLExecutionStart}
import org.apache.spark.sql.internal.SQLConf
Expand Down Expand Up @@ -68,6 +69,17 @@ object SQLExecution extends Logging {
}
}

private def extractShuffleIds(plan: SparkPlan): Seq[Int] = {
plan match {
case ae: AdaptiveSparkPlanExec =>
ae.context.shuffleIds.asScala.keys.toSeq
case nonAdaptivePlan =>
nonAdaptivePlan.collect {
case exec: ShuffleExchangeLike => exec.shuffleId
}
}
}

/**
* Wrap an action that will execute "queryExecution" to track all Spark jobs in the body so that
* we can connect them with an execution.
Expand Down Expand Up @@ -177,13 +189,10 @@ object SQLExecution extends Logging {
if (queryExecution.shuffleCleanupMode != DoNotCleanup
&& isExecutedPlanAvailable) {
val shuffleIds = queryExecution.executedPlan match {
case ae: AdaptiveSparkPlanExec =>
ae.context.shuffleIds.asScala.keys
case nonAdaptivePlan =>
nonAdaptivePlan.collect {
case exec: ShuffleExchangeLike =>
exec.shuffleId
}
case command: V2CommandExec =>
Copy link
Contributor

@cloud-fan cloud-fan Aug 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how about other commands?

Copy link
Contributor Author

@karuppayya karuppayya Aug 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have handled the DataWritingCommandExec,
I think it might be tricky for ExecutedCommandExec/V1Commands, since the spark plan is created internally.
Should we enumerate each SparkPlan type and handle them?

command.children.flatMap(extractShuffleIds)
case plan =>
extractShuffleIds(plan)
}
shuffleIds.foreach { shuffleId =>
queryExecution.shuffleCleanupMode match {
Expand Down
8 changes: 4 additions & 4 deletions sql/core/src/main/scala/org/apache/spark/sql/internal/BaseSessionStateBuilder.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.classic.{SparkSession, Strategy, StreamingQueryManager, UDFRegistration}
import org.apache.spark.sql.connector.catalog.CatalogManager
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.execution.{ColumnarRule, CommandExecutionMode, QueryExecution, SparkOptimizer, SparkPlanner, SparkSqlParser}
import org.apache.spark.sql.execution.{ColumnarRule, CommandExecutionMode, QueryExecution, ShuffleCleanupMode, SparkOptimizer, SparkPlanner, SparkSqlParser}
import org.apache.spark.sql.execution.adaptive.AdaptiveRulesHolder
import org.apache.spark.sql.execution.aggregate.{ResolveEncodersInScalaAgg, ScalaUDAF}
import org.apache.spark.sql.execution.analysis.DetectAmbiguousSelfJoin
Expand Down Expand Up @@ -406,9 +406,9 @@ abstract class BaseSessionStateBuilder(
* Create a query execution object.
*/
protected def createQueryExecution:
(LogicalPlan, CommandExecutionMode.Value) => QueryExecution =
(plan, mode) => new QueryExecution(session, plan, mode = mode,
shuffleCleanupMode = QueryExecution.determineShuffleCleanupMode(session.sessionState.conf))
(LogicalPlan, CommandExecutionMode.Value, ShuffleCleanupMode) => QueryExecution =
(plan, mode, cleanupMode) => new QueryExecution(session, plan, mode = mode,
shuffleCleanupMode = cleanupMode)

/**
* Interface to start and stop streaming queries.
Expand Down
8 changes: 5 additions & 3 deletions sql/core/src/main/scala/org/apache/spark/sql/internal/SessionState.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,8 @@ private[sql] class SessionState(
val streamingQueryManagerBuilder: () => StreamingQueryManager,
val listenerManager: ExecutionListenerManager,
resourceLoaderBuilder: () => SessionResourceLoader,
createQueryExecution: (LogicalPlan, CommandExecutionMode.Value) => QueryExecution,
createQueryExecution: (LogicalPlan, CommandExecutionMode.Value,
ShuffleCleanupMode) => QueryExecution,
createClone: (SparkSession, SessionState) => SessionState,
val columnarRules: Seq[ColumnarRule],
val adaptiveRulesHolder: AdaptiveRulesHolder,
Expand Down Expand Up @@ -135,8 +136,9 @@ private[sql] class SessionState(

def executePlan(
plan: LogicalPlan,
mode: CommandExecutionMode.Value = CommandExecutionMode.ALL): QueryExecution =
createQueryExecution(plan, mode)
mode: CommandExecutionMode.Value = CommandExecutionMode.ALL,
shuffleCleanUpMode: ShuffleCleanupMode = DoNotCleanup): QueryExecution =
createQueryExecution(plan, mode, shuffleCleanUpMode)
}

private[sql] object SessionState {
Expand Down
18 changes: 17 additions & 1 deletion sql/core/src/test/scala/org/apache/spark/sql/execution/QueryExecutionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ import scala.collection.mutable
import scala.io.Source
import scala.util.Try

import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator}
import org.apache.spark.sql.{AnalysisException, ExtendedExplainGenerator, FastOperator, SaveMode}
import org.apache.spark.sql.catalyst.{QueryPlanningTracker, QueryPlanningTrackerCallback, TableIdentifier}
import org.apache.spark.sql.catalyst.analysis.{CurrentNamespace, UnresolvedFunction, UnresolvedRelation}
import org.apache.spark.sql.catalyst.expressions.{Alias, UnsafeRow}
Expand Down Expand Up @@ -327,6 +327,22 @@ class QueryExecutionSuite extends SharedSparkSession {
}
}

test("SPARK-53413: Cleanup shuffle dependencies for commands") {
Seq(true, false).foreach { adaptiveEnabled => {
withSQLConf((SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, adaptiveEnabled.toString),
(SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED.key, "true")) {
val plan = spark.range(100).repartition(10).logicalPlan
val df = Dataset.ofRows(spark, plan, RemoveShuffleFiles)
df.write.format("noop").mode(SaveMode.Overwrite).save()

val blockManager = spark.sparkContext.env.blockManager
assert(blockManager.migratableResolver.getStoredShuffles().isEmpty)
assert(blockManager.diskBlockManager.getAllBlocks().isEmpty)
}
}
}
}

test("SPARK-47764: Cleanup shuffle dependencies - DoNotCleanup mode") {
Seq(true, false).foreach { adaptiveEnabled => {
withSQLConf((SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, adaptiveEnabled.toString)) {
Expand Down