viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3816667048
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2604,6 +2687,36 @@ private[spark] class DAGScheduler(
.getOrElse(new Properties())
addPySparkConfigsToProperties(stage, properties)
+ // For a pipelined PRODUCER stage, tell its tasks which reduce partitions
the job actually
+ // reads (the result stage's partitions). The in-process channel writer
drops records
+ // routed to partitions no consumer will drain -- otherwise a partial-read
job (LIMIT /
+ // executeTake reads a subset) fills the unread partitions' bounded queues
and deadlocks
+ // the writer. The result stage is created before submitStage, so its
partitions are known
+ // here. Only the map-side producer needs this; a regular full-read job's
result stage
+ // covers every partition, so the property (if written) lists them all and
drops nothing.
+ // The live set is per-SHUFFLE-EDGE, not per-job: it is the reduce
partitions the
+ // consumer of THIS shuffle reads. It equals the job's result partitions
ONLY for the
+ // producer whose shuffle the result stage reads DIRECTLY (result
partition i maps to
+ // that producer's reduce partition i). A middle pipelined exchange in a
chain (e.g. a
+ // subquery's hash below a single-partition agg) is consumed by another
map stage that
+ // reads ALL its partitions, so it must stay fully live -- setting the
result's subset
+ // there would make it drop partitions the downstream stage still needs,
deadlocking.
+ // So set the property only on the result-feeding producer.
+ stage match {
+ case sms: ShuffleMapStage if isPipelinedProducer(stage) =>
Review Comment:
Fixed. Replaced the reachability check with `readsShuffleByIdentity`, which
follows only 1:1 same-index `OneToOneDependency` hops from the result RDD to
the shuffle. Your `coalesce` example (and a `union` whose shuffle branch is
offset) is a `NarrowDependency`/`RangeDependency`, so the walk stops there, the
live set is left unset (all partitions live), and the operator drains every
reduce partition -- no hang, no dropped rows. The identity chain also covers
the `MapPartitionsRDD` wrapper `executeTake`/`collect` add, which a "direct
dependency only" attempt had wrongly excluded. Added a regression test with the
`coalesce` shape.
--
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]
