Skip to content

[SPARK-53469][SQL] Ability to cleanup shuffle in Thrift server #52213

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

Open
wants to merge 6 commits into
base: master
Choose a base branch
from
+17 −7
Open

[SPARK-53469][SQL] Ability to cleanup shuffle in Thrift server #52213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
5103f76
[SPARK-53469] Abiliy to cleanup shuffle in Thrift server
karuppayya Sep 3, 2025
f7c5a78
Fix test
karuppayya Sep 5, 2025
b6914bb
Fix test: IjectRuntimeFilterSuite
karuppayya Sep 6, 2025
2841ec3
Address review comment: add code comment
karuppayya Sep 21, 2025
23176ed
Workaround to make test pass, correct fix is to fix adaptive execution
karuppayya Sep 29, 2025
fee6ad9
Revert testdata changes
karuppayya Oct 14, 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
3 changes: 2 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/classic/SparkSession.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -496,7 +496,8 @@ class SparkSession private(
parsedPlan
}
}
Dataset.ofRows(self, plan, tracker)
Dataset.ofRows(self, plan, tracker,
Copy link
Contributor Author

Choose a reason for hiding this comment

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

SparkExecuteStatementOperation's code eventually leads here.

QueryExecution.determineShuffleCleanupMode(sessionState.conf))
}

/** @inheritdoc */
Expand Down
16 changes: 13 additions & 3 deletions sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,18 @@ class AdaptiveQueryExecSuite

setupTestData()

protected override def beforeAll(): Unit = {
super.beforeAll()
// Tests depend on intermediate results that would otherwise be cleaned up when
Copy link
Contributor

Choose a reason for hiding this comment

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

ditto here, AQE is turned on by default now and let's make sure it's compatible with shuffle cleanup

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I think my wording might not have not been clear.
Thge query execution indeed succeeds and there is no correctness issue.
The issue is with how we assert in tests
In org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite#checkNumLocalShuffleReads

  private def checkNumLocalShuffleReads(
      plan: SparkPlan, numShufflesWithoutLocalRead: Int = 0): Unit = {
    val numShuffles = collect(plan) {
      case s: ShuffleQueryStageExec => s
    }.length

    val numLocalReads = collect(plan) {
      case read: AQEShuffleReadExec if read.isLocalRead => read
    }
    numLocalReads.foreach { r =>
      val rdd = r.execute()
      val parts = rdd.partitions
      assert(parts.forall(rdd.preferredLocations(_).nonEmpty))
    }
    assert(numShuffles === (numLocalReads.length + numShufflesWithoutLocalRead))
  }

Specifically rdd.preferredLocations(_).nonEmpty) will be empty after the cleanup and the assertion fails.
if shuffle clean up is enabled, i dont think this assertion should pass.

Copy link
Contributor

Choose a reason for hiding this comment

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

can we disable cleanup only for that test case?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

This is an assertion method used by all tests.

// shuffle clean up is enabled, causing test failures.
sqlConf.setConf(SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED, false)
Copy link
Contributor

Choose a reason for hiding this comment

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

ditto

Copy link
Contributor Author

Choose a reason for hiding this comment

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

done

}

protected override def afterAll(): Unit = {
sqlConf.setConf(SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED, true)
super.afterAll()
}

private def runAdaptiveAndVerifyResult(query: String,
skipCheckAnswer: Boolean = false): (SparkPlan, SparkPlan) = {
var finalPlanCnt = 0
Expand Down Expand Up @@ -1871,9 +1883,7 @@ class AdaptiveQueryExecSuite
}

withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
SQLConf.SHUFFLE_PARTITIONS.key -> "5",
// Disabling cleanup as the test assertions depend on them
SQLConf.CLASSIC_SHUFFLE_DEPENDENCY_FILE_CLEANUP_ENABLED.key -> "false") {
SQLConf.SHUFFLE_PARTITIONS.key -> "5") {
val df = sql(
"""
|SELECT * FROM (
Expand Down
5 changes: 2 additions & 3 deletions ...e-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkSQLDriver.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.COMMAND
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.plans.logical.CommandResult
import org.apache.spark.sql.classic.ClassicConversions._
import org.apache.spark.sql.execution.{QueryExecution, QueryExecutionException, SQLExecution}
import org.apache.spark.sql.execution.{QueryExecution, QueryExecutionException, SQLExecution}
import org.apache.spark.sql.execution.HiveResult.hiveResultString
import org.apache.spark.sql.internal.{SQLConf, VariableSubstitution}
import org.apache.spark.util.Utils
Expand Down Expand Up @@ -67,7 +66,7 @@ private[hive] class SparkSQLDriver(val sparkSession: SparkSession = SparkSQLEnv.
new VariableSubstitution().substitute(command)
}
sparkSession.sparkContext.setJobDescription(substitutorCommand)
val execution = sparkSession.sessionState.executePlan(sparkSession.sql(command).logicalPlan)
val execution = sparkSession.sql(command).queryExecution
Copy link
Contributor

Choose a reason for hiding this comment

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

So this PR does not enable shuffle cleanup in thriftserver yet?

Copy link
Contributor Author

@karuppayya karuppayya Sep 16, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

I am not sure if i fully understand.
But on why i changed this line
Before: the query execution object was created with the default DonotCleanup, without honoring the shuffle cleanup configs
After: the query execution is created with the DataSet APIs, where the configs are honored when set, giving an opportunity to cleanup generated shuffle after the SQL.
Did i answer your question?

Copy link
Contributor

Choose a reason for hiding this comment

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

Yea, so it's more like a bug fix to make thriftserver to respect this config, but the config is off by default so thriftserver does not do shuffle cleanup by default.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes, thats right

// The SQL command has been executed above via `executePlan`, therefore we don't need to
// wrap it again with a new execution ID when getting Hive result.
execution.logical match {
Expand Down