viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3832927885
##########
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:
Follow-up: I've reworked this since my last reply. Instead of the
identity-only walk, the live set is now mapped through
`NarrowDependency.getParents` (`liveReduceSet`), which threads the
read-partition subset correctly through `OneToOne` (identity),
`RangeDependency` (union's per-branch offset), and join fan-in -- so a
`union`/`join` + `LIMIT` now maps its live set instead of falling back to
all-live. When the narrow chain reaches the shuffle but the mapping is
uncomputable and the read is partial, it fails fast rather than hanging.
Separately, your `coalesce` example turned out to deadlock the transport
regardless of the live set (one reduce task drains several reduce partitions
sequentially against a bounded queue), so a `coalesce` over a shuffle is now
left as a regular shuffle rather than pipelined. Added union/join + LIMIT and
coalesce-fallback tests.
--
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]
