jerrypeng commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4406405959
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1516,6 +1515,29 @@ private[spark] class DAGScheduler(
checkPipelinedStageCreation(finalRDD, partitions, jobId)
}
+ /** Reject persistence that could bypass a native pipelined shuffle reader.
*/
+ private def checkPipelinedCacheSupported(finalRDD: RDD[_]): Unit = {
+ val persisted = new HashSet[RDD[_]]
+ traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
+ if (rdd.getStorageLevel != StorageLevel.NONE ||
+
rdd.checkpointData.exists(_.isInstanceOf[LocalRDDCheckpointData[_]])) {
+ persisted += rdd
+ }
+ rdd.dependencies.foreach(dep => enqueue(dep.rdd))
+ }
+ // Cache locations cannot establish a stable cut: eviction can revive an
upstream producer,
+ // and a partial cache hit skips a reader whose producer still needs its
termination ACK.
+ // Check storage intent instead, including pending local checkpoints and
narrow descendants.
+ val visited = new HashSet[RDD[_]]
+ if (persisted.exists(rdd => rddChainReadsPipelinedShuffle(rdd, visited))) {
Review Comment:
Agreed. `AttachDistributedSequenceExec` creates its temporary cache directly
through RDD
`persist()` or `localCheckpoint()`, so it bypasses `CacheManager`'s
protection. I added
`AttachDistributedSequenceExec(cache = true)` to the shared
`isUnsupportedConsumer` predicate.
Both the AQE and non-AQE rules now retain regular input shuffles, allowing
the indexing job
to execute without hitting the new scheduler rejection.
Using the consumer check preserves pipelining for `cache = false` and for
cached indexing
on shuffle-free input that feeds a downstream shuffle. The regression tests
execute the
indexing operator with AQE on/off, default and memory caching, local
checkpointing, and
`NONE`, including a shuffle above the indexing operator. They verify all
4,000 input IDs
and consecutive index values, and cover both supported pipelined cases.
The restriction deliberately follows the operator's `cache` flag, so
eligibility does not
depend on reading mutable pandas configuration during planning. It
conservatively retains
regular input shuffles for `cache = true` even with `NONE` or a single input
partition,
where execution would not persist. The existing all-or-nothing rules can
also keep other
exchanges in the plan regular. This is a compatibility fix; I have not
benchmarked its
performance impact.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
