dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3816785113
##########
core/src/main/scala/org/apache/spark/ContextCleaner.scala:
##########
@@ -262,7 +262,23 @@ private[spark] class ContextCleaner(
listeners.asScala.foreach(_.shuffleCleaned(shuffleId))
logDebug("Cleaned pipelined shuffle " + shuffleId)
} else {
- logDebug("Asked to cleanup non-existent shuffle (maybe it was already
removed)")
+ // The shuffle is in NEITHER tracker. This is either a genuinely
non-existent shuffle
+ // (already removed) OR an in-process pipelined shuffle whose manager
keeps NO output
+ // tracker
(PipelinedChannelShuffleManager.usesStreamingShuffleOutputTracker = false):
+ // such a shuffle registers with no tracker at all, so the two
branches above miss it,
+ // yet its process-wide rendezvous queues (ChannelShuffleRendezvous)
still need freeing.
+ // Call shuffleDriverComponents.removeShuffle unconditionally: the
RemoveShuffle it issues
+ // routes to SparkEnv.unregisterShuffleFromAllManagers, which reaches
that manager's
+ // unregisterShuffle (-> ChannelShuffleRendezvous.removeShuffle) and
frees the queues.
+ // Safe for a truly non-existent id:
BlockManagerMasterEndpoint.removeShuffle finds no
+ // blocks and unregisterShuffle for an unknown id is a no-op on every
manager. No
+ // shuffle-manager type check here on purpose -- the cleaner stays
transport-agnostic and
+ // just balances the registerShuffle the ShuffleDependency constructor
issues for every
+ // shuffle.
+ logDebug("Cleaning tracker-less shuffle " + shuffleId)
+ shuffleDriverComponents.removeShuffle(shuffleId, blocking)
Review Comment:
This is the one change in the PR that is not behind the opt-in flag, and it
changes behavior for regular shuffles.
The previous code logged and did nothing here. Now every shuffle id that is
in neither tracker issues a `RemoveShuffle` (a broadcast to all executors via
`BlockManagerMasterEndpoint`) and fires `shuffleCleaned` on every registered
`CleanerListener`.
That branch is reached for already-cleaned shuffles, not just tracker-less
pipelined ones. `RDD.cleanShuffleDependencies` (a public `@DeveloperApi`, and
the path `spark.sql.classic.shuffleDependency.fileCleanup.enabled` takes at the
end of a classic SQL execution) calls `doCleanupShuffle` eagerly, which
unregisters from the `MapOutputTracker`. When the `ShuffleDependency` is later
collected, the reference queue calls `doCleanupShuffle` a second time — and now
lands here. So an ordinary deployment with the feature off gets a duplicate
cluster-wide `RemoveShuffle` RPC and a duplicate `shuffleCleaned` callback per
eagerly-cleaned shuffle.
This contradicts "No behavior change by default" in the PR description. The
comment says the absence of a manager-type check is deliberate to keep the
cleaner transport-agnostic, which I understand — but the cost is a change that
reaches users who never enable this. Gating on
`!SparkEnv.get.pipelinedShuffleManager.usesStreamingShuffleOutputTracker`, or
having the channel manager expose the set of shuffle ids it holds queues for,
would keep the new call on the pipelined path only.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,280 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.local.pipelined
+
+import java.util.Arrays
+
+import org.apache.spark.{SparkContext, SparkEnv, TaskContext}
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleReader,
ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+
+/**
+ * Map-side of the in-process pipelined shuffle. Each input record is routed
to the reduce
+ * partition its key hashes to and accumulated in a per-partition batch; a
FULL batch (an
+ * `Array[AnyRef]` of `batchSize` pairs) is pushed onto that partition's
shared queue in one
+ * queue operation (see [[ChannelShuffleRendezvous]]) -- the consumer stage,
running
+ * concurrently, drains it batch by batch. No serialization, no disk, no
network.
+ *
+ * Batching is what makes the transport viable for large unaggregated
shuffles: the queue
+ * costs a lock acquisition per operation (~hundreds of ns under
producer/consumer
+ * contention), so handing rows across one at a time costs that PER ROW --
measured at ~19x
+ * slower than a regular shuffle on a 20M-row repartition. Batching divides
the lock traffic
+ * by `batchSize`, the same lesson as any object-batch transport. A batch
+ * array is handed off to the consumer and never touched again by the writer
(a fresh array
+ * is allocated after each put), so ownership transfer is clean across threads.
+ */
+private[spark] class ChannelShuffleWriter[K, V](
+ handle: BaseShuffleHandle[K, V, _],
+ mapId: Long,
+ batchSize: Int,
+ writeMetrics: ShuffleWriteMetricsReporter)
+ extends ShuffleWriter[K, V] {
+
+ require(batchSize > 0, s"batchSize must be positive, got $batchSize")
+
+ private val dep = handle.dependency
+ private val partitioner = dep.partitioner
+ private val numPartitions = partitioner.numPartitions
+ private val shuffleId = handle.shuffleId
+
+ // The reduce partitions this job actually reads, from the producer stage's
task property
+ // (set by the DAGScheduler from the result stage's partitions). A record
routed to a
+ // partition NOT in this set has no consumer -- putting it would fill that
partition's
+ // bounded queue and, because the writer interleaves all partitions on one
thread, block
+ // the writer before it can feed even the read partitions or emit their
end-of-stream,
+ // deadlocking the job. So such records are dropped. Absent property (None)
means every
+ // partition is live (the normal full-read case: collect, count, a
full-partition job) and
+ // nothing is dropped.
+ private val liveReducePartitions: Option[Set[Int]] =
+ Option(TaskContext.get())
+ .flatMap(tc =>
+
Option(tc.getLocalProperty(SparkContext.SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS)))
+ .map(_.split(",").filter(_.nonEmpty).map(_.toInt).toSet)
+
+ // Per-partition liveness, precomputed ONCE from the (static) live set: true
iff a consumer
+ // reads this reduce partition at all. This is the hot-path gate -- checked
per input record --
+ // so it is a plain Array[Boolean] load, not a boxed Set lookup: on a large
repartition the
+ // per-record path must not allocate (the transport's whole point is
amortizing per-row cost).
+ // Absent property means every partition is live. The OTHER half of "worth
writing" --
+ // abandonment, which happens at runtime when a reader departs early (e.g.
LIMIT) -- is dynamic
+ // and is checked where it matters (at hand-off, in putUnlessAbandoned), NOT
per record:
+ // accumulating a few more rows into an in-memory batch for a
since-abandoned partition is
+ // harmless because that batch is never put (putUnlessAbandoned drops it).
+ private val liveMask: Array[Boolean] = {
+ val mask = Array.fill(numPartitions)(true)
+ liveReducePartitions.foreach { live =>
+ var p = 0
+ while (p < numPartitions) { mask(p) = live.contains(p); p += 1 }
+ }
+ mask
+ }
+
+ // Hand a batch to a partition's queue, but do NOT block forever if its
reader departs:
+ // poll with a short timeout and bail out the moment the partition becomes
abandoned. This
+ // is the cooperative unblock for the early-stop case -- abandon() also
drains the queue to
+ // release a parked put, and this re-check ensures the writer then stops
rather than
+ // re-filling. Returns false if the partition was abandoned before the batch
was accepted.
+ // On a successful hand-off, records the batch's records and the time spent
(including any
+ // backpressure wait) against the write metrics; a dropped/abandoned batch
counts nothing,
+ // since those records are never shuffled out. `records` is the number of
pairs in `batch`
+ // (a full batch is `batchSize`, a trimmed tail is shorter; the
end-of-stream marker is 0).
+ private def putUnlessAbandoned(pid: Int, batch: AnyRef, records: Int):
Boolean = {
+ val q = ChannelShuffleRendezvous.queue(shuffleId, pid)
+ val start = System.nanoTime()
+ while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)) {
+ if (q.offer(batch, 100, java.util.concurrent.TimeUnit.MILLISECONDS)) {
+ // A successful offer can race abandon(): abandon does `add(mark)`
then `q.clear()`, so if
+ // it ran between the isAbandoned check above and this offer, our
batch lands AFTER the
+ // clear and would be stranded in the queue (no reader will ever drain
it). Re-check and
+ // clear it ourselves so nothing is left behind. The reader has
departed, so discarding is
+ // correct; and it keeps the queue empty for removeShuffle rather than
pinning a batch.
+ if (ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)) {
+ q.clear()
+ return false
+ }
+ if (records > 0) {
+ writeMetrics.incRecordsWritten(records.toLong)
+ writeMetrics.incWriteTime(System.nanoTime() - start)
+ }
+ return true
+ }
+ }
+ false
+ }
+
+ override def write(records: Iterator[Product2[K, V]]): Unit = {
+ // Abandoned marks left by an EARLIER run of this shuffleId (a
RangePartitioner sampling job
+ // then the main job; executeTake batches) are reset by the DAGScheduler
when it submits this
+ // producer stage -- before any map task of this run starts (see
+ // ChannelShuffleRendezvous.clearAbandonedForShuffle). The writer must NOT
clear them itself:
+ // map tasks of the same run are concurrent, and a late one clearing a
mark would erase a
+ // departure a sibling's reader had already recorded for this run,
re-hanging the writer.
+
+ // One in-progress batch per reduce partition, plus its fill count.
+ val batches = Array.fill(numPartitions)(new Array[AnyRef](batchSize))
+ val sizes = new Array[Int](numPartitions)
+
+ while (records.hasNext) {
+ val rec = records.next()
+ val pid = partitioner.getPartition(rec._1)
+ // Only accumulate for partitions a consumer reads (liveMask).
Abandonment is not checked
+ // here -- it is handled at hand-off in putUnlessAbandoned (see
liveMask's comment).
+ if (liveMask(pid)) {
+ // Records must already be detached from the producer's reused row
buffers by the time
+ // they reach here (the producer reuses its output UnsafeRow across
iterations, and the
+ // consumer reads on another thread). The copy is done in the SQL
layer's
+ // ShuffleWriteProcessor for the pipelined path -- where
InternalRow.copy() is available
+ // -- rather than here, because this class lives in `core` and cannot
reference SQL rows,
+ // and the UnsafeRow serializer offers no single-object copy. So batch
the pair as-is.
+ batches(pid)(sizes(pid)) = (rec._1, rec._2)
+ sizes(pid) += 1
+ if (sizes(pid) == batchSize) {
+ putUnlessAbandoned(pid, batches(pid), batchSize)
+ batches(pid) = new Array[AnyRef](batchSize)
+ sizes(pid) = 0
+ }
+ }
+ }
+
+ // Flush partial batches (trimmed so the reader can iterate array length
directly), then
+ // signal end-of-stream to every partition still wanted, so each live
reader can count
+ // this map task as done. Same thread, same queue: data always precedes
the marker. A
+ // partition that is dead (no reader) or abandoned (reader departed) gets
neither -- its
+ // queue is left for removeShuffle to drop.
+ var p = 0
+ while (p < numPartitions) {
+ if (liveMask(p)) {
+ if (sizes(p) > 0) {
+ putUnlessAbandoned(p, Arrays.copyOf(batches(p), sizes(p)), sizes(p))
+ }
+ // Re-check: the reader may have departed while the trimmed batch was
being put.
+ if (!ChannelShuffleRendezvous.isAbandoned(shuffleId, p)) {
+ putUnlessAbandoned(p, ChannelShuffleRendezvous.EndOfStream, records
= 0)
+ }
+ }
+ p += 1
+ }
+ }
+
+ override def stop(success: Boolean): Option[MapStatus] = {
+ // A pipelined reducer never reads partition lengths, but the
ShuffleWriter contract
+ // still requires a MapStatus. Return an all-zero placeholder, mirroring
the RPC
+ // streaming writer.
+ Some(MapStatus(
+ SparkEnv.get.blockManager.shuffleServerId,
+ Array.fill(numPartitions)(0L),
+ mapId))
+ }
+
+ override def getPartitionLengths(): Array[Long] =
Array.fill(numPartitions)(0L)
+}
+
+/**
+ * Reduce-side of the in-process pipelined shuffle. Drains the shared queue
for this reduce
+ * partition batch by batch, handing rows to the consumer stage as the map
tasks produce
+ * them, until every map task has signalled end-of-stream.
+ *
+ * `numMaps` is the number of map tasks feeding this shuffle; the reader stops
after it has
+ * observed that many [[ChannelShuffleRendezvous.EndOfStream]] markers on its
queue.
+ *
+ * ONE reduce partition per reader task ONLY: `endPartition - startPartition`
must be 1. The
+ * channel transport cannot serve a coalesced multi-partition range. The
reader would have to
+ * drain the range's queues in some order, but the map-side writer interleaves
all partitions
+ * on ONE thread and blocks on a full bounded queue; if the reader drains
partition `start` to
+ * completion before touching `start+1` while the writer has parked filling
`start+1`, the two
+ * deadlock with no timeout escape. Spark never sends a coalesced spec here
today -- AQE keeps a
+ * pipelined exchange out of any ShuffleQueryStage, so
CoalesceShufflePartitions never coalesces
+ * it, and both the AQE and non-AQE readers use width-1
CoalescedPartitionSpec(i, i+1). The
+ * `require` below makes that a hard, fail-loud invariant rather than a silent
hang if a future
+ * change ever lets a coalesced spec reach a pipelined dependency.
+ */
+private[spark] class ChannelShuffleReader[K, C](
+ handle: BaseShuffleHandle[K, _, C],
+ startPartition: Int,
+ endPartition: Int,
+ numMaps: Int,
+ readMetrics: ShuffleReadMetricsReporter)
+ extends ShuffleReader[K, C] {
+
+ require(endPartition - startPartition == 1,
+ s"ChannelShuffleReader supports exactly one reduce partition per task, got
" +
+ s"[$startPartition, $endPartition); the in-process channel transport
does not support " +
+ "coalesced multi-partition reads (see class doc).")
+
+ // On task completion (normal end, early stop like LIMIT, or failure) mark
this reader's
+ // partition as abandoned, so a writer still feeding it stops and does not
wedge on its bounded
+ // queue. Registered once here; fires whether or not the iterator was
drained to the end.
+ Option(TaskContext.get()).foreach { tc =>
+ tc.addTaskCompletionListener[Unit] { _ =>
+ var p = startPartition
+ while (p < endPartition) {
+ ChannelShuffleRendezvous.abandon(handle.shuffleId, p)
+ p += 1
+ }
+ }
+ }
+
+ override def read(): Iterator[Product2[K, C]] =
+ (startPartition until endPartition).iterator.flatMap(drainQueue)
+
+ private def drainQueue(reducePartitionId: Int): Iterator[Product2[K, C]] = {
+ val q = ChannelShuffleRendezvous.queue(handle.shuffleId, reducePartitionId)
+ new Iterator[Product2[K, C]] {
+ // The current batch being handed out, and the cursor into it. A null
batch after
+ // advance() means every map task has signalled end-of-stream: iteration
is over.
+ private var batch: Array[AnyRef] = _
+ private var pos = 0
+ private var endOfStreamSeen = 0
+ advance()
+
+ // Blocking-drain until the next non-empty data batch, or until every
map task has
+ // signalled end-of-stream for this queue (then leave `batch` null to
end iteration).
+ private def advance(): Unit = {
+ batch = null
+ pos = 0
+ // A producer with zero map tasks (numMaps == 0, e.g. a pipelined
shuffle over an empty
+ // RDD) enqueues nothing and no end-of-stream marker ever arrives;
without this guard the
+ // take() below would block forever. Terminate immediately with an
empty iterator. The
+ // check is also correct for numMaps > 0 once every marker has been
seen (advance is not
+ // called again after batch stays null, but this keeps the invariant
explicit).
+ if (endOfStreamSeen >= numMaps) return
+ var item = q.take()
Review Comment:
`take()` here is uninterruptible, so a reader whose producer dies never
wakes up.
If a map task fails mid-`write()`, it emits no `EndOfStream` for the
partitions it had not flushed yet. The reduce task for such a partition parks
in this `take()`. The DAGScheduler aborts the group and kills the reduce task,
but `killAllTaskAttempts` passes `shouldInterruptTaskThread(job)`, which reads
`spark.job.interruptOnCancel` and defaults to `false`. The thread is never
interrupted, so it stays parked here.
The job itself does fail and the user sees the error, so this is not a
silent hang — but the reduce task thread keeps its executor slot for the
lifetime of the application. In `local[N]` that permanently costs one core; a
few producer failures in a session exhaust the pool and every later pipelined
query then fails gang admission.
The writer already has the cooperative escape for exactly this class of
problem — `putUnlessAbandoned` polls with a 100ms timeout and re-checks a flag
instead of blocking forever. The reader needs the symmetric treatment, e.g.
```scala
var item: AnyRef = null
while (item == null) {
Option(TaskContext.get()).foreach(_.killTaskIfInterrupted())
item = q.poll(100, TimeUnit.MILLISECONDS)
}
```
`killTaskIfInterrupted` works without thread interruption, since the
executor sets the interrupted flag on `TaskContext` on kill either way.
Worth a regression test where a producer map task throws:
`PipelinedLimitHangSuite` covers the early-stop (normal) path but nothing
covers producer failure.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala:
##########
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.local.pipelined
+
+import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue}
+
+/**
+ * Process-wide rendezvous between the map (writer) and reduce (reader) sides
of an
+ * in-process pipelined shuffle. One bounded queue exists per
+ * `(shuffleId, reducePartitionId)`; every map task writing to a given reduce
partition
+ * shares the queue with the single reduce task that drains it. Queue elements
are BATCHES
+ * of records (`Array[AnyRef]` of pairs, see [[ChannelShuffleWriter]]) or the
+ * [[EndOfStream]] marker, so the queue's per-operation lock cost is paid per
batch, not
+ * per row.
+ *
+ * This is correct only when producer and consumer tasks are co-resident in
the same JVM,
+ * i.e. a single executor (local mode). The concurrent-stage scheduler
co-schedules the two
+ * stages so both are running, but it does not by itself guarantee
co-location; the
+ * [[PipelinedChannelShuffleManager]] is only intended for single-executor
deployments,
+ * where co-location is automatic. Cross-executor pipelined shuffle is served
by the RPC
+ * streaming shuffle instead.
+ *
+ * The queues are bounded, so a fast producer blocks on `put` when the
consumer lags --
+ * this is the backpressure that keeps the pipelined hand-off memory-bounded.
+ */
+private[spark] object ChannelShuffleRendezvous {
+
+ /**
+ * Marker placed on a queue by each map task when it finishes writing to
that reduce
+ * partition. A reader stops once it has seen one marker per map task.
+ */
+ val EndOfStream: AnyRef = new AnyRef
+
+ // Keyed by (shuffleId, reducePartitionId). Values are AnyRef because the
queue carries
+ // both record batches (Array[AnyRef]) and the EndOfStream marker.
+ private val queues =
+ new ConcurrentHashMap[(Int, Int), LinkedBlockingQueue[AnyRef]]()
+
+ // (shuffleId, reducePartitionId) keys whose reader has departed (its reduce
task finished)
+ // and will drain no more. A writer stops feeding an abandoned partition and
drops the rest.
+ // This covers the LIVE-partition early-stop case (e.g. a LIMIT reader that
pulled enough
+ // and quit): without it the writer fills the partition's bounded queue and
blocks forever.
+ private val abandoned =
+ java.util.concurrent.ConcurrentHashMap.newKeySet[(Int, Int)]()
+
+ /**
+ * Default per-queue capacity in BATCHES (not rows): with the default
1024-row batch this
+ * bounds the in-flight hand-off at ~64K rows per reduce partition.
+ */
+ private val DefaultCapacity = 64
Review Comment:
`batchSize` is a registered conf but the queue depth that multiplies it is a
hard-coded constant, and the two together decide how much heap the transport
pins.
Worst-case in-flight residency for one shuffle is `numPartitions * 64 *
batchSize` rows, held as strong references in this `object` and invisible to
the memory manager (no `MemoryConsumer`, no spill). Gang admission bounds
`numPartitions` in practice, so this is not the multi-million-row figure the
raw arithmetic suggests, but it is still tens of MB of live rows that a user
tuning `spark.shuffle.pipelined.channel.batchSize` upward has no way to reason
about — the doc string mentions "per-partition buffering" without saying it is
multiplied by 64.
Either promote the capacity to a conf next to `batchSize`, or state the `64
* batchSize * numPartitions` product in that conf's doc.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,280 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.local.pipelined
+
+import java.util.Arrays
+
+import org.apache.spark.{SparkContext, SparkEnv, TaskContext}
+import org.apache.spark.scheduler.MapStatus
+import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleReader,
ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter}
+
+/**
+ * Map-side of the in-process pipelined shuffle. Each input record is routed
to the reduce
+ * partition its key hashes to and accumulated in a per-partition batch; a
FULL batch (an
+ * `Array[AnyRef]` of `batchSize` pairs) is pushed onto that partition's
shared queue in one
+ * queue operation (see [[ChannelShuffleRendezvous]]) -- the consumer stage,
running
+ * concurrently, drains it batch by batch. No serialization, no disk, no
network.
+ *
+ * Batching is what makes the transport viable for large unaggregated
shuffles: the queue
+ * costs a lock acquisition per operation (~hundreds of ns under
producer/consumer
+ * contention), so handing rows across one at a time costs that PER ROW --
measured at ~19x
+ * slower than a regular shuffle on a 20M-row repartition. Batching divides
the lock traffic
+ * by `batchSize`, the same lesson as any object-batch transport. A batch
+ * array is handed off to the consumer and never touched again by the writer
(a fresh array
+ * is allocated after each put), so ownership transfer is clean across threads.
+ */
+private[spark] class ChannelShuffleWriter[K, V](
+ handle: BaseShuffleHandle[K, V, _],
+ mapId: Long,
+ batchSize: Int,
+ writeMetrics: ShuffleWriteMetricsReporter)
+ extends ShuffleWriter[K, V] {
+
+ require(batchSize > 0, s"batchSize must be positive, got $batchSize")
+
+ private val dep = handle.dependency
+ private val partitioner = dep.partitioner
+ private val numPartitions = partitioner.numPartitions
+ private val shuffleId = handle.shuffleId
+
+ // The reduce partitions this job actually reads, from the producer stage's
task property
+ // (set by the DAGScheduler from the result stage's partitions). A record
routed to a
+ // partition NOT in this set has no consumer -- putting it would fill that
partition's
+ // bounded queue and, because the writer interleaves all partitions on one
thread, block
+ // the writer before it can feed even the read partitions or emit their
end-of-stream,
+ // deadlocking the job. So such records are dropped. Absent property (None)
means every
+ // partition is live (the normal full-read case: collect, count, a
full-partition job) and
+ // nothing is dropped.
+ private val liveReducePartitions: Option[Set[Int]] =
+ Option(TaskContext.get())
+ .flatMap(tc =>
+
Option(tc.getLocalProperty(SparkContext.SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS)))
+ .map(_.split(",").filter(_.nonEmpty).map(_.toInt).toSet)
+
+ // Per-partition liveness, precomputed ONCE from the (static) live set: true
iff a consumer
+ // reads this reduce partition at all. This is the hot-path gate -- checked
per input record --
+ // so it is a plain Array[Boolean] load, not a boxed Set lookup: on a large
repartition the
+ // per-record path must not allocate (the transport's whole point is
amortizing per-row cost).
+ // Absent property means every partition is live. The OTHER half of "worth
writing" --
+ // abandonment, which happens at runtime when a reader departs early (e.g.
LIMIT) -- is dynamic
+ // and is checked where it matters (at hand-off, in putUnlessAbandoned), NOT
per record:
+ // accumulating a few more rows into an in-memory batch for a
since-abandoned partition is
+ // harmless because that batch is never put (putUnlessAbandoned drops it).
+ private val liveMask: Array[Boolean] = {
+ val mask = Array.fill(numPartitions)(true)
+ liveReducePartitions.foreach { live =>
+ var p = 0
+ while (p < numPartitions) { mask(p) = live.contains(p); p += 1 }
+ }
+ mask
+ }
+
+ // Hand a batch to a partition's queue, but do NOT block forever if its
reader departs:
+ // poll with a short timeout and bail out the moment the partition becomes
abandoned. This
+ // is the cooperative unblock for the early-stop case -- abandon() also
drains the queue to
+ // release a parked put, and this re-check ensures the writer then stops
rather than
+ // re-filling. Returns false if the partition was abandoned before the batch
was accepted.
+ // On a successful hand-off, records the batch's records and the time spent
(including any
+ // backpressure wait) against the write metrics; a dropped/abandoned batch
counts nothing,
+ // since those records are never shuffled out. `records` is the number of
pairs in `batch`
+ // (a full batch is `batchSize`, a trimmed tail is shorter; the
end-of-stream marker is 0).
+ private def putUnlessAbandoned(pid: Int, batch: AnyRef, records: Int):
Boolean = {
+ val q = ChannelShuffleRendezvous.queue(shuffleId, pid)
+ val start = System.nanoTime()
+ while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)) {
+ if (q.offer(batch, 100, java.util.concurrent.TimeUnit.MILLISECONDS)) {
+ // A successful offer can race abandon(): abandon does `add(mark)`
then `q.clear()`, so if
+ // it ran between the isAbandoned check above and this offer, our
batch lands AFTER the
+ // clear and would be stranded in the queue (no reader will ever drain
it). Re-check and
+ // clear it ourselves so nothing is left behind. The reader has
departed, so discarding is
+ // correct; and it keeps the queue empty for removeShuffle rather than
pinning a batch.
+ if (ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)) {
+ q.clear()
+ return false
+ }
+ if (records > 0) {
+ writeMetrics.incRecordsWritten(records.toLong)
+ writeMetrics.incWriteTime(System.nanoTime() - start)
+ }
+ return true
+ }
+ }
+ false
+ }
+
+ override def write(records: Iterator[Product2[K, V]]): Unit = {
+ // Abandoned marks left by an EARLIER run of this shuffleId (a
RangePartitioner sampling job
+ // then the main job; executeTake batches) are reset by the DAGScheduler
when it submits this
+ // producer stage -- before any map task of this run starts (see
+ // ChannelShuffleRendezvous.clearAbandonedForShuffle). The writer must NOT
clear them itself:
+ // map tasks of the same run are concurrent, and a late one clearing a
mark would erase a
+ // departure a sibling's reader had already recorded for this run,
re-hanging the writer.
+
+ // One in-progress batch per reduce partition, plus its fill count.
+ val batches = Array.fill(numPartitions)(new Array[AnyRef](batchSize))
+ val sizes = new Array[Int](numPartitions)
+
+ while (records.hasNext) {
+ val rec = records.next()
+ val pid = partitioner.getPartition(rec._1)
+ // Only accumulate for partitions a consumer reads (liveMask).
Abandonment is not checked
+ // here -- it is handled at hand-off in putUnlessAbandoned (see
liveMask's comment).
+ if (liveMask(pid)) {
+ // Records must already be detached from the producer's reused row
buffers by the time
+ // they reach here (the producer reuses its output UnsafeRow across
iterations, and the
+ // consumer reads on another thread). The copy is done in the SQL
layer's
+ // ShuffleWriteProcessor for the pipelined path -- where
InternalRow.copy() is available
+ // -- rather than here, because this class lives in `core` and cannot
reference SQL rows,
+ // and the UnsafeRow serializer offers no single-object copy. So batch
the pair as-is.
+ batches(pid)(sizes(pid)) = (rec._1, rec._2)
+ sizes(pid) += 1
+ if (sizes(pid) == batchSize) {
+ putUnlessAbandoned(pid, batches(pid), batchSize)
+ batches(pid) = new Array[AnyRef](batchSize)
+ sizes(pid) = 0
+ }
+ }
+ }
+
+ // Flush partial batches (trimmed so the reader can iterate array length
directly), then
+ // signal end-of-stream to every partition still wanted, so each live
reader can count
+ // this map task as done. Same thread, same queue: data always precedes
the marker. A
+ // partition that is dead (no reader) or abandoned (reader departed) gets
neither -- its
+ // queue is left for removeShuffle to drop.
+ var p = 0
+ while (p < numPartitions) {
+ if (liveMask(p)) {
+ if (sizes(p) > 0) {
+ putUnlessAbandoned(p, Arrays.copyOf(batches(p), sizes(p)), sizes(p))
+ }
+ // Re-check: the reader may have departed while the trimmed batch was
being put.
+ if (!ChannelShuffleRendezvous.isAbandoned(shuffleId, p)) {
+ putUnlessAbandoned(p, ChannelShuffleRendezvous.EndOfStream, records
= 0)
+ }
+ }
+ p += 1
+ }
+ }
+
+ override def stop(success: Boolean): Option[MapStatus] = {
+ // A pipelined reducer never reads partition lengths, but the
ShuffleWriter contract
+ // still requires a MapStatus. Return an all-zero placeholder, mirroring
the RPC
+ // streaming writer.
+ Some(MapStatus(
+ SparkEnv.get.blockManager.shuffleServerId,
+ Array.fill(numPartitions)(0L),
+ mapId))
+ }
+
+ override def getPartitionLengths(): Array[Long] =
Array.fill(numPartitions)(0L)
+}
+
+/**
+ * Reduce-side of the in-process pipelined shuffle. Drains the shared queue
for this reduce
+ * partition batch by batch, handing rows to the consumer stage as the map
tasks produce
+ * them, until every map task has signalled end-of-stream.
+ *
+ * `numMaps` is the number of map tasks feeding this shuffle; the reader stops
after it has
+ * observed that many [[ChannelShuffleRendezvous.EndOfStream]] markers on its
queue.
+ *
+ * ONE reduce partition per reader task ONLY: `endPartition - startPartition`
must be 1. The
+ * channel transport cannot serve a coalesced multi-partition range. The
reader would have to
+ * drain the range's queues in some order, but the map-side writer interleaves
all partitions
+ * on ONE thread and blocks on a full bounded queue; if the reader drains
partition `start` to
+ * completion before touching `start+1` while the writer has parked filling
`start+1`, the two
+ * deadlock with no timeout escape. Spark never sends a coalesced spec here
today -- AQE keeps a
+ * pipelined exchange out of any ShuffleQueryStage, so
CoalesceShufflePartitions never coalesces
+ * it, and both the AQE and non-AQE readers use width-1
CoalescedPartitionSpec(i, i+1). The
+ * `require` below makes that a hard, fail-loud invariant rather than a silent
hang if a future
+ * change ever lets a coalesced spec reach a pipelined dependency.
+ */
+private[spark] class ChannelShuffleReader[K, C](
+ handle: BaseShuffleHandle[K, _, C],
+ startPartition: Int,
+ endPartition: Int,
+ numMaps: Int,
+ readMetrics: ShuffleReadMetricsReporter)
+ extends ShuffleReader[K, C] {
+
+ require(endPartition - startPartition == 1,
+ s"ChannelShuffleReader supports exactly one reduce partition per task, got
" +
+ s"[$startPartition, $endPartition); the in-process channel transport
does not support " +
+ "coalesced multi-partition reads (see class doc).")
+
+ // On task completion (normal end, early stop like LIMIT, or failure) mark
this reader's
+ // partition as abandoned, so a writer still feeding it stops and does not
wedge on its bounded
+ // queue. Registered once here; fires whether or not the iterator was
drained to the end.
+ Option(TaskContext.get()).foreach { tc =>
+ tc.addTaskCompletionListener[Unit] { _ =>
+ var p = startPartition
+ while (p < endPartition) {
+ ChannelShuffleRendezvous.abandon(handle.shuffleId, p)
+ p += 1
+ }
+ }
+ }
+
+ override def read(): Iterator[Product2[K, C]] =
+ (startPartition until endPartition).iterator.flatMap(drainQueue)
+
+ private def drainQueue(reducePartitionId: Int): Iterator[Product2[K, C]] = {
+ val q = ChannelShuffleRendezvous.queue(handle.shuffleId, reducePartitionId)
+ new Iterator[Product2[K, C]] {
+ // The current batch being handed out, and the cursor into it. A null
batch after
+ // advance() means every map task has signalled end-of-stream: iteration
is over.
+ private var batch: Array[AnyRef] = _
+ private var pos = 0
+ private var endOfStreamSeen = 0
+ advance()
+
+ // Blocking-drain until the next non-empty data batch, or until every
map task has
+ // signalled end-of-stream for this queue (then leave `batch` null to
end iteration).
+ private def advance(): Unit = {
+ batch = null
+ pos = 0
+ // A producer with zero map tasks (numMaps == 0, e.g. a pipelined
shuffle over an empty
+ // RDD) enqueues nothing and no end-of-stream marker ever arrives;
without this guard the
+ // take() below would block forever. Terminate immediately with an
empty iterator. The
+ // check is also correct for numMaps > 0 once every marker has been
seen (advance is not
+ // called again after batch stays null, but this keeps the invariant
explicit).
+ if (endOfStreamSeen >= numMaps) return
+ var item = q.take()
+ while (item eq ChannelShuffleRendezvous.EndOfStream) {
+ endOfStreamSeen += 1
+ if (endOfStreamSeen >= numMaps) return
+ item = q.take()
+ }
+ batch = item.asInstanceOf[Array[AnyRef]]
+ // Count the records handed to the consumer as this batch is fetched.
Local, so this
+ // is the read-side records metric; there is no remote fetch and no
wire bytes.
+ readMetrics.incRecordsRead(batch.length.toLong)
Review Comment:
Only `incRecordsRead` is reported on the read side — there is no
`incFetchWaitTime`, so the reader's blocking wait on the channel never shows up.
For this transport that wait is the interesting read-side number: it is the
backpressure signal that tells a user whether the consumer or the producer is
the bottleneck, which is the main thing a batch-size or queue-depth tuning
decision turns on. The class doc for the transport and the PR description both
say "records + read/write time are reported to the shuffle metrics", which
currently holds for the writer (`incWriteTime` in `putUnlessAbandoned`) but not
for the reader.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleBenchmark.scala:
##########
@@ -0,0 +1,452 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.exchange
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.benchmark.SqlBasedBenchmark
+
+/**
+ * Benchmark to compare the in-process pipelined channel shuffle against the
regular
+ * (materializing) shuffle on simple batch queries (SPARK-57399).
+ *
+ * Fair-comparison constraints, read before trusting any number:
+ * - Runs on `local[N]` with N = physical cores, and only queries whose
pipelined
+ * whole-group slot demand is <= N, so the concurrently-scheduled
pipelined stages do NOT
+ * oversubscribe the cores. A demand > cores run would measure thread
thrash, not the
+ * transport, and is intentionally avoided here.
+ * - Pipelined overlaps map+reduce stages (uses more concurrent slots) vs
the baseline's
+ * sequential map-then-reduce; with demand <= cores neither is
slot-limited, so the delta
+ * reflects stage overlap minus channel overhead, which is the honest
comparison.
+ *
+ * To run this benchmark:
+ * {{{
+ * 1. build/sbt "sql/Test/runMain
+ * org.apache.spark.sql.execution.exchange.PipelinedShuffleBenchmark"
+ * 2. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt
+ * "sql/Test/runMain
+ * org.apache.spark.sql.execution.exchange.PipelinedShuffleBenchmark"
+ * Results will be written to
"benchmarks/PipelinedShuffleBenchmark-results.txt".
+ * }}}
+ */
+object PipelinedShuffleBenchmark extends SqlBasedBenchmark {
+
+ override def getSparkSession: SparkSession = {
+ SparkSession.builder()
+ .master("local[1]")
+ .appName(this.getClass.getCanonicalName)
+ .config("spark.ui.enabled", "false")
+ .getOrCreate()
+ }
+
+ private val cores = Runtime.getRuntime.availableProcessors()
+ private val numRows = 20000000L // 20M: large enough that transport cost
dominates startup
+ private val inputParts = 6 // demand = 6 (input) + 8 (shuffle) = 14 <=
16 cores
Review Comment:
`inputParts` is pinned to a 16-core assumption while `cores` right above it
is read from the machine, so the benchmark is not portable.
On anything below 14 cores the pipelined configurations do not produce
slower numbers — they fail gang admission with
`CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT` and the run dies. Since the whole
point of the fair-comparison preamble is that demand stays under the core
count, deriving `inputParts` from `cores` (and skipping, with a clear message,
when even the minimum shape does not fit) would make the checked-in results
reproducible outside the machine they were generated on.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -570,7 +570,12 @@ object ShuffleExchangeExec {
rddWithPartitionIds,
new PartitionIdPassthrough(part.numPartitions),
serializer,
- shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics),
+ // Copy rows only for a transport that hands object references to a
concurrent
+ // consumer (the in-process channel). The RPC streaming transport
detaches rows by
+ // serializing them promptly and must not pay an extra per-row copy
on its path.
+ shuffleWriterProcessor = createShuffleWriteProcessor(
+ writeMetrics,
+ copyRows =
SparkEnv.get.pipelinedShuffleManager.requiresDetachedRecords),
Review Comment:
Two small redundancies on the pipelined write path:
1. When `needToCopyObjectsBeforeShuffle(part)` is true,
`prepareShuffleDependency` already emits `(pid, row.copy())`, and this then
copies every row a second time. `copyRows = requiresDetachedRecords &&
!needToCopyObjectsBeforeShuffle(part)` avoids that.
2. The write processor builds a fresh `Tuple2` per record, and
`ChannelShuffleWriter.write` immediately builds another one
(`batches(pid)(sizes(pid)) = (rec._1, rec._2)`). Given `requiresDetachedRecords
= true` is the contract that the pair reaching the writer is already detached,
the re-wrap in the writer is redundant — storing `rec` directly removes one
allocation per record from the hot loop the batching design exists to protect.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2604,6 +2687,60 @@ 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) =>
+ // Notify the pipelined manager that this producer stage is being
(re)submitted, before
+ // any of its map tasks start. A transport with per-run state keyed by
shuffleId resets it
+ // here -- the one point with no live task of the new run, so the
reset cannot race the
+ // run's own writers/readers. The in-process channel transport clears
a prior run's
+ // abandoned-partition marks; the RPC streaming manager keeps no such
state (no-op default).
+ // The scheduler stays transport-agnostic: it calls the
PipelinedShuffleManager trait, not
+ // a concrete transport.
+ SparkEnv.get.pipelinedShuffleManager
+ .onPipelinedProducerStageSubmit(sms.shuffleDep.shuffleId)
+ jobIdToActiveJob.get(jobId).map(_.finalStage).collect { case rs:
ResultStage => rs }
+ .foreach { rs =>
+ // The live set is the result stage's partition ids. Those are the
shuffle's reduce
+ // partition ids -- so partition i means reduce partition i --
ONLY when the result RDD
+ // reaches this shuffle through an IDENTITY-PRESERVING chain:
every hop a 1:1,
+ // same-index OneToOneDependency (e.g. the mapPartitions wrapper
executeTake/collect
+ // adds over the ShuffledRowRDD, which preserves partition count
and index). It is NOT
+ // enough for the shuffle to be merely reachable: a narrow
operator that REMAPS
+ // partitions (coalesce -> a custom NarrowDependency, union -> a
RangeDependency offset)
+ // makes result partition ids differ from reduce partition ids,
and treating them as
+ // the live set would drop partitions a downstream operator still
pulls -- hanging the
+ // reader. When the chain is not identity-preserving the property
is left unset (all
+ // partitions live) and that operator drains every reduce
partition.
+ def readsShuffleByIdentity(rdd: RDD[_]): Boolean =
rdd.dependencies match {
Review Comment:
The identity check is the right shape, but RDD-level identity does not by
itself imply partition-index identity at the one place it matters.
The RDD matched by the first case is `ShuffledRowRDD`, which is precisely
the RDD that maps its partition `i` to an arbitrary reducer index through
`partitionSpecs`. With the default constructor the specs are
`CoalescedPartitionSpec(i, i + 1)` and the identity holds, but that is assumed
here rather than established.
`ChannelShuffleReader`'s `require(endPartition - startPartition == 1)`
guards the *width* of a spec, not its *offset*. A `CoalescedPartitionSpec(5,
6)` sitting at result index 0 passes that `require`, while the live set
computed here would be `{0}` — the writer drops everything routed to reduce
partition 5 and emits no end-of-stream for it, and the reader for partition 5
blocks.
Nothing produces that today (AQE keeps a pipelined exchange out of
`ShuffleQueryStageExec`, so `getShuffleRDD(specs)` is never reached for one),
so this is defense-in-depth rather than a live bug. But since the failure mode
is a hang and the guard is cheap, it would be worth also requiring
`rs.rdd.partitions.length == sd.partitioner.numPartitions` before writing the
property, so a future change that opens the spec path degrades to "all
partitions live" instead of hanging.
##########
core/src/test/scala/org/apache/spark/shuffle/PipelinedShuffleRoutingSuite.scala:
##########
@@ -85,7 +85,13 @@ private class DefaultRecordingManager(conf: SparkConf,
isDriver: Boolean)
override def shuffleBlockResolver: ShuffleBlockResolver =
mock(classOf[ShuffleBlockResolver])
}
private class IncrementalRecordingManager(conf: SparkConf, isDriver: Boolean)
- extends RecordingShuffleManager(conf, isDriver) with PipelinedShuffleManager
+ extends RecordingShuffleManager(conf, isDriver) with PipelinedShuffleManager
{
+ // A non-streaming pipelined manager (an in-process transport): it needs no
+ // StreamingShuffleOutputTracker, unlike the RPC streaming manager that the
trait defaults to.
+ // The "SparkEnv does not initialize the tracker when the incremental
manager is not streaming"
+ // and fail-loud tests rely on this manager reporting false here.
+ override def usesStreamingShuffleOutputTracker: Boolean = false
Review Comment:
Dropping `"createShuffleMapStage fails loud for a pipelined dependency when
no streaming tracker"` in this file is a net loss of coverage rather than a
wash.
The old invariant is genuinely relaxed, so the test as written had to go.
But the replacement behavior — a pipelined shuffle whose manager reports
`usesStreamingShuffleOutputTracker = false` registers with no output tracker at
all and the job still runs — is exactly what this suite exists to pin, and it
is now asserted nowhere here. The `None` arm added in
`DAGScheduler.outputTrackerMaster` is also unexercised.
Converting the deleted test into the positive case (this manager configured,
job succeeds, `streamingShuffleOutputTracker.isEmpty`, and the shuffle is
absent from `mapOutputTracker`) keeps the coverage where the routing rules are
tested.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleManager.scala:
##########
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.local.pipelined
+
+import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.{BaseShuffleHandle, PipelinedShuffleManager,
ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter,
ShuffleWriteMetricsReporter, ShuffleWriter}
+
+/**
+ * A pipelined shuffle manager whose writer -> reader transport is an
in-process bounded
+ * channel (see [[ChannelShuffleRendezvous]]) rather than the RPC streaming
shuffle. It
+ * serves a [[org.apache.spark.PipelinedShuffleDependency]] on a single
executor, letting
+ * the concurrent-stage scheduler run a shuffle's map and reduce stages at the
same time
+ * while records flow between them in memory -- the in-process pipelined
shuffle execution model.
+ *
+ * Selected via `spark.shuffle.manager.incremental`.
+ *
+ * Unlike the RPC streaming manager, this one needs no
`StreamingShuffleOutputTracker`: it
+ * finds each reader/writer pair through the JVM-local
[[ChannelShuffleRendezvous]] rather
+ * than a directory of writer host/port locations. It therefore declares
+ * `usesStreamingShuffleOutputTracker = false`, so `SparkEnv` creates no
tracker and the
+ * scheduler registers the shuffle with none (a pipelined stage's availability
is tracked on
+ * the stage itself, not in any output tracker). This is why it implements the
+ * `PipelinedShuffleManager` trait directly instead of subclassing the
concrete streaming
+ * manager.
+ *
+ * This manager deliberately keeps NO per-shuffle registry. An early version
recorded each
+ * shuffle's map-task count at registration and looked it up in getReader --
and lost it when
+ * an unregisterShuffle arrived BETWEEN registration and the job running,
which happens
+ * legitimately: Dataset.rdd builds the RDD inside a SQL execution scope that
ends (and, with
+ * spark.sql.classic.shuffleDependency.fileCleanup.enabled, removes the
shuffle from every
+ * manager) before any job has run. The reader then saw a missing entry as
numMaps = 0 and
+ * silently under-read the channel. The count is instead stamped into the
shuffle handle at
+ * registration ([[ChannelShuffleHandle.numMaps]]): the handle travels with
the dependency
+ * into every task, a plain Int field survives task serialization (the
dependency's own `rdd`
+ * reference is @transient and is null inside a deserialized task), and no
later unregister
+ * can take it away.
+ */
+private[spark] class PipelinedChannelShuffleManager(conf: SparkConf)
+ extends PipelinedShuffleManager {
+
+ // The in-process rendezvous is JVM-local: on a multi-executor deployment
each executor would
+ // get its own empty queue map, and every reader would block forever on data
written in some
+ // other JVM -- a silent hang. Refuse to construct anywhere but local mode,
so a
+ // misconfiguration fails loudly at startup instead.
+ require(org.apache.spark.util.Utils.isLocalMaster(conf),
+ "PipelinedChannelShuffleManager is an in-process (single-JVM) transport
and requires " +
+ s"local mode; got master '${conf.get("spark.master", "")}'")
+
+ // Rows accumulated per output partition before a batch is handed across the
channel in one
+ // queue operation. Batching amortizes the queue's per-operation lock cost;
per-row hand-off
+ // measured ~19x slower than a regular shuffle on a 20M-row repartition.
+ private val batchSize = conf.get(config.SHUFFLE_PIPELINED_CHANNEL_BATCH_SIZE)
+
+ override def usesStreamingShuffleOutputTracker: Boolean = false
+
+ // Records cross the channel as object references read by a concurrent
consumer thread; the
+ // SQL layer must detach each row from the producer's reused buffer before
the writer sees it.
+ override def requiresDetachedRecords: Boolean = true
+
+ // Reset a prior run's abandoned-partition marks when this shuffle's
producer stage is
+ // (re)submitted -- before any map task of the new run starts, so it cannot
race the run's own
+ // writers/readers. A shuffleId is re-run within one query (a
RangePartitioner sample job then
+ // the main job; executeTake's per-batch jobs), and abandonment is a per-run
fact.
+ override def onPipelinedProducerStageSubmit(shuffleId: Int): Unit =
+ ChannelShuffleRendezvous.clearAbandonedForShuffle(shuffleId)
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle =
+ new ChannelShuffleHandle(shuffleId, dependency,
dependency.rdd.partitions.length)
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] =
+ new ChannelShuffleWriter[K, V](
+ handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, batchSize,
metrics)
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ // A reduce task reads the reduce-partition range [startPartition,
endPartition). Core
+ // ShuffledRDD uses width 1, but SQL's ShuffledRowRDD may coalesce several
reduce
Review Comment:
This comment contradicts the reader it describes: it says `ShuffledRowRDD`
may coalesce several reduce partitions into one reader task "so the reader must
honor the whole range", while `ChannelShuffleReader` opens with
`require(endPartition - startPartition == 1)` and its class doc explains at
length why a coalesced range cannot be served. Worth reconciling so the next
reader of this file does not trust the wrong one.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleManager.scala:
##########
@@ -0,0 +1,133 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.local.pipelined
+
+import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext}
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.{BaseShuffleHandle, PipelinedShuffleManager,
ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter,
ShuffleWriteMetricsReporter, ShuffleWriter}
+
+/**
+ * A pipelined shuffle manager whose writer -> reader transport is an
in-process bounded
+ * channel (see [[ChannelShuffleRendezvous]]) rather than the RPC streaming
shuffle. It
+ * serves a [[org.apache.spark.PipelinedShuffleDependency]] on a single
executor, letting
+ * the concurrent-stage scheduler run a shuffle's map and reduce stages at the
same time
+ * while records flow between them in memory -- the in-process pipelined
shuffle execution model.
+ *
+ * Selected via `spark.shuffle.manager.incremental`.
+ *
+ * Unlike the RPC streaming manager, this one needs no
`StreamingShuffleOutputTracker`: it
+ * finds each reader/writer pair through the JVM-local
[[ChannelShuffleRendezvous]] rather
+ * than a directory of writer host/port locations. It therefore declares
+ * `usesStreamingShuffleOutputTracker = false`, so `SparkEnv` creates no
tracker and the
+ * scheduler registers the shuffle with none (a pipelined stage's availability
is tracked on
+ * the stage itself, not in any output tracker). This is why it implements the
+ * `PipelinedShuffleManager` trait directly instead of subclassing the
concrete streaming
+ * manager.
+ *
+ * This manager deliberately keeps NO per-shuffle registry. An early version
recorded each
+ * shuffle's map-task count at registration and looked it up in getReader --
and lost it when
+ * an unregisterShuffle arrived BETWEEN registration and the job running,
which happens
+ * legitimately: Dataset.rdd builds the RDD inside a SQL execution scope that
ends (and, with
+ * spark.sql.classic.shuffleDependency.fileCleanup.enabled, removes the
shuffle from every
+ * manager) before any job has run. The reader then saw a missing entry as
numMaps = 0 and
+ * silently under-read the channel. The count is instead stamped into the
shuffle handle at
+ * registration ([[ChannelShuffleHandle.numMaps]]): the handle travels with
the dependency
+ * into every task, a plain Int field survives task serialization (the
dependency's own `rdd`
+ * reference is @transient and is null inside a deserialized task), and no
later unregister
+ * can take it away.
+ */
+private[spark] class PipelinedChannelShuffleManager(conf: SparkConf)
+ extends PipelinedShuffleManager {
+
+ // The in-process rendezvous is JVM-local: on a multi-executor deployment
each executor would
+ // get its own empty queue map, and every reader would block forever on data
written in some
+ // other JVM -- a silent hang. Refuse to construct anywhere but local mode,
so a
+ // misconfiguration fails loudly at startup instead.
+ require(org.apache.spark.util.Utils.isLocalMaster(conf),
+ "PipelinedChannelShuffleManager is an in-process (single-JVM) transport
and requires " +
+ s"local mode; got master '${conf.get("spark.master", "")}'")
+
+ // Rows accumulated per output partition before a batch is handed across the
channel in one
+ // queue operation. Batching amortizes the queue's per-operation lock cost;
per-row hand-off
+ // measured ~19x slower than a regular shuffle on a 20M-row repartition.
+ private val batchSize = conf.get(config.SHUFFLE_PIPELINED_CHANNEL_BATCH_SIZE)
+
+ override def usesStreamingShuffleOutputTracker: Boolean = false
+
+ // Records cross the channel as object references read by a concurrent
consumer thread; the
+ // SQL layer must detach each row from the producer's reused buffer before
the writer sees it.
+ override def requiresDetachedRecords: Boolean = true
+
+ // Reset a prior run's abandoned-partition marks when this shuffle's
producer stage is
+ // (re)submitted -- before any map task of the new run starts, so it cannot
race the run's own
+ // writers/readers. A shuffleId is re-run within one query (a
RangePartitioner sample job then
+ // the main job; executeTake's per-batch jobs), and abandonment is a per-run
fact.
+ override def onPipelinedProducerStageSubmit(shuffleId: Int): Unit =
+ ChannelShuffleRendezvous.clearAbandonedForShuffle(shuffleId)
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle =
+ new ChannelShuffleHandle(shuffleId, dependency,
dependency.rdd.partitions.length)
+
+ override def getWriter[K, V](
+ handle: ShuffleHandle,
+ mapId: Long,
+ context: TaskContext,
+ metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] =
+ new ChannelShuffleWriter[K, V](
+ handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, batchSize,
metrics)
+
+ override def getReader[K, C](
+ handle: ShuffleHandle,
+ startMapIndex: Int,
+ endMapIndex: Int,
+ startPartition: Int,
+ endPartition: Int,
+ context: TaskContext,
+ metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
+ // A reduce task reads the reduce-partition range [startPartition,
endPartition). Core
+ // ShuffledRDD uses width 1, but SQL's ShuffledRowRDD may coalesce several
reduce
+ // partitions into one reader task, so the reader must honor the whole
range. Map-index
+ // bounds are irrelevant to the channel transport (all map tasks share
each partition's
+ // queue). The final 5-arg getReader forwards here with the correct
partition range.
+ //
+ // numMaps -- how many end-of-stream markers to expect per queue -- was
stamped into the
+ // handle at registration, never kept in mutable manager state (see class
scaladoc).
+ val h = handle.asInstanceOf[ChannelShuffleHandle[K, _, C]]
+ new ChannelShuffleReader[K, C](h, startPartition, endPartition, h.numMaps,
metrics)
+ }
+
+ override def unregisterShuffle(shuffleId: Int): Boolean = {
+ ChannelShuffleRendezvous.removeShuffle(shuffleId)
+ true
+ }
+
+ override def stop(): Unit = {}
Review Comment:
`stop()` being empty lets `ChannelShuffleRendezvous` outlive the
`SparkContext` that filled it, and the rendezvous is keyed only by `(shuffleId,
reducePartitionId)` with no application scoping.
Queues are freed only through the GC-driven `ContextCleaner` ->
`unregisterShuffle` path, so anything not collected before the context stops
stays in the `object` for the rest of the JVM. Stop a context and create
another in the same JVM — a test fork, or a REPL/notebook restarting the
context — and shuffle ids start over at 0. The new context's reduce task for
shuffle 0 then drains whatever the previous context left in that queue: extra
rows, or an `EndOfStream` count that never reaches `numMaps`.
This is not hypothetical for the suites in this PR:
`PipelinedChannelShuffleSuite` calls `clearForTesting()` per context, but
`PipelinedShuffleSqlSuite`, `AQEPipelinedShuffleSuite`, and
`PipelinedLimitHangSuite` each build and stop their own session without
clearing it.
`SparkEnv.stop()` already calls `_pipelinedShuffleManager.stop()`, so the
fix is local:
```scala
override def stop(): Unit = ChannelShuffleRendezvous.clear()
```
(promoting `clearForTesting` to a real method rather than a test-only one).
--
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]
