You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "viirya (via GitHub)" <gi...@apache.org> on 2026/08/19 02:30:24 UTC
[PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
viirya opened a new pull request, #58097:
URL: https://github.com/apache/spark/pull/58097
### What changes were proposed in this pull request?
This PR adds an opt-in, in-process **channel transport** for pipelined shuffles and wires it
into the SQL layer, so a local-mode batch query can run its shuffle exchanges as pipelined
shuffles served entirely within one JVM -- producer and consumer stages co-scheduled by the
concurrent-stage scheduler, with records flowing through bounded in-memory queues instead of
shuffle files.
It builds on the already-merged pipelined-shuffle infrastructure -- `PipelinedShuffleDependency`
and dependency-type shuffle routing (SPARK-58185), concurrent-stage scheduling (SPARK-58263),
group-atomic failure and fail-fast rejection (SPARK-58398), and the MapOutputTracker decoupling
(SPARK-58454) -- and adds the pieces specific to a local, in-process transport.
New `core` transport (`org.apache.spark.shuffle.local.pipelined`):
- `ChannelShuffleRendezvous`: a process-wide rendezvous holding one bounded `LinkedBlockingQueue`
per `(shuffleId, reducePartitionId)`. Every map task writing a reduce partition shares the
queue with the single reduce task that drains it. Queue elements are BATCHES of records (an
`Array` of pairs) or an end-of-stream marker, so the queue's per-operation lock cost is paid
per batch, not per row.
- `ChannelShuffleWriter` / `ChannelShuffleReader`: the `ShuffleWriter`/`ShuffleReader` for the
transport. The writer batches records per reduce partition, pushes full batches onto the
queues, and emits one end-of-stream marker per partition; the reader drains its partition's
queue until it has seen `numMaps` end-of-stream markers. Records + read/write time are
reported to the shuffle metrics (no byte metrics: an in-process transport serializes nothing,
so there is no wire-byte count).
- `PipelinedChannelShuffleManager`: a `PipelinedShuffleManager` that mints the writer/reader.
It declares `usesStreamingShuffleOutputTracker = false` (the reader/writer find each other by
`(shuffleId, partition)` in-JVM, so no writer-location directory is needed) and
`requiresDetachedRecords = true` (records cross to a concurrent consumer thread, so the SQL
layer must copy each row off the producer's reused buffer). It `require`s local mode in its
constructor -- a cross-executor deployment would give each executor its own empty queue map
and hang every reader, so it fails loud at startup instead.
SQL integration:
- `EnablePipelinedShuffle` (non-AQE) and `AQEEnablePipelinedShuffle` (AQE) rewrite eligible
`ShuffleExchangeExec` nodes to `pipelined = true`. Both gate on the opt-in flag, on local
mode, and on the in-process channel manager actually being the configured pipelined manager;
otherwise they leave the plan regular. Both leave a plan regular when it contains a reused
exchange (a pipelined producer cannot fan out to more than one consumer).
- `ShuffleExchangeExec` gains a `pipelined` flag and, for the pipelined path, copies each row
off the producer's reused buffer before it is handed across the channel.
- Two registered configs: `spark.sql.pipelinedShuffle.enabled` (default false) turns the
rewrite on; `spark.shuffle.pipelined.channel.batchSize` (default 1024) sets the per-partition
batch size. `spark.shuffle.manager.incremental` selects the channel manager.
Scheduler:
- `DAGScheduler`'s job-shape classification is relaxed to admit a MATERIALIZED-PREFIX MIXED job:
a job may mix regular and pipelined shuffles when every regular boundary reachable from the
final RDD is fully materialized and no pipelined shuffle sits below a regular one. This is the
shape adaptive execution produces (prior map-stage jobs materialize the prefix; the final job
runs the pipelined tail). An unmaterialized regular prefix, and a pipelined shuffle below a
regular boundary, stay rejected fail-fast. The relaxation is a strict superset: previously
rejected shapes now run, and every previously-valid job classifies and schedules identically
(an all-pipelined job is unchanged, pinned by a test). The `DAGScheduler` also passes the
result stage's live reduce partitions to the producer, so a partial-read job (LIMIT /
executeTake reads a subset) does not fill and wedge the queues of partitions no consumer will
drain.
### Why are the changes needed?
Local repartition (a shuffle whose producer and consumer are co-located in one JVM) does not
need the durable, file-based, cross-executor machinery of a regular shuffle. Serving it through
an in-process channel -- records handed directly from writer to a concurrently running reader --
avoids the shuffle-file write/read, block-manager, and serialization/fetch startup costs, which
dominate for the small-to-medium shuffles typical of a single-executor deployment.
Crucially, this reuses Spark's own scheduling machinery (the already-merged concurrent-stage /
pipelined-shuffle infrastructure) rather than introducing a parallel mechanism outside the
shuffle framework: the two shuffle sides remain real, separately scheduled stages, so the shuffle
boundary stays visible in the UI and to AQE, and the same scheduler serves both regular and
pipelined shuffles routed by dependency type. Each transport serves the workload it was built
for -- the RPC streaming transport is latency-optimized for cross-executor streaming, while
the in-process channel is throughput-optimized for local batch (a three-way transport benchmark
in this PR shows the channel beating a regular shuffle while the RPC streaming transport loses to
it on batch shapes, which is why local batch needs its own transport rather than reusing the
streaming one).
### Does this PR introduce _any_ user-facing change?
No behavior change by default: the feature is off unless `spark.sql.pipelinedShuffle.enabled` is
set to true (default false) AND the in-process channel manager is configured via
`spark.shuffle.manager.incremental`, both in local mode. With those set, eligible shuffle
exchanges of a batch query run as in-process pipelined shuffles; results are unchanged, and the
shuffle appears in the UI as concurrently scheduled producer/consumer stages rather than a
materialized boundary. Two new configs are added (both documented, `spark.sql.pipelinedShuffle.enabled`
and `spark.shuffle.pipelined.channel.batchSize`).
### How was this patch tested?
New unit and end-to-end suites, all passing:
- `PipelinedChannelShuffleSuite` (core): the channel transport loses/duplicates no rows and
routes correctly; matches a regular shuffle's grouping; slot admission fits/rejects a group;
the manager refuses to construct outside local mode; a materialized regular prefix runs
end-to-end while an unmaterialized one is rejected; `ContextCleaner` frees the channel's queues
for a tracker-less pipelined shuffle; and deterministic unit tests for the abandon /
end-of-stream-counting logic (abandon marks and drains; `clearAbandoned`/`removeShuffle` reset;
the reader stops after exactly `numMaps` markers).
- `PipelinedShuffleSqlSuite`, `AQEPipelinedShuffleSuite` (sql): batch queries -- repartition,
keyed groupBy, range partitioning, sort-merge join, single-partition aggregate, chains through
SinglePartition -- run end-to-end through the channel and produce correct results against an
independently computed ground truth, under both the non-AQE and AQE rules; over-wide plans fail
loud at admission; cross-subquery reuse cannot create a shared pipelined exchange.
- `PipelinedLimitHangSuite` (sql): a LIMIT over a pipelined shuffle completes and returns the
correct rows in both AQE modes (an early-stopping reader must not wedge the writer).
- `DAGSchedulerSuite` (core): the materialized-prefix relaxation accepts a fully-materialized
mixed job and rejects unmaterialized-prefix / pipelined-below-regular; a FetchFailed on a
pipelined group member (including one reading an external materialized prefix) aborts the whole
group rather than resubmitting a lone stage; an all-pipelined job classifies identically
under the relaxation (a previously-valid shape is unchanged).
- `PipelinedShuffleRoutingSuite`, `ContextCleanerSuite`: routing by dependency type and cleanup.
A `PipelinedShuffleBenchmark` (on the standard `SqlBasedBenchmark` framework, with a checked-in
results file) compares the channel against the regular shuffle and against the RPC streaming
transport across batch shapes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5614877245
I also prefer the narrower scope. BTW, could you resolve the conflicts, @viirya ?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3817078964
##########
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:
Both fixed. (1) `copyRows = requiresDetachedRecords && !needToCopyObjectsBeforeShuffle(part)`, so the write processor no longer copies a row that `prepareShuffleDependency` already copied. (2) `ChannelShuffleWriter` now stores the incoming pair directly instead of re-wrapping it in a fresh `Tuple2` -- one fewer allocation per record.
##########
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:
Fixed by promoting it to a conf: `spark.shuffle.pipelined.channel.queueCapacity` (default 64), and its doc states the `queueCapacity * batchSize * numPartitions` heap-residency product so a user tuning either knob can reason about it. The channel manager sets the rendezvous capacity from it at construction.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3985505299
##########
python/pyspark/sql/tests/test_pipelined_shuffle.py:
##########
@@ -0,0 +1,105 @@
+#
+# 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.
+#
+
+import unittest
+
+from pyspark.sql.functions import udf
+from pyspark.sql.types import LongType
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+
+
+class PipelinedShuffleTests(ReusedSQLTestCase):
Review Comment:
Fixed in `bbb1cefaa94`. Added `pyspark.sql.tests.test_pipelined_shuffle` to `pyspark_sql.python_test_goals` and verified that `determine_dangling_python_tests` no longer reports it.
This was an integration step I missed. The earlier local suite results did not establish that CI could start; thanks for catching it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicBoolean
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, DeserializeToObjectExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.QueryStageExec
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Shared environment gate for the two pipelined-shuffle enabling rules
+ * ([[EnablePipelinedShuffle]] non-AQE and `AQEEnablePipelinedShuffle` under AQE). This is a
+ * CORRECTNESS gate, not cosmetics: flipping an exchange to pipelined while the incremental manager
+ * is still the RPC streaming one would route it to an untested transport that reports
+ * `requiresDetachedRecords = false`, so the SQL layer would skip the row copy and silently corrupt
+ * rows shared across the writer/reader threads. Both rules must apply the identical gate, so it
+ * lives here rather than being copy-pasted into each `apply` (where the two could drift and split
+ * AQE vs non-AQE behavior). Each rule keeps only its own plan-shape logic.
+ */
+private[sql] object PipelinedShuffleEligibility extends Logging {
+
+ // The flag/manager mismatch is a start-up misconfiguration, so warn once per JVM rather than on
+ // every query planned in the session.
+ private val mismatchWarned = new AtomicBoolean(false)
+
+ /** Operators whose consumers cannot safely drain a bounded, single-reader channel. */
+ def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match {
+ case _: CoalesceExec | _: CartesianProductExec | _: DeserializeToObjectExec |
+ _: CollectLimitExec | _: CollectTailExec | _: TakeOrderedAndProjectExec => true
+ case _ => false
+ }
+
+ /**
+ * Estimate the complete group's task demand without executing or materializing the plan.
+ * Each exchange contributes its producer width; the root contributes the consumer width.
+ * Unknown widths conservatively retain regular execution. The scheduler still checks actual
+ * RDD widths and free slots at submission, including contention from other jobs.
+ */
+ def fitsLocalCapacity(plan: SparkPlan, candidates: Set[Int]): Boolean = {
+ val exchanges = plan.collect {
+ case s: ShuffleExchangeExec if s.pipelined || candidates.contains(s.id) => s
+ }.groupBy(_.id).values.map(_.head).toSeq
+ val widths = plan.outputPartitioning.numPartitions +:
+ exchanges.map(_.child.outputPartitioning.numPartitions)
+ val sc = plan.session.sparkContext
+ val slots = sc.maxNumConcurrentTasks(sc.resourceProfileManager.defaultResourceProfile)
+ val demand = widths.map(_.toLong).sum
+ val fits = widths.forall(_ > 0) && demand <= slots
+ if (!fits) {
+ logDebug(s"Pipelined shuffle: estimated stage widths ${widths.mkString(",")} cannot " +
+ s"fit $slots local task slots; retaining regular shuffles.")
+ }
+ fits
+ }
+
+ /**
+ * Whether the pipelined channel transport may be used for `plan` at all, independent of plan
+ * shape. Requires: the opt-in flag on; single-executor local mode (the in-process channel
+ * transport needs producer and consumer in one JVM); and the configured incremental manager
+ * actually being the in-process channel manager. Returns false (with a DEBUG diagnostic) when
+ * any gate fails, so the caller leaves the plan regular.
+ */
+ def enabled(plan: SparkPlan, conf: SQLConf): Boolean = {
+ if (!conf.localPipelinedShuffleEnabled) {
+ return false
+ }
+ if (plan.session == null || !plan.session.sparkContext.isLocal) {
+ return false
+ }
+ // Batch only. `IncrementalExecution.preparations` inherits QueryExecution's list, so without
+ // this gate a streaming plan would be rewritten here: every micro-batch exchange (the
+ // state-store shuffles, the static side of a stream-static join) would be flipped to pipelined
+ // BEFORE `MarkPipelinedShuffleForRealTimeMode` runs. That contradicts what the Real-Time Mode
+ // rule deliberately does -- it leaves the static side regular, because pulling it into the gang
+ // would demand slots for stages that must instead finish first, failing admission. Streaming
+ // marks its own pipelined boundaries; this opt-in batch path must not pre-empt that decision.
+ // (`logicalLink.exists(_.isStreaming)` is the same signal InsertAdaptiveSparkPlan uses to keep
+ // AQE off streaming plans.)
+ // Dataset.rdd exposes arbitrary consumers beyond the SQL plan, including RDD shuffles,
+ // multi-partition reads and repeated reads of the same partition. Cached inputs also hide
+ // shuffle lineage: cache hits skip those readers, while misses may require regular stages.
+ def hasUnsupportedBoundary(p: SparkPlan): Boolean = p match {
+ case _: DeserializeToObjectExec | _: InMemoryTableScanExec => true
Review Comment:
Added `RDDScanExec` and `ExternalRDDScanExec` to the shared boundary exclusions, and applied `withRegularShuffle` at all four `UnionLoopExec` RDD access sites.
The new RDD-input test runs `createDataFrame` over an RDD-level `reduceByKey`, followed by a SQL aggregation and `collect`, without passing through the Dataset RDD exit helper.
The recursive-CTE test includes repartitioned iterations, disables conversion to LocalRelation, and uses a small channel capacity. It checks regular RDD lineage and the complete result. Both tests pass with AQE on and off.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicBoolean
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, DeserializeToObjectExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.QueryStageExec
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Shared environment gate for the two pipelined-shuffle enabling rules
+ * ([[EnablePipelinedShuffle]] non-AQE and `AQEEnablePipelinedShuffle` under AQE). This is a
+ * CORRECTNESS gate, not cosmetics: flipping an exchange to pipelined while the incremental manager
+ * is still the RPC streaming one would route it to an untested transport that reports
+ * `requiresDetachedRecords = false`, so the SQL layer would skip the row copy and silently corrupt
+ * rows shared across the writer/reader threads. Both rules must apply the identical gate, so it
+ * lives here rather than being copy-pasted into each `apply` (where the two could drift and split
+ * AQE vs non-AQE behavior). Each rule keeps only its own plan-shape logic.
+ */
+private[sql] object PipelinedShuffleEligibility extends Logging {
+
+ // The flag/manager mismatch is a start-up misconfiguration, so warn once per JVM rather than on
+ // every query planned in the session.
+ private val mismatchWarned = new AtomicBoolean(false)
+
+ /** Operators whose consumers cannot safely drain a bounded, single-reader channel. */
+ def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match {
+ case _: CoalesceExec | _: CartesianProductExec | _: DeserializeToObjectExec |
+ _: CollectLimitExec | _: CollectTailExec | _: TakeOrderedAndProjectExec => true
+ case _ => false
+ }
+
+ /**
+ * Estimate the complete group's task demand without executing or materializing the plan.
+ * Each exchange contributes its producer width; the root contributes the consumer width.
+ * Unknown widths conservatively retain regular execution. The scheduler still checks actual
+ * RDD widths and free slots at submission, including contention from other jobs.
+ */
+ def fitsLocalCapacity(plan: SparkPlan, candidates: Set[Int]): Boolean = {
+ val exchanges = plan.collect {
+ case s: ShuffleExchangeExec if s.pipelined || candidates.contains(s.id) => s
+ }.groupBy(_.id).values.map(_.head).toSeq
+ val widths = plan.outputPartitioning.numPartitions +:
Review Comment:
Agreed. I retained the conservative unknown-width fallback for this PR rather than adding scan-specific width estimation.
SQLConf and the PR description now explicitly state that ordinary non-bucketed file scans remain regular. Added a Parquet-read/repartition test in both AQE modes that verifies correct results and no pipelined exchange.
This makes the current limitation explicit; it does not add pipelining support for those scans.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala:
##########
@@ -409,6 +409,23 @@ class QueryExecution(
def assertExecutedPlanPrepared(): Unit = executedPlan
+ /**
+ * RDD consumers and partition-at-a-time iterators can outlive one SQL job. Give them a
+ * separate plan whose shuffles retain output, without changing this execution's plan.
+ * Keep the cloned session for lazy planning and AQE, not just for this method's call.
+ */
+ private[sql] def withRegularShuffle: QueryExecution = {
+ val session = SparkSession.getOrCloneSessionWithConfigsOff(
Review Comment:
The fallback is now memoized per QueryExecution. It returns the original execution when the feature is disabled, the shared eligibility gate rejects the plan, or there is no shuffle. The shuffle check looks through adaptive plans and query stages.
Added tests verifying that shuffle-free queries retain the original QE and repeated fallback requests return the same QE/session. Updated the iterator tests to verify reuse of retained regular-shuffle output across repeated actions.
I did not call `SparkSession.close()` after an action: in classic Spark that stops the shared SparkContext. The fallback session remains alive with its QE, and artifact cleanup uses the existing ArtifactManager Cleaner. This reduces unnecessary clones and repeated planning; it does not introduce a new explicit session-disposal mechanism.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -0,0 +1,728 @@
+/*
+ * 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 java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit}
+
+import org.apache.spark.{PipelinedShuffleDependency, SparkEnv, SparkFunSuite}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.shuffle.local.pipelined.ChannelShuffleRendezvous
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.storage.{RDDBlockId, StorageLevel}
+
+/**
+ * End-to-end SQL coverage of the pipelined channel path: a batch query whose hash
+ * exchange is rewritten to a pipelined shuffle (EnablePipelinedShuffle) and served by the
+ * in-process channel manager (PipelinedChannelShuffleManager), run through the
+ * concurrent-stage scheduler on a single executor. Self-manages its SparkSession because the
+ * shuffle manager and AQE-off gate are start-up configs.
+ */
+class PipelinedShuffleSqlSuite extends SparkFunSuite
+ with AdaptiveSparkPlanHelper with PipelinedShuffleTestSession {
+
+ private def withPipelinedSession(body: SparkSession => Unit): Unit =
+ withPipelinedSession("pipelined-shuffle-sql", aqe = false)(body)
+
+ test("batch repartition($k) runs end-to-end through the pipelined channel shuffle") {
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 10)).repartition($"k")
+
+ // Single action only: a pipelined shuffle is single-shot, so collect exactly once and
+ // derive everything from that one result.
+ val rows = df.select($"id").as[Long].collect()
+ val ids = rows.toSet
+
+ // The rule fired and the exchange is pipelined.
+ val pipelinedExchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec if s.pipelined => s
+ }
+ assert(pipelinedExchanges.nonEmpty,
+ s"expected a pipelined ShuffleExchangeExec; plan was:\n${df.queryExecution.executedPlan}")
+
+ // Correctness: the same 1000 ids, repartitioned, all present exactly once.
+ assert(rows.length === 1000, s"expected 1000 rows, got ${rows.length}")
+ assert(ids === (0L until 1000L).toSet)
+ }
+ }
+
+ test("a single keyed groupBy runs end-to-end through the pipelined channel shuffle") {
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7))
+ .groupBy($"k").count()
+
+ // Single action only (pipelined shuffle is single-shot).
+ val counts = df.as[(Long, Long)].collect().toMap
+ val pipelined = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec if s.pipelined => s
+ }
+ assert(pipelined.nonEmpty,
+ s"expected pipelined exchange; plan:\n${df.queryExecution.executedPlan}")
+ // Each residue class 0..6 of 0..999.
+ val expected = (0L until 1000L).groupBy(_ % 7).map { case (k, vs) => (k, vs.size.toLong) }
+ assert(counts === expected)
+ }
+ }
+
+ test("groupBy with ORDER BY (hash + range exchanges) is all-pipelined") {
+ // A trailing ORDER BY adds a RANGE exchange (global sort with 4 shuffle partitions ->
+ // RangePartitioning) on top of the groupBy's hash exchange. The relaxed rule pipelines
+ // BOTH (a mixed pipelined/regular job would be rejected). Range is the interesting case:
+ // RangePartitioner construction runs a SAMPLE job over the exchange's child -- which here
+ // reads the pipelined hash shuffle -- before the main job runs, so this also exercises
+ // two successive jobs over the same single-shot pipelined producer.
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7))
+ .groupBy($"k").count().orderBy($"k")
+
+ val rows = df.as[(Long, Long)].collect()
+ val exchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.nonEmpty && exchanges.forall(_.pipelined),
+ s"every exchange should be pipelined; plan:\n${df.queryExecution.executedPlan}")
+ // Pin the partitioning shapes so this test can't silently stop covering range.
+ val partitionings = exchanges.map(_.outputPartitioning.getClass.getSimpleName).sorted
+ assert(exchanges.exists(_.outputPartitioning.isInstanceOf[
+ org.apache.spark.sql.catalyst.plans.physical.RangePartitioning]),
+ s"expected a RangePartitioning exchange, got: $partitionings; " +
+ s"plan:\n${df.queryExecution.executedPlan}")
+ // Result is correct AND globally ordered by k.
+ val expected = (0L until 1000L).groupBy(_ % 7).map { case (k, vs) => (k, vs.size.toLong) }
+ .toSeq.sortBy(_._1)
+ assert(rows.toSeq === expected)
+ }
+ }
+
+ test("repartitionByRange (pure range exchange) runs through the pipelined channel shuffle") {
+ // A range exchange directly over the scan: RangePartitioner samples the scan (a job with
+ // no shuffle at all), then the main job runs the pipelined range shuffle. Verifies the
+ // channel transport is agnostic to the partitioner kind, and rows land range-partitioned.
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 100))
+ .repartitionByRange($"k")
+ // spark_partition_id() records which output partition each row landed in without
+ // leaving the DataFrame API (Dataset.rdd would execute a separate QueryExecution).
+ .select($"k", org.apache.spark.sql.functions.spark_partition_id().as("p"))
+
+ val partitioned = df.as[(Long, Int)].collect().map { case (k, p) => (p, k) }
+ val exchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.nonEmpty && exchanges.forall(_.pipelined),
+ s"expected a pipelined exchange; plan:\n${df.queryExecution.executedPlan}")
+ assert(exchanges.exists(_.outputPartitioning.isInstanceOf[
+ org.apache.spark.sql.catalyst.plans.physical.RangePartitioning]),
+ s"expected RangePartitioning; plan:\n${df.queryExecution.executedPlan}")
+
+ // No rows lost, and the partitioning is a genuine range split: key ranges of distinct
+ // partitions must not overlap.
+ assert(partitioned.length === 1000)
+ val ranges = partitioned.groupBy(_._1).map { case (p, rows) =>
+ (p, rows.map(_._2).min, rows.map(_._2).max)
+ }.toSeq.sortBy(_._2)
+ ranges.sliding(2).foreach {
+ case Seq((p1, _, max1), (p2, min2, _)) =>
+ assert(max1 <= min2, s"partitions $p1 and $p2 overlap: max($p1)=$max1 > min($p2)=$min2")
+ case _ =>
+ }
+ }
+ }
+
+ test("sort-merge join (both sides hash-exchanged) is all-pipelined and correct") {
+ // A shuffled join is the last TPC-DS transport shape not yet covered: both join inputs
+ // get a hash ShuffleExchangeExec. Disable broadcast so the join is a SortMergeJoin with
+ // two real shuffles; the relaxed rule pipelines both, and the concurrent-stage group
+ // (two producers + the join stage) runs together.
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
+ // Two structurally DIFFERENT inputs so exchange reuse does not collapse them into one
+ // ReusedExchange (which the rule would skip). Different ranges + key expressions.
+ val left = spark.range(0, 200, 1, 2).withColumn("k", ($"id" % 10))
+ .select($"k", $"id".as("lv"))
+ val right = spark.range(0, 120, 1, 2).withColumn("k", ($"id" % 6))
+ .select($"k", $"id".as("rv"))
+ val joined = left.join(right, "k")
+
+ val rows = joined.select($"k", $"lv", $"rv").as[(Long, Long, Long)].collect()
+ val exchanges = collect(joined.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.length >= 2 && exchanges.forall(_.pipelined),
+ s"both join inputs should be pipelined; plan:\n${joined.queryExecution.executedPlan}")
+
+ // Ground truth: an equi-join on k over the two relations.
+ val l = (0L until 200L).map(i => (i % 10, i))
+ val r = (0L until 120L).map(i => (i % 6, i))
+ val expected = (for ((lk, lv) <- l; (rk, rv) <- r if lk == rk) yield (lk, lv, rv)).toSet
+ assert(rows.toSet === expected)
+ }
+ }
+
+ test("global aggregate (single-partition exchange) runs through the pipelined channel") {
+ // An ungrouped aggregate requires AllTuples, planned as a SinglePartition exchange: the
+ // channel's numPartitions == 1 degenerate case (everything routes to queue 0).
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).agg(org.apache.spark.sql.functions.sum($"id"))
+
+ val result = df.as[Long].collect()
+ val exchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.nonEmpty && exchanges.forall(_.pipelined),
+ s"expected a pipelined exchange; plan:\n${df.queryExecution.executedPlan}")
+ assert(exchanges.exists(_.outputPartitioning ==
+ org.apache.spark.sql.catalyst.plans.physical.SinglePartition),
+ s"expected a SinglePartition exchange; plan:\n${df.queryExecution.executedPlan}")
+ assert(result.toSeq === Seq((0L until 1000L).sum))
+ }
+ }
+
+ private def assertRegularRDD(root: RDD[_]): Unit = {
+ val visited = scala.collection.mutable.Set.empty[Int]
+ val pending = scala.collection.mutable.Stack[RDD[_]](root)
+ while (pending.nonEmpty) {
+ val rdd = pending.pop()
+ if (visited.add(rdd.id)) {
+ rdd.dependencies.foreach { dep =>
+ assert(!dep.isInstanceOf[PipelinedShuffleDependency[_, _, _]])
+ pending.push(dep.rdd)
+ }
+ }
+ }
+ }
+
+ for (aqe <- Seq(false, true)) {
+ test(s"typed and Python RDD exports retain regular shuffle lineage with AQE=$aqe") {
+ withPipelinedSession("pipelined-rdd-exports", aqe) { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 4000, 1, 2).repartition(4)
+ assert(df.collect().length === 4000)
+ val typed = df.groupByKey(_ % 4).mapGroups { (key, rows) =>
+ (key, rows.size.toLong)
+ }
+ val rdd = typed.rdd
+ assertRegularRDD(rdd)
+ assert(rdd.coalesce(1).collect().map(_._2).sum === 4000L)
+ val pythonRDD = df.asInstanceOf[org.apache.spark.sql.classic.Dataset[Long]]
+ .javaToPython.rdd
+ assertRegularRDD(pythonRDD)
+ assert(pythonRDD.coalesce(1).count() > 0)
+ val fromRDD = spark.createDataFrame(rdd).repartition(4).rdd
+ assertRegularRDD(fromRDD)
+ assert(fromRDD.coalesce(1).count() === 4L)
+ }
+ }
+
+ test(s"lazy checkpoints retain regular shuffle lineage with AQE=$aqe") {
+ withPipelinedSession("pipelined-checkpoint", aqe) { spark =>
+ withTempDir { dir =>
+ spark.sparkContext.setCheckpointDir(dir.getCanonicalPath)
+ for (reliable <- Seq(false, true)) {
+ val df = spark.range(0, 4000, 1, 2).repartition(4)
+ assert(df.collect().length === 4000)
+ val checkpointed = if (reliable) {
+ df.checkpoint(eager = false)
+ } else {
+ df.localCheckpoint(eager = false)
+ }
+ val rdd = checkpointed.rdd
+ assertRegularRDD(rdd)
+ assert(rdd.coalesce(1).count() === 4000L)
+ assert(checkpointed.collect().sorted === (0L until 4000L).toArray)
+ }
+ }
+ }
+ }
+
+ test(s"SQL cursor computes its producer once with AQE=$aqe") {
+ withPipelinedSession("pipelined-cursor", aqe) { spark =>
+ spark.conf.set("spark.sql.scripting.enabled", "true")
+ spark.conf.set("spark.sql.scripting.cursorEnabled", "true")
+ spark.conf.set("spark.sql.classic.shuffleDependency.fileCleanup.enabled", "false")
+ val evaluated = spark.sparkContext.longAccumulator("cursor producer rows")
+ spark.udf.register("record_cursor_row", (id: Long) => {
+ evaluated.add(1)
+ id
+ })
+ val result = spark.sql(
+ """BEGIN
+ | DECLARE v BIGINT;
+ | DECLARE total BIGINT DEFAULT 0;
+ | DECLARE i INT DEFAULT 0;
+ | DECLARE c CURSOR FOR
+ | SELECT /*+ REPARTITION(4) */ record_cursor_row(id) FROM range(100);
+ | OPEN c;
+ | WHILE i < 100 DO
+ | FETCH c INTO v;
+ | SET total = total + v;
+ | SET i = i + 1;
+ | END WHILE;
+ | CLOSE c;
+ | VALUES (total);
+ |END""".stripMargin)
+ assert(result.collect().head.getLong(0) === (0L until 100L).sum)
+ assert(evaluated.value === 100L)
+ }
+ }
+
+ test(s"Dataset.rdd supports narrow and shuffle consumers with AQE=$aqe") {
+ withPipelinedSession("pipelined-rdd-boundary", aqe) { spark =>
+ val rdd = spark.range(0, 4000, 1, 2).repartition(4).toDF().rdd
+ // Assert eligibility before executing a shape that would hang with the channel.
+ assertRegularRDD(rdd)
+ assert(rdd.coalesce(2).count() === 4000L)
+ assert(rdd.union(rdd).count() === 8000L)
+ assert(rdd.zip(rdd).count() === 4000L)
+ val grouped = rdd.map(r => (r.getLong(0) % 2, 1L)).reduceByKey(_ + _)
+ assert(grouped.collect().toMap === Map(0L -> 2000L, 1L -> 2000L))
+ assert(rdd.repartition(2).count() === 4000L)
+ }
+ }
+
+ test(s"cache construction and partial eviction use regular shuffles with AQE=$aqe") {
+ withPipelinedSession("pipelined-cache", aqe) { spark =>
+ val df = spark.range(0, 4000, 1, 2).repartition(4).persist(StorageLevel.MEMORY_ONLY)
+ try {
+ val classicDf = df.asInstanceOf[org.apache.spark.sql.classic.Dataset[_]]
+ val cached = spark.sharedState.cacheManager.lookupCachedData(classicDf).get
+ .cachedRepresentation.cacheBuilder
+ assert(collect(cached.cachedPlan) {
+ case s: ShuffleExchangeExec if s.pipelined => s
+ }.isEmpty)
+ val expected = (0L until 4000L).toArray
+ assert(df.collect().sorted === expected)
+ val buffers = cached.cachedColumnBuffers
+ assert(buffers.getNumPartitions > 1)
+ SparkEnv.get.blockManager.removeBlock(RDDBlockId(buffers.id, 0))
Review Comment:
Strengthened the cache test with `batchSize = 8` and `queueCapacity = 1`, so the data volume exceeds what an unread channel partition could buffer.
An accumulator in the cached projection now records 4,000 evaluations on initial population, 5,000 after evicting partition 0, and still 5,000 after another read. This verifies that only the evicted partition is recomputed. The test retains the plan-policy assertion before execution.
Also corrected the AQE repeated-action comment: local cleanup makes the stage unavailable and the next action recomputes it. That test does not exercise a genuine `FetchFailed` retry.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,359 @@
+/*
+ * 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] with org.apache.spark.internal.Logging {
+
+ 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
+
+ // Per-run epoch (the jobId), read from the job-level local property the DAGScheduler set for a
+ // pipelined job. The reader of this gang reads the SAME value, so both address the same
+ // per-run queues in the process-wide rendezvous; a re-run of this shuffleId is a different job
+ // and gets a different epoch, keeping its queues physically separate. Absent (a core-RDD test
+ // path that never sets it, and never re-runs a shuffleId concurrently) means epoch 0.
+ private val runEpoch = ChannelShuffleRendezvous.epochOf(TaskContext.get())
+
+ // 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, runEpoch, pid)
+ val start = System.nanoTime()
+ while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, runEpoch, pid)) {
+ // Wake on a task kill even without thread interruption. `isAbandoned` is set only by a
+ // reduce task that actually STARTED (its completion listener); if the reader for pid never
+ // started -- the group aborts while this producer is already filling queues (an
+ // unserializable consumer task, an early failure of another member, a job cancel) -- the
+ // mark never appears and this offer loop would park forever, pinning the executor slot
+ // (spark.job.interruptOnCancel defaults to false, so the kill does not interrupt the
+ // thread). Checking the TaskContext interrupt flag each cycle is the symmetric escape to
+ // the reader's takeItem.
+ Option(TaskContext.get()).foreach(_.killTaskIfInterrupted())
+ 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, runEpoch, 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 = {
+ // No stale-state reset is needed here: this run's queues and abandoned marks are keyed by
+ // runEpoch (the jobId), so an EARLIER run of this shuffleId (a RangePartitioner sampling job
+ // then the main job; executeTake batches; a re-executed classic plan) used a different epoch
+ // and its leftovers are physically separate -- this run starts against empty per-epoch state.
+
+ // One in-progress batch per reduce partition, plus its fill count. A partition's batch array
+ // is allocated LAZILY, on its first record (batches(pid) starts null), so a map task pays for
+ // only the partitions it actually writes -- a wide shuffle (thousands of partitions) or a
+ // partial read (liveMask leaves most partitions dead) does not eagerly allocate
+ // numPartitions * batchSize empty slots up front.
+ val batches = new Array[Array[AnyRef]](numPartitions)
+ val sizes = new Array[Int](numPartitions)
+ // Records skipped because their reduce partition has no consumer (see liveMask).
+ // Reported at the end of write().
+ var droppedRecords = 0L
+
+ 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)) {
+ // This record is routed to a reduce partition the driver said no consumer reads, so it is
+ // dropped -- see liveMask. Count it: dropping is CORRECT only if the live set was computed
+ // correctly, and everything else here fails loudly (a wrong width fails a require, a
+ // reader-less live partition hangs the writer), while an under-approximated live set would
+ // instead lose rows quietly. A non-zero count at the end of a job whose result looks wrong
+ // is the thread to pull.
+ droppedRecords += 1
+ } else {
+ if (batches(pid) == null) batches(pid) = new Array[AnyRef](batchSize)
+ // 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
Review Comment:
Corrected the comment to point to `ShuffleExchangeExec.prepareShuffleDependency` and `needToCopyObjectsBeforeShuffle` for the pipelined copy decision. The reference to the removed `ShuffleWriteProcessor` override is gone.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5499453217
Thanks for the updates, @viirya. CI is green and I re-read the whole diff on `b74f813`. Two new items (1, 2) are interactions with the existing Real-Time Mode path that I missed in earlier rounds; item 3 is a regression in a fix from the last round; the rest are smaller.
## 1. `EnablePipelinedShuffle` is applied to streaming plans too
`IncrementalExecution.preparations` is `state +: (super.preparations :+ PrepareTransformWithStateInitialStateForRealTimeMode :+ MarkPipelinedShuffleForRealTimeMode)`, and `super.preparations` is `QueryExecution.preparations(...)`, which this PR now ends with `Seq(EnablePipelinedShuffle())`. The rule has no streaming guard, so in local mode with the flag and the channel manager on, every exchange in a micro-batch or RTM plan -- the state-store shuffles, the static side of a stream-static join -- is flipped to `pipelined = true` *before* `MarkPipelinedShuffleForRealTimeMode` runs.
That contradicts what `markStreamingPath` deliberately does: it leaves the static side regular because "marking it pipelined would pull it into the group and demand slots for stages that must instead finish, which fails admission". With this PR, the blanket rule has already marked it by the time RTM looks.
The rule (or `PipelinedShuffleEligibility`) needs to bail out for a streaming `QueryExecution`, or `IncrementalExecution` needs to drop it from the inherited list. Either way, a test that a streaming query in a feature-on session is left untouched would pin it.
## 2. The new DAGScheduler paths run for the RPC streaming transport as well
Everything this PR adds to the scheduler is keyed on `PipelinedShuffleDependency` / `isPipelinedProducer`, not on the manager, so it also runs for Real-Time Mode jobs on the default streaming manager:
- `handleJobSubmitted` now clones the job `Properties` and stamps `SPARK_PIPELINED_RUN_EPOCH` for every RTM job.
- `submitMissingTasks` runs `liveReduceSet` + `rddReachesShuffle` (and, on `None`, a second `rddReachesShuffle`) for every RTM producer stage, and stamps `SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS` into its task properties.
- On an unmappable partial read it calls `abortStage(...)` with a message telling the user to "disable spark.sql.shuffle.localPipelined.enabled" -- advice that does not apply to RTM.
The streaming writer never reads either property (nothing under `core/.../shuffle/streaming` references them), so on that path this is pure overhead plus a new failure mode, none of it exercised by the streaming suites. The `requiresDetachedRecords` / `usesStreamingShuffleOutputTracker` pattern already in this PR is the right shape: a capability on `PipelinedShuffleManager` (say `supportsLiveReducePartitionHints`, default `false`) that gates the live-set stamping and the abort, leaving the RTM scheduler path byte-for-byte as it is on master.
## 3. The `classifyJobShuffleShape` early-out never fires
From the last round:
> `classifyJobShuffleShape` returns immediately when no pipelined manager is configured, so a default cluster no longer pays for the graph walk itself on every job submission.
```scala
// DAGScheduler.scala:1170
if (SparkEnv.get == null || SparkEnv.get.pipelinedShuffleManager == null) {
```
`spark.shuffle.manager.incremental` defaults to `"streaming"`, and `SparkEnv.initializeShuffleManager()` instantiates it unconditionally -- `SparkEnv`'s own comment in `initializeStreamingShuffleOutputTracker` says "the manager is non-null here". So on a driver this branch is unreachable and every job submission on every deployment pays the new `(RDD, Boolean)`-keyed walk (a tuple allocation per visit plus the boundary `HashMap`), which is strictly more than the old `classifyJobShuffleKinds` over `traverseRDDGraph`'s `HashSet[RDD[_]]`. The comment claiming otherwise should not ship either way.
A cheap fix that keeps the structure: run the old single-context walk first (`hasPipelined` / `hasRegular` only), and run the two-context walk and the materialization loop only when both are true. All-regular and all-pipelined jobs then cost exactly what they cost on master.
## 4. `ContextCleaner` names a concrete transport
```scala
// ContextCleaner.scala:306
val mgr = SparkEnv.get.pipelinedShuffleManager
mgr != null && !mgr.usesStreamingShuffleOutputTracker &&
org.apache.spark.shuffle.local.pipelined.ChannelShuffleRendezvous.holdsShuffle(shuffleId)
```
Removing the `holdsShuffle` trait method fixed the leak but made core's generic cleanup depend on one transport by fully-qualified name. The registry was the bug, not the abstraction: `def holdsShuffle(shuffleId: Int): Boolean = false` on `PipelinedShuffleManager`, with the channel manager delegating to the rendezvous, gives the same behavior and keeps `ContextCleaner` transport-agnostic. Also `ContextCleaner` already holds `sc` and uses `sc.env.broadcastManager` below; `sc.env.pipelinedShuffleManager` is the consistent form (same for `DAGScheduler.outputTrackerMaster`, which mixes `SparkEnv.get` and `sc.env` on adjacent lines).
## 5. The rendezvous key forces a full scan per cleanup
`queues` and `abandoned` are keyed by the flat `(shuffleId, epoch, reducePartitionId)`, so `holdsShuffle` and `removeShuffle` iterate every live entry. `holdsShuffle` is now called from `doCleanupShuffle` for **every** shuffle cleaned in a feature-on session, including the regular prefix shuffles this feature produces, so a long-lived local session pays O(N x live entries). Keying by `shuffleId` first makes both O(1) and makes `removeShuffle` one atomic `remove`.
## 6. `liveReduceSet`'s identity check is a count heuristic, and a wrong answer is silent data loss
```scala
if (cur.partitions.length == sd.partitioner.numPartitions) resultReduce ++= live
else unmappable = true
```
Equal counts do not imply an identity index mapping (one skew-split plus one coalesce of two gives the same length and a different mapping). Today the AQE rule keeps pipelined exchanges out of `ShuffleQueryStageExec` so neither `OptimizeSkewedJoin` nor `CoalesceShufflePartitions` can produce that -- but that is a scheduler-level invariant resting on a SQL-rule behavior, with no check between them. What makes it worth tightening is the failure mode: the writer's `liveMask` silently drops every record routed outside the computed set, so an under-approximation yields wrong results, not the hang the mechanism exists to prevent. Everything else in this PR fails loud; this path fails quiet. At minimum, count dropped records in the writer and log them at debug; better, make the identity property explicit (restrict the match to `ShuffledRDD` / width-1 `CoalescedPartitionSpec`, or have the reader RDD report the reduce id for a partition index).
## 7. Reader fetch-wait time truncates to zero on fast polls
```scala
// ChannelShuffleWriterReader.scala:292
readMetrics.incFetchWaitTime((System.nanoTime() - start) / 1000000L)
```
This converts each `poll` interval to ms separately; a normal hand-off returns in microseconds, so every sub-millisecond wait contributes 0 and a consumer that is genuinely waiting reports ~0 -- the opposite of the intent in the comment above it. Accumulate nanos and convert once, as `ShuffleBlockFetcherIterator.withFetchWaitTimeTracked` does. (The write side is fine: `writeTime` is in nanoseconds.)
## 8. The benchmark cannot be regenerated by the project's benchmark workflow
`runBenchmarkSuite` returns immediately when `cores < 11`. `.github/workflows/benchmark.yml` runs on `ubuntu-latest` (4 cores), so the workflow would print `[skip]` and never write a results file -- which is why the checked-in file was produced on an "Apple M4 Max / Mac OS X 26.5.2" rather than by the workflow, and why only `-jdk21-results.txt` exists where every other benchmark under `sql/core/benchmarks` has `-results.txt`, `-jdk21-`, and `-jdk25-`. Either shrink the shapes so the gang fits the standard runner (it is `inputParts + 8 + 1`; with `spark.sql.shuffle.partitions=2` the minimum shape fits in 5 slots), or don't check in results that no committer can reproduce through the normal path.
## 9. Config naming and surface
- The streaming transport's knobs are `spark.shuffle.streaming.*`; for symmetry the channel transport's should be `spark.shuffle.channel.batchSize` / `spark.shuffle.channel.queueCapacity`, not `spark.shuffle.pipelined.channel.*`.
- `spark.sql.shuffle.localPipelined.enabled` is public, documented "Experimental.", and does nothing unless `spark.shuffle.manager.incremental` is also set to the channel manager -- the mismatch is only a `logDebug`. Either mark it `.internal()` while the feature bakes, or make the flag self-sufficient and warn once when the manager does not match.
## 10. Small things
- `EnablePipelinedShuffle`'s scaladoc says "a production version would be a targeted, cost/shape-aware replacement rather than a blanket rewrite". Either the blanket rule is what we ship (say so), or that sentence names a follow-up JIRA; it should not ship as-is.
- `EnablePipelinedShuffle()` and `AQEEnablePipelinedShuffle()` are no-field case classes; `object` is the usual form.
- `ShuffleExchangeExec.stringArgs`'s comment, "`pipelined` is only meaningful for a Real-Time Mode plan", is no longer true after this PR.
- `PipelinedShuffleSqlSuite`, `AQEPipelinedShuffleSuite`, and `PipelinedLimitHangSuite` each carry a near-identical ~30-line session harness, and each begins by stopping and clearing whatever active/default `SparkSession` the JVM holds. I could not find precedent for that under `sql/core/src/test`; please factor it into one trait and reconsider tearing down a session the suite did not create.
- The new files run 46-65% comment lines, and several narrate review history ("This replaces an earlier design that...", "An early version recorded each shuffle's map-task count..."). That belongs in the JIRA and commit message; the genuinely subtle invariants (epoch keying, width-1 reads, the abandon/offer re-check) are worth every line.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "tgravescs (via GitHub)" <gi...@apache.org>.
tgravescs commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3898081674
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1269,6 +1256,18 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val PIPELINED_SHUFFLE_ENABLED = buildConf("spark.sql.pipelinedShuffle.enabled")
Review Comment:
I think a lot of shuffle related ones are spark.sql.shuffle.XXX
Also it might be better to make it explicity localPipelinedShuffle to make it clear to the user when it applies.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5621160276
Pushed `022f26af749` to close the RDD and iterator boundary gaps identified in the latest review, including the PySpark entry points, typed Dataset RDD exports, checkpoints, and SQL cursors.
Validation: 135 SQL/cursor tests and the PySpark regression suite passed, along with Scala style and Python lint checks. Details in the inline reply.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3823446866
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,304 @@
+/*
+ * 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)) {
Review Comment:
`putUnlessAbandoned` has no kill/interrupt check, so this is the writer-side twin of the reader hang fixed in the last round.
The only exit from this loop is `isAbandoned`, and that mark is set exclusively by `ChannelShuffleReader`'s task-completion listener -- i.e. only if the reduce task for `pid` actually **started**. If it never starts, the loop never terminates.
That is reachable on the opt-in path, because the producer is submitted before the consumer. In `submitStage` the recursion runs `submitStage(parent)` first and the consumer's `submitMissingTasks` afterwards, so the producer's map tasks can already be filling queues before the consumer TaskSet exists. Anything that aborts the group in that window -- `submitMissingTasks` throwing on the consumer (unserializable task), an early failure of another group member, a job cancel -- leaves the producer running with no reader for some or all partitions. The abort then goes through `cancelTasks` / `killAllTaskAttempts` with `shouldInterruptTaskThread(job)`, which reads `spark.job.interruptOnCancel` and defaults to `false`, so the thread is never interrupted and this `offer` stays parked.
The consequence is the one you fixed on the reader: the task never dies and pins an executor slot for the life of the application. A partial abort is enough -- with 4 reduce partitions and only one reduce task launched before the abort, the writers for the other three park here.
The symmetric fix is the one already applied to `takeItem`: check the kill flag each cycle, e.g. `Option(TaskContext.get()).foreach(_.killTaskIfInterrupted())` at the top of the loop body.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala:
##########
@@ -0,0 +1,155 @@
+/*
+ * 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)]()
+
+ /**
+ * Per-queue capacity in BATCHES (not rows), the backpressure bound and the heap-residency
+ * knob (see spark.shuffle.pipelined.channel.queueCapacity). Set once by the channel manager at
+ * construction from that conf; defaults to 64 (with the default 1024-row batch, ~64K rows per
+ * reduce partition in flight) until a manager sets it. `@volatile` because the manager sets it
+ * on the driver while writer/reader threads read it.
+ */
+ @volatile private var capacity = 64
+
+ /** Set the per-queue capacity in batches. Called by the channel manager from its conf. */
+ private[pipelined] def setCapacity(batches: Int): Unit = { capacity = batches }
+
+ /** The queue for one `(shuffleId, reducePartitionId)`, created on first access. */
+ def queue(shuffleId: Int, reducePartitionId: Int): LinkedBlockingQueue[AnyRef] =
+ queues.computeIfAbsent(
+ (shuffleId, reducePartitionId),
+ _ => new LinkedBlockingQueue[AnyRef](capacity))
+
+ /** Whether this reduce partition's reader has departed (see [[abandon]]). */
+ def isAbandoned(shuffleId: Int, reducePartitionId: Int): Boolean =
+ abandoned.contains((shuffleId, reducePartitionId))
+
+ /**
+ * Clear ALL abandoned marks for a shuffle. A pipelined producer is RE-RUN for the same
+ * shuffleId within one query (a RangePartitioner sampling job then the main job; executeTake's
+ * per-batch jobs), and abandonment is a PER-JOB fact -- a mark left by a previous run's reader
+ * must not make the re-run's fresh writers think a partition is dead.
+ *
+ * This is called ONCE by the DAGScheduler when it submits the producer stage, i.e. BEFORE any
+ * map task of the new run has started, so it can never race a concurrently running writer or
+ * reader of the SAME run. (An earlier design cleared marks from inside each writer's write();
+ * that raced across the run's own map tasks -- a late-starting map task erased a departure a
+ * sibling's reader had legitimately recorded, re-hanging the writer. Clearing at the stage's
+ * submission, the one point with no live task of that run, removes the race by construction.)
+ * Queues are left intact -- only the marks are reset.
+ */
+ def clearAbandonedForShuffle(shuffleId: Int): Unit = {
Review Comment:
Resetting the marks without the queues leaves a previous run's data readable by the next run of the same `shuffleId`.
The scaladoc is explicit that "Queues are left intact", and `abandon` clears a queue only for a partition whose reader actually departed. A partition whose reduce task never started therefore keeps whatever the writers pushed into it.
In classic mode a `Dataset` holds one `executedPlan`, so `ShuffleExchangeExec.shuffleDependency` -- and its `shuffleId` -- is reused across actions. Sequence:
1. An action runs and the group aborts partway (a consumer task fails, or the job is cancelled). Partitions whose readers never started keep their batches, and -- if their writers got that far -- their `EndOfStream` markers too.
2. The user re-runs the action. Same `shuffleId`, new producer stage; `onPipelinedProducerStageSubmit` clears the marks only.
3. The new writers append to the same queues.
4. The new reader drains the stale batches first and counts the stale `EndOfStream` markers toward `numMaps`, so it can stop before the new run's data arrives.
That is a silent wrong result (duplicated and/or dropped rows) rather than a hang.
Calling `removeShuffle(shuffleId)` here instead of `clearAbandonedForShuffle` would cover it, since the hook is documented to run before any map task of the run starts -- but only if a pipelined producer stage can never be resubmitted with partial progress, otherwise the clear would drop in-flight data. A sturdier option is to scope the rendezvous keys by a per-run epoch, `(shuffleId, epoch, reducePartitionId)`. That also closes the related race the mark reset still has: right after step 2 clears the marks, a straggler writer from the aborted run that has not yet noticed its kill sees `isAbandoned == false` again and can push into the new run's queue.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3825469940
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,304 @@
+/*
+ * 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)) {
Review Comment:
Fixed. `putUnlessAbandoned` now calls `killTaskIfInterrupted()` at the top of each loop cycle, so a writer whose reader never started wakes on the group abort interrupt flag instead of parking on `offer` forever -- symmetric to `takeItem`. Added a regression test: a producer map task that throws now fails the job promptly (under a deadline) rather than pinning the slot.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala:
##########
@@ -0,0 +1,155 @@
+/*
+ * 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)]()
+
+ /**
+ * Per-queue capacity in BATCHES (not rows), the backpressure bound and the heap-residency
+ * knob (see spark.shuffle.pipelined.channel.queueCapacity). Set once by the channel manager at
+ * construction from that conf; defaults to 64 (with the default 1024-row batch, ~64K rows per
+ * reduce partition in flight) until a manager sets it. `@volatile` because the manager sets it
+ * on the driver while writer/reader threads read it.
+ */
+ @volatile private var capacity = 64
+
+ /** Set the per-queue capacity in batches. Called by the channel manager from its conf. */
+ private[pipelined] def setCapacity(batches: Int): Unit = { capacity = batches }
+
+ /** The queue for one `(shuffleId, reducePartitionId)`, created on first access. */
+ def queue(shuffleId: Int, reducePartitionId: Int): LinkedBlockingQueue[AnyRef] =
+ queues.computeIfAbsent(
+ (shuffleId, reducePartitionId),
+ _ => new LinkedBlockingQueue[AnyRef](capacity))
+
+ /** Whether this reduce partition's reader has departed (see [[abandon]]). */
+ def isAbandoned(shuffleId: Int, reducePartitionId: Int): Boolean =
+ abandoned.contains((shuffleId, reducePartitionId))
+
+ /**
+ * Clear ALL abandoned marks for a shuffle. A pipelined producer is RE-RUN for the same
+ * shuffleId within one query (a RangePartitioner sampling job then the main job; executeTake's
+ * per-batch jobs), and abandonment is a PER-JOB fact -- a mark left by a previous run's reader
+ * must not make the re-run's fresh writers think a partition is dead.
+ *
+ * This is called ONCE by the DAGScheduler when it submits the producer stage, i.e. BEFORE any
+ * map task of the new run has started, so it can never race a concurrently running writer or
+ * reader of the SAME run. (An earlier design cleared marks from inside each writer's write();
+ * that raced across the run's own map tasks -- a late-starting map task erased a departure a
+ * sibling's reader had legitimately recorded, re-hanging the writer. Clearing at the stage's
+ * submission, the one point with no live task of that run, removes the race by construction.)
+ * Queues are left intact -- only the marks are reset.
+ */
+ def clearAbandonedForShuffle(shuffleId: Int): Unit = {
Review Comment:
Fixed with per-run epoch isolation. The rendezvous is now keyed `(shuffleId, epoch, reducePartitionId)`, where the epoch is the jobId: `DAGScheduler.handleJobSubmitted` stamps it into a pipelined job properties before any stage, and `submitMissingTasks` clones it per stage so the gang producer and consumer read the same value. A re-run of the same shuffleId is a different job, hence a different epoch and physically separate queues -- so it never drains a prior run leftover batches or counts its `EndOfStream` markers toward `numMaps`, and a straggler writer from the aborted run (still on the old epoch) cannot push into the new run queue. This removes the whole "clear the marks between runs" mechanism: `clearAbandonedForShuffle` and the `onPipelinedProducerStageSubmit` hook are both deleted, and `removeShuffle(shuffleId)` now drops every epoch of a shuffle. New regression test asserts the isolation.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3874738545
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,141 @@
+/*
+ * 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.SparkEnv
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, SparkPlan, UnaryExecNode}
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together. This is the minimal SQL entry point that lets a batch query exercise the
+ * pipelined channel execution path; a production version would be a targeted,
+ * cost/shape-aware replacement rather than a blanket rewrite.
+ *
+ * Enabled only when `spark.sql.pipelinedShuffle.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * Two shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - coalesce over a shuffle: a `CoalesceExec` (user `.coalesce(n)`, a narrow no-shuffle
+ * partition reduction) reading from a shuffle makes ONE reduce task drain SEVERAL reduce
+ * partitions sequentially (a core `CoalescedRDD` over the `ShuffledRowRDD`). The channel
+ * transport cannot serve that -- the map-side writer interleaves all partitions on one
+ * thread and parks on a full bounded queue, so a reader draining partition `start` to
+ * completion before touching `start + 1` deadlocks with the parked writer, and there is no
+ * "combine adjacent partitions" narrow read the transport could substitute without giving
+ * up the bounded-queue backpressure the transport relies on. `coalesce`'s API contract is a
+ * narrow dependency that merges adjacent partitions, so we cannot honor it by re-hashing to
+ * `n` partitions either. So if any shuffle in the plan is read by a coalesce (directly or
+ * through a narrow chain) the rule leaves the whole plan regular; the query runs correctly,
+ * just not pipelined. (Leaving only that one exchange regular would put a pipelined exchange
+ * below a regular boundary, which the scheduler rejects -- so it must be all-or-nothing.)
+ */
+case class EnablePipelinedShuffle() extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.pipelinedShuffleEnabled) return plan
+ // Single-executor only: the validated transport (the in-process channel
+ // manager) requires producer and consumer in one JVM. (The pipelined machinery itself is
+ // not local-only -- the RPC streaming transport is cross-executor -- but batch queries
+ // over it are unexplored territory, so the rule stays conservative.)
+ if (plan.session == null || !plan.session.sparkContext.isLocal) return plan
+
+ // The SQL flag alone does not pick a transport: the pipelined manager is set separately by
+ // spark.shuffle.manager.incremental, and DEFAULTS to the RPC StreamingShuffleManager. Only
+ // the in-process channel manager is validated for batch queries in local mode; flipping to
+ // pipelined while the incremental manager is still the RPC streaming one would route these
+ // exchanges to an untested transport (and, because that manager reports
+ // requiresDetachedRecords = false, also skip the row copy). So require the channel manager
+ // to be active; otherwise leave the plan regular, mirroring the reuse fallback below.
+ if (!SparkEnv.get.pipelinedShuffleManager.isInstanceOf[PipelinedChannelShuffleManager]) {
+ logDebug("EnablePipelinedShuffle: spark.sql.pipelinedShuffle.enabled is on but the " +
+ "incremental shuffle manager is not the in-process channel manager " +
+ "(spark.shuffle.manager.incremental); leaving the plan regular.")
+ return plan
+ }
+
+ val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
+ if (shuffles.isEmpty) return plan
+
+ // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so
+ // leave the whole plan regular rather than produce a rejected job. Check subquery plans
+ // too (plan.exists walks the operator tree only): today no SQL shape can place a reused
+ // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse
+ // never fires because the subquery's own preparation pass (PlanSubqueries ->
+ // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE
+ // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery
+ // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the
+ // second mechanism is an accident of rule ordering and the third is optimizer behavior,
+ // so this gate does not rely on either.
+ if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) {
+ // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output,
+ // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has
+ // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined")
+ // rather than WARN, which would fire on every reuse-bearing query and read as a fault.
+ logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " +
+ "avoid a fan-out pipelined job.")
+ return plan
+ }
+
+ // A CoalesceExec reading from a shuffle would make one reduce task drain several reduce
+ // partitions sequentially, which the channel transport cannot serve (see class doc). Leave
+ // the whole plan regular -- like the reuse fallback, this is a normal, expected outcome the
+ // user has nothing to act on, so log at DEBUG rather than WARN.
+ if (readsShuffleByCoalesce(plan)) {
+ logDebug("EnablePipelinedShuffle: a coalesce reads from a shuffle; leaving the plan " +
+ "regular to avoid a coalesced multi-partition read the channel transport cannot serve.")
+ return plan
+ }
+
+ plan.transformUp {
+ case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true)
Review Comment:
**[correctness / CONFIRMED]** The blanket flip has no guard for a shuffle consumed N-to-1 through `CartesianProductExec`, so N concurrent `ChannelShuffleReader`s drain the same `(shuffleId, epoch, pid)` queue.
With AQE off, the flag on, and broadcast not applying, `df.repartition($"k").crossJoin(df2).collect()` plans `CartesianProductExec` directly over the flipped exchange. `UnsafeCartesianRDD`'s `NarrowDependency` (`getParents = id / numPartitionsInRdd2`) computes each left reduce partition once per right partition, and each compute mints a fresh reader on the SAME rendezvous queue (same job, same epoch). Rows and `EndOfStream` markers are then split nondeterministically between concurrent readers — wrong results — and a reader that received fewer than `numMaps` markers polls forever in `takeItem`; worse, the first reader to finish fires `abandon()`, which clears the queue and stops the writer, silently discarding the other readers' data. Nothing rejects the job: the DAGScheduler fan-out check counts distinct consumer RDDs (here 1, the single `ShuffledRowRDD`), and the width-1 `require` guards range width, not reader count.
Fix direction: bail out (leave the plan regular) when a shuffle is consumed through a `CartesianProductExec` — or, more generally, through any narrow dependency that maps several output partitions to one reduce partition.
##########
core/src/main/scala/org/apache/spark/internal/config/package.scala:
##########
@@ -1852,6 +1852,34 @@ package object config {
.stringConf
.createWithDefault("streaming")
+ private[spark] val SHUFFLE_PIPELINED_CHANNEL_BATCH_SIZE =
+ ConfigBuilder("spark.shuffle.pipelined.channel.batchSize")
+ .doc("Number of records the in-process pipelined channel shuffle accumulates per output " +
+ "partition before handing a batch across its queue in one operation. Larger batches " +
+ "amortize the queue's per-operation lock cost at the price of higher hand-off latency " +
+ "and per-partition buffering. Only used when spark.shuffle.manager.incremental is the " +
+ "in-process channel manager.")
+ .version("4.3.0")
Review Comment:
**[conventions / CONFIRMED]** The three new configs — `spark.shuffle.pipelined.channel.batchSize` here, `spark.shuffle.pipelined.channel.queueCapacity` below (line 1877), and `spark.sql.pipelinedShuffle.enabled` in SQLConf — all declare `.version("4.3.0")`, but branch-4.3 is already cut (branch-4.x is at 4.4.0-SNAPSHOT, master at 5.0.0-SNAPSHOT). A config merged now cannot first ship in 4.3.0; the docs would claim these configs exist in a release that never had them. They should be `4.4.0` (or `5.0.0` if this lands master-only). The adjacent `spark.shuffle.manager.incremental` legitimately carries 4.3.0 because it already shipped there — likely where the copied value came from.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1749,14 +1846,75 @@ private[spark] class DAGScheduler(
}
/**
- * The total concurrent-task demand of an all-pipelined job, computed from the RDD graph BEFORE
- * any stage is created (so a rejection based on it leaves no partial scheduler state, exactly as
- * the barrier slot check and the speculation/DA reject do). Because a job is either all-regular
- * or all-pipelined, an all-pipelined job's whole stage graph is one pipelined group; its members
- * are the final result stage plus every pipelined producer. Each member's task count is its RDD's
- * partition count (`rdd.partitions.length`), matching how `createShuffleMapStage` derives
- * `numTasks`. `finalNumPartitions` is the result stage's task count (the number of partitions the
- * job runs, which may be a subset of `finalRDD.partitions`).
+ * The set of reduce partitions of `targetShuffleId` that the result stage actually reads, given
+ * the result RDD and the subset of ITS partitions the job runs (`liveResultPartitions`). Used to
+ * tell a pipelined producer which of its reduce partitions have a consumer, so it can drop the
+ * rest (a partial read -- LIMIT / executeTake -- runs only some result partitions).
+ *
+ * Walk from `rdd` toward the target shuffle, threading the live partition-index set. Each hop is
+ * either a `NarrowDependency` -- map the live set through its generic `getParents(p)` contract
+ * (OneToOne is identity, RangeDependency is an offset, a coalesce dependency is a range, etc.;
+ * no per-operator special-casing) and recurse into the parent -- or the target
+ * `ShuffleDependency` itself, at which point the reader RDD's partition index equals the reduce
+ * partition index (the pipelined reader always uses a width-1 CoalescedPartitionSpec(i, i+1);
+ * the count check below guards against a future offset spec). A node with several dependencies
+ * (a join's ZippedPartitionsRDD) contributes from every branch that reaches the target shuffle.
+ *
+ * Returns None if an edge cannot be mapped (a non-target ShuffleDependency in the path -- which a
+ * result-feeding producer never has, since pipelined-below-regular is rejected earlier -- or a
+ * reader RDD whose partition count does not match the shuffle's, i.e. a non-identity spec).
+ * The caller treats None as "cannot determine": safe to ignore for a full read, fail-fast for a
+ * partial read.
+ */
+ private def liveReduceSet(
+ rdd: RDD[_], liveResultPartitions: Set[Int], targetShuffleId: Int): Option[Set[Int]] = {
+ // Reached the target shuffle: the reader RDD's live partition indices ARE the reduce indices
+ // (width-1 reader spec), verified by the partition-count identity guard.
+ val reachesTarget = rdd.dependencies.collectFirst {
+ case sd: ShuffleDependency[_, _, _] if sd.shuffleId == targetShuffleId => sd
+ }
+ reachesTarget match {
+ case Some(sd) =>
+ if (rdd.partitions.length == sd.partitioner.numPartitions) Some(liveResultPartitions)
+ else None // a non-identity reader spec (e.g. an offset coalesce) -- cannot map safely
+ case None =>
+ // Follow every dependency whose subtree reaches the target shuffle, mapping the live set
+ // through that (narrow) dependency's getParents. A non-narrow edge that is not the target
+ // is unmappable.
+ val branches = rdd.dependencies.filter(d => dependencyReachesShuffle(d, targetShuffleId))
+ if (branches.isEmpty) {
+ None
+ } else {
+ val mapped = branches.map {
+ case nd: NarrowDependency[_] =>
+ val parentLive = liveResultPartitions.flatMap(nd.getParents)
+ liveReduceSet(nd.rdd, parentLive, targetShuffleId)
+ case _ => None // a non-target ShuffleDependency on the path -- unmappable
+ }
+ if (mapped.contains(None)) None else Some(mapped.flatMap(_.get).toSet)
+ }
+ }
+ }
+
+ /** Whether the RDD graph reachable through `dep` contains `targetShuffleId`. */
+ private def dependencyReachesShuffle(dep: Dependency[_], targetShuffleId: Int): Boolean =
Review Comment:
**[correctness / CONFIRMED]** `dependencyReachesShuffle` recurses with no visited set, and `liveReduceSet` re-invokes it (a full subtree re-walk) for every dependency at every recursion level, itself also unmemoized: O(n²) on a deep narrow chain, exponential on narrow diamonds (shared ancestors via `zip`/nested unions), plus unbounded recursion depth — all on the single-threaded `dag-scheduler-event-loop` at pipelined producer-stage submission.
A pipelined producer under ~30 levels of shared narrow fan-in explores ~2^30 paths inside `submitMissingTasks`, freezing all scheduling in the application; a chain thousands of operators deep instead throws `StackOverflowError` on the event loop. Every other traversal in this file uses the deduped `traverseRDDGraph` for exactly this reason (the PR's own `classifyJobShuffleShape` comment cites eliminating O(K x graph) re-walks as motivation).
Fix direction: compute reachability once with a memoized `HashMap[RDD, Boolean]` (iteratively), and thread it through `liveReduceSet` instead of re-querying per branch.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleManager.scala:
##########
@@ -0,0 +1,156 @@
+/*
+ * 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)
+
+ // Per-queue depth in batches (backpressure bound + heap-residency knob). Set the process-wide
+ // rendezvous from the conf at construction, before any writer/reader creates a queue.
+ ChannelShuffleRendezvous.setCapacity(conf.get(config.SHUFFLE_PIPELINED_CHANNEL_QUEUE_CAPACITY))
+
+ 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
+
+ // Shuffle ids this manager currently holds rendezvous state for: added at registration (driver
+ // side, once per shuffle), removed at unregister. A channel shuffle registers with NO output
+ // tracker (usesStreamingShuffleOutputTracker = false), so the ContextCleaner cannot tell one of
+ // ours apart from an already-cleaned regular shuffle by tracker membership alone; it asks
+ // holdsShuffle instead, so its tracker-less cleanup arm fires only for a shuffle we actually
+ // own. This is the ONE registry the manager keeps -- deliberately just a membership set, not the
+ // per-shuffle metadata an earlier design kept and lost to a mid-job unregister (see class doc);
+ // a stale entry cannot mis-serve a reader (numMaps lives in the handle), only scope cleanup.
+ private val registeredShuffleIds =
+ java.util.concurrent.ConcurrentHashMap.newKeySet[Int]()
+
+ override def registerShuffle[K, V, C](
+ shuffleId: Int,
+ dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
+ registeredShuffleIds.add(shuffleId)
+ new ChannelShuffleHandle(shuffleId, dependency, dependency.rdd.partitions.length)
+ }
+
+ override def holdsShuffle(shuffleId: Int): Boolean = registeredShuffleIds.contains(shuffleId)
+
+ 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). The channel
+ // transport serves exactly ONE reduce partition per reader task -- ChannelShuffleReader
+ // require()s endPartition - startPartition == 1 (its class doc explains why a coalesced
+ // multi-partition range cannot be drained safely). A width-1 spec is what actually arrives:
+ // the SQL rules (EnablePipelinedShuffle / AQEEnablePipelinedShuffle) refuse to pipeline any
+ // shuffle read by a CoalesceExec, and AQE also keeps a pipelined exchange out of a
+ // ShuffleQueryStage so CoalesceShufflePartitions never coalesces it -- so both core ShuffledRDD
+ // and SQL's ShuffledRowRDD hand a width-1 spec here. The require makes any future wider range
+ // fail loud rather than deadlock. Map-index bounds are irrelevant (all map tasks share each
+ // partition's queue). The final 5-arg getReader forwards here with the 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 = {
Review Comment:
**[correctness / CONFIRMED]** The unregister-before-run sequence this class's own scaladoc documents (Dataset.rdd under `spark.sql.classic.shuffleDependency.fileCleanup.enabled` ends the SQL scope and removes the shuffle from every manager before any job runs) leaves a cleanup hole: `unregisterShuffle` removes the id from `registeredShuffleIds` and nothing ever re-adds it (`registerShuffle` runs once, in the dependency constructor). When the job then runs, `ChannelShuffleRendezvous.queue()`'s `computeIfAbsent` recreates the queues — but `holdsShuffle` is now permanently false, so at GC time `ContextCleaner.heldByTracklessPipelinedManager` returns false, the tracker-less arm never fires (the shuffle is in neither tracker), and the recreated queues plus `abandoned` marks leak until `SparkContext.stop()`. The PR fixed the correctness half of this sequence (numMaps stamped into the handle) but not the cleanup half.
Impact is bounded — the leaked queues are empty (the reader's completion listener drains them), so it is ~numPartitions map entries + abandoned key tuples per run, growing unboundedly over a long-lived session.
Fix direction: re-add the id on first rendezvous access for the shuffle (or have the writer/reader re-register), or make the cleaner's tracker-less arm fire whenever the rendezvous holds state for the id (e.g. a `ChannelShuffleRendezvous.holdsShuffle` check) rather than relying on the manager's registry.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -741,15 +746,40 @@ private[spark] class DAGScheduler(
* enforces the same invariant, see StreamingShuffleReader).
*/
private def outputTrackerMaster(
- shuffleDep: ShuffleDependency[_, _, _]): ShuffleOutputTrackerMaster = {
+ shuffleDep: ShuffleDependency[_, _, _]): Option[ShuffleOutputTrackerMaster] = {
if (shuffleDep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) {
- sc.env.streamingShuffleOutputTracker
- .getOrElse(throw new IllegalStateException(
- s"A pipelined shuffle (id ${shuffleDep.shuffleId}) requires a " +
- "StreamingShuffleOutputTracker, but none is configured"))
- .asInstanceOf[StreamingShuffleOutputTrackerMaster]
+ // A pipelined shuffle uses the StreamingShuffleOutputTracker only when its manager
+ // discovers writers over RPC. An in-process transport declares it needs no tracker
+ // (SparkEnv then creates none); such a shuffle registers with no output tracker at all,
+ // and its map-stage availability is tracked locally on the ShuffleMapStage. When a
+ // tracker IS expected (the RPC streaming manager) but absent, that is a real
+ // misconfiguration, so keep failing loudly.
+ sc.env.streamingShuffleOutputTracker match {
Review Comment:
**[correctness / PLAUSIBLE, low severity]** The `Some` arm never consults `usesStreamingShuffleOutputTracker`, and `SparkEnv.initializeStreamingShuffleOutputTracker` creates the tracker for `blockingIsMulti` too. So with the (deprecated) `MultiShuffleManager` as blocking manager plus the channel manager as incremental, the tracker exists and a channel-served pipelined shuffle IS registered in the `StreamingShuffleOutputTrackerMaster` — contradicting the invariant the comment below states ("SparkEnv then creates none; such a shuffle registers with no output tracker at all") and the one `PipelinedShuffleRoutingSuite` pins. The 'two decisions locked consistent / UNREACHABLE by construction' analysis omits the `blockingIsMulti` OR-clause.
Consequences are cosmetic on inspection: the registration is inert (the channel reader/writer never consult the tracker) and cleanup still routes through `shuffleDriverComponents.removeShuffle` in the streaming arm, so nothing leaks. Still, either the `Some` arm should consult the flag, or the comment/test should carry the caveat for this combination.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,233 @@
+/*
+ * 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.adaptive
+
+import scala.collection.mutable
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{BinaryExecNode, CoalesceExec, SparkPlan}
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.ShuffledJoin
+
+/**
+ * Opt-in (SPARK-57399). Flips eligible [[ShuffleExchangeExec]]
+ * nodes to `pipelined = true` under AQE, the adaptive counterpart of the non-AQE
+ * `EnablePipelinedShuffle` preparation rule (which is a no-op once the plan is wrapped in
+ * `AdaptiveSparkPlanExec`). Runs in `AdaptiveSparkPlanExec.queryStagePreparationRules`, so
+ * it is re-applied on every replanning round; the decision is deterministic on plan shape,
+ * and already-flipped exchanges are left alone.
+ *
+ * Placement policy. A flipped exchange has no
+ * map output statistics (it never materializes as a query stage: the DAGScheduler
+ * gang-runs it inline with its consumer in the final job -- see the pipelined case in
+ * `AdaptiveSparkPlanExec.createNonResultQueryStages`), so only exchanges whose statistics
+ * no AQE decision consumes are flipped:
+ *
+ * - "free" candidate: the path from the candidate to the plan root crosses no
+ * stats-sensitive node ([[BinaryExecNode]], another [[ShuffleExchangeExec]], or a
+ * query stage). Its own coalescing/skew handling is given up; nothing above needed its
+ * stats.
+ * - "join-paired" candidate: the immediate shuffle inputs of a [[ShuffledJoin]] whose
+ * path to the root is otherwise free, flipped only as a symmetric pair (an asymmetric
+ * flip would leave one side participating in AQE coalesce/skew and the other fixed).
+ * - everything else stays regular and materializes as usual -- those stages form the
+ * fully-materialized prefix the scheduler's mixed-job shape requires.
+ *
+ * A pipelined exchange supports every
+ * partitioning, so a SinglePartition exchange in a free position is simply a candidate
+ * itself. The walk stops below a flipped candidate: exchanges underneath stay regular and
+ * keep full AQE treatment. Candidates whose canonicalized form occurs more than once in
+ * the plan (including inside materialized stages and subqueries) are skipped: flipping
+ * them would trade AQE's stage reuse for duplicate recomputation, and a pipelined producer
+ * cannot be consumed twice.
+ */
+case class AQEEnablePipelinedShuffle() extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.pipelinedShuffleEnabled) return plan
+ // Single-executor only: the validated transport (the in-process
+ // channel manager) requires producer and consumer in one JVM.
+ if (plan.session == null || !plan.session.sparkContext.isLocal) return plan
+ // The SQL flag alone does not pick a transport: the pipelined manager is set separately by
+ // spark.shuffle.manager.incremental and defaults to the RPC StreamingShuffleManager. Only the
+ // in-process channel manager is validated here; require it, else leave the plan regular (same
+ // reasoning as the non-AQE EnablePipelinedShuffle rule).
+ if (!SparkEnv.get.pipelinedShuffleManager.isInstanceOf[PipelinedChannelShuffleManager]) {
+ logDebug("AQEEnablePipelinedShuffle: spark.sql.pipelinedShuffle.enabled is on but the " +
+ "incremental shuffle manager is not the in-process channel manager; leaving regular.")
+ return plan
+ }
+
+ flipEligibleExchanges(plan)
+ }
+
+ /**
+ * The plan-shape core of the rule, factored out of [[apply]]'s environment guards (opt-in flag,
+ * local mode, channel manager) so it can be unit-tested on a hand-built plan directly. Collects
+ * the eligible exchanges and returns the plan with each flipped to `pipelined = true`.
+ */
+ private[adaptive] def flipEligibleExchanges(plan: SparkPlan): SparkPlan = {
+ val shared = if (conf.exchangeReuseEnabled) duplicatedShuffleForms(plan) else Set.empty[Any]
+ // Collect the exchanges to flip BY IDENTITY (SparkPlan.id, unique per instance), not by the
+ // node itself: TreeNode overrides hashCode but not equals, so a HashSet[ShuffleExchangeExec]
+ // matches structurally, and the transformDown below would then flip EVERY exchange
+ // structurally equal to a collected one -- including a twin the collector deliberately left
+ // regular on a blocked path. That twin, if it sits below a regular boundary, makes
+ // classifyJobShuffleShape reject the whole job. Keying on the instance id flips exactly the
+ // nodes the collector chose, regardless of spark.sql.exchange.reuse (duplicatedShuffleForms,
+ // the only other guard, is empty when reuse is off). transformDown matches each ORIGINAL node
+ // before rebuilding it, so its id is the same instance id the collector recorded.
+ val toFlip = mutable.HashSet.empty[Int]
+ collectCandidates(plan, blocked = false, shared, toFlip)
+ if (toFlip.isEmpty) return plan
+
+ // transformDown, NOT transformUp: candidates can be nested (a SinglePartition candidate
+ // above a hash candidate). transformUp rebuilds children first, so by the time it
+ // reaches the upper candidate that node is a NEW instance whose (already flipped) child
+ // no longer matches the collected original structurally, and the upper flip is silently
+ // dropped -- leaving a regular exchange above a pipelined one, which the scheduler then
+ // rejects. transformDown hands each candidate to the pattern before its subtree is
+ // rebuilt, so both nested flips apply.
+ plan.transformDown {
+ case s: ShuffleExchangeExec if toFlip.contains(s.id) => s.copy(pipelined = true)
+ }
+ }
+
+ private def isCandidate(s: ShuffleExchangeExec, shared: Set[Any]): Boolean =
+ !s.pipelined && !shared.contains(s.canonicalized)
+
+ /**
+ * Top-down walk collecting exchanges to flip. `blocked` is true once the path from the
+ * root has crossed a stats-sensitive node.
+ */
+ private def collectCandidates(
+ plan: SparkPlan,
+ blocked: Boolean,
+ shared: Set[Any],
+ out: mutable.HashSet[Int]): Unit = plan match {
+ case s: ShuffleExchangeExec =>
+ val flipped = !blocked && isCandidate(s, shared)
+ if (flipped) {
+ out += s.id
+ }
+ // A flipped SinglePartition exchange keeps the walk going: AQE makes no decision at
+ // it (it cannot be coalesced or skew-split), so free candidates BELOW it flip too,
+ // forming a pipelined chain: all exchanges in such a chain must flip together (a
+ // SinglePartition exchange cannot be coalesced or skew-split, so AQE makes no decision
+ // at it and free candidates below it flip too). Below any OTHER exchange (flipped or
+ // not) the walk stops: what is underneath either materializes as the prefix or feeds
+ // a regular exchange whose stats AQE uses, and keeps full AQE treatment either way.
+ if (flipped &&
+ s.outputPartitioning == org.apache.spark.sql.catalyst.plans.physical.SinglePartition) {
+ collectCandidates(s.child, blocked = false, shared, out)
+ }
+
+ case _: QueryStageExec => // already materialized; a leaf here
+
+ case j: ShuffledJoin if !blocked =>
+ // Flip the join's immediate shuffle inputs only as a symmetric pair.
+ val leftCandidate = immediateShuffleInput(j.left, shared)
+ val rightCandidate = immediateShuffleInput(j.right, shared)
+ (leftCandidate, rightCandidate) match {
+ case (Some(l), Some(r)) =>
+ out += l.id
+ out += r.id
+ case _ => // asymmetric (a broadcast side, a materialized stage, no clean input): skip
+ }
+ // Anything deeper is below a join input; blocked either way.
+
+ case p if isStatsSensitive(p) =>
+ p.children.foreach(collectCandidates(_, blocked = true, shared, out))
+
+ case p =>
+ p.children.foreach(collectCandidates(_, blocked, shared, out))
+ }
+
+ /**
+ * The single eligible [[ShuffleExchangeExec]] at the top of one join input, looking
+ * through unary non-stats-sensitive forwarders. None if the input is anything else.
+ */
+ private def immediateShuffleInput(
+ plan: SparkPlan,
+ shared: Set[Any]): Option[ShuffleExchangeExec] = plan match {
+ case s: ShuffleExchangeExec => Some(s).filter(isCandidate(_, shared))
+ case _: QueryStageExec => None
+ case p if isStatsSensitive(p) => None
+ case p if p.children.size == 1 => immediateShuffleInput(p.children.head, shared)
+ case _ => None
+ }
+
+ /**
+ * Nodes below which a candidate exchange must NOT be flipped. Two reasons a node lands here:
+ * - AQE consumes map output statistics from the stages below it ([[BinaryExecNode]], another
+ * [[ShuffleExchangeExec]]); flipping below it would give up stats a decision needs.
+ * - [[CoalesceExec]] reads its child shuffle multi-partition-per-task (a `CoalescedRDD` over
+ * the `ShuffledRowRDD`), which the channel transport cannot serve; the shuffle it reads
+ * must stay regular. Blocking here keeps that shuffle -- and everything deeper -- regular,
+ * so no pipelined exchange ends up below the coalesce's regular boundary (which the
+ * scheduler would reject). See the non-AQE `EnablePipelinedShuffle` for the full rationale.
+ */
+ private def isStatsSensitive(plan: SparkPlan): Boolean = plan match {
Review Comment:
**[correctness / CONFIRMED]** Neither flip rule guards against operators that build *hidden* regular shuffles inside `doExecute`: `CollectLimitExec`, `CollectTailExec`, and `TakeOrderedAndProjectExec` each call `ShuffleExchangeExec.prepareShuffleDependency` with the default `pipelined = false` and no plan node, so the plan walk cannot see them. A flipped exchange below them puts a pipelined shuffle under an unmaterialized regular boundary — the exact shape `classifyJobShuffleShape` rejects (`hasPipelinedBelowRegular`) — so the query hard-fails with the mixing error instead of falling back to regular.
Repros (flag on, either AQE mode): `df.groupBy("k").count().orderBy("cnt").limit(5).write.parquet(...)` — `SpecialLimits` plans `TakeOrderedAndProjectExec` for the non-root sort+limit; it is a `UnaryExecNode` not matched by `isStatsSensitive`, so the exchange below it flips, then `doExecute` builds the hidden SinglePartition regular shuffle (child has >1 partition) and job submission throws. Likewise `df.repartition($"k").limit(10).toLocalIterator()` via `CollectLimitExec.doExecute`. The PR's tests only exercise `limit(n).collect()`, which takes `executeTake` and never hits `doExecute` — which is why this went unnoticed.
Fix direction: treat `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec` as blocking in both rules (like `CoalesceExec`), or teach those operators to emit a pipelined-aware dependency.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,141 @@
+/*
+ * 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.SparkEnv
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, SparkPlan, UnaryExecNode}
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together. This is the minimal SQL entry point that lets a batch query exercise the
+ * pipelined channel execution path; a production version would be a targeted,
+ * cost/shape-aware replacement rather than a blanket rewrite.
+ *
+ * Enabled only when `spark.sql.pipelinedShuffle.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * Two shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - coalesce over a shuffle: a `CoalesceExec` (user `.coalesce(n)`, a narrow no-shuffle
+ * partition reduction) reading from a shuffle makes ONE reduce task drain SEVERAL reduce
+ * partitions sequentially (a core `CoalescedRDD` over the `ShuffledRowRDD`). The channel
+ * transport cannot serve that -- the map-side writer interleaves all partitions on one
+ * thread and parks on a full bounded queue, so a reader draining partition `start` to
+ * completion before touching `start + 1` deadlocks with the parked writer, and there is no
+ * "combine adjacent partitions" narrow read the transport could substitute without giving
+ * up the bounded-queue backpressure the transport relies on. `coalesce`'s API contract is a
+ * narrow dependency that merges adjacent partitions, so we cannot honor it by re-hashing to
+ * `n` partitions either. So if any shuffle in the plan is read by a coalesce (directly or
+ * through a narrow chain) the rule leaves the whole plan regular; the query runs correctly,
+ * just not pipelined. (Leaving only that one exchange regular would put a pipelined exchange
+ * below a regular boundary, which the scheduler rejects -- so it must be all-or-nothing.)
+ */
+case class EnablePipelinedShuffle() extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ if (!conf.pipelinedShuffleEnabled) return plan
+ // Single-executor only: the validated transport (the in-process channel
+ // manager) requires producer and consumer in one JVM. (The pipelined machinery itself is
+ // not local-only -- the RPC streaming transport is cross-executor -- but batch queries
+ // over it are unexplored territory, so the rule stays conservative.)
+ if (plan.session == null || !plan.session.sparkContext.isLocal) return plan
+
+ // The SQL flag alone does not pick a transport: the pipelined manager is set separately by
+ // spark.shuffle.manager.incremental, and DEFAULTS to the RPC StreamingShuffleManager. Only
+ // the in-process channel manager is validated for batch queries in local mode; flipping to
+ // pipelined while the incremental manager is still the RPC streaming one would route these
+ // exchanges to an untested transport (and, because that manager reports
+ // requiresDetachedRecords = false, also skip the row copy). So require the channel manager
+ // to be active; otherwise leave the plan regular, mirroring the reuse fallback below.
+ if (!SparkEnv.get.pipelinedShuffleManager.isInstanceOf[PipelinedChannelShuffleManager]) {
+ logDebug("EnablePipelinedShuffle: spark.sql.pipelinedShuffle.enabled is on but the " +
+ "incremental shuffle manager is not the in-process channel manager " +
+ "(spark.shuffle.manager.incremental); leaving the plan regular.")
+ return plan
+ }
+
+ val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
+ if (shuffles.isEmpty) return plan
+
+ // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so
+ // leave the whole plan regular rather than produce a rejected job. Check subquery plans
+ // too (plan.exists walks the operator tree only): today no SQL shape can place a reused
+ // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse
+ // never fires because the subquery's own preparation pass (PlanSubqueries ->
+ // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE
+ // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery
+ // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the
+ // second mechanism is an accident of rule ordering and the third is optimizer behavior,
+ // so this gate does not rely on either.
+ if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) {
+ // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output,
+ // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has
+ // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined")
+ // rather than WARN, which would fire on every reuse-bearing query and read as a fault.
+ logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " +
+ "avoid a fan-out pipelined job.")
+ return plan
+ }
+
+ // A CoalesceExec reading from a shuffle would make one reduce task drain several reduce
+ // partitions sequentially, which the channel transport cannot serve (see class doc). Leave
+ // the whole plan regular -- like the reuse fallback, this is a normal, expected outcome the
+ // user has nothing to act on, so log at DEBUG rather than WARN.
+ if (readsShuffleByCoalesce(plan)) {
+ logDebug("EnablePipelinedShuffle: a coalesce reads from a shuffle; leaving the plan " +
+ "regular to avoid a coalesced multi-partition read the channel transport cannot serve.")
+ return plan
+ }
+
+ plan.transformUp {
+ case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true)
+ }
+ }
+
+ /**
+ * True if any [[ShuffleExchangeExec]] in `plan` is read by a [[CoalesceExec]] above it through
+ * a chain of only narrow (unary, non-exchange) operators. Such a shuffle would be drained
+ * multi-partition-per-task by a `CoalescedRDD` and cannot be pipelined (see class doc). A
+ * shuffle underneath that shuffle is NOT affected: its own reader reads one reduce partition
+ * per task, so the walk from a coalesce stops at the FIRST shuffle it reaches.
+ */
+ private def readsShuffleByCoalesce(plan: SparkPlan): Boolean = {
+ // Does a narrow chain from `p` reach a ShuffleExchangeExec before any other exchange?
+ def reachesShuffle(p: SparkPlan): Boolean = p match {
+ case _: ShuffleExchangeExec => true
+ case u: UnaryExecNode => reachesShuffle(u.child)
+ case _ => false
Review Comment:
**[correctness / CONFIRMED]** `reachesShuffle` walks only `UnaryExecNode` chains, so a `CoalesceExec` above a `UnionExec` (or any `BinaryExecNode`) whose subtree contains shuffles is not detected, and those exchanges get flipped — producing exactly the coalesced multi-partition read the `ChannelShuffleReader` class doc says must never reach the transport.
Repro shape (AQE off, flag on, channel manager): `df1.repartition($"k").union(df2.repartition($"k")).coalesce(2).collect()` with >65K rows per reduce partition. `c.child` is `UnionExec`, not a `UnaryExecNode`, so the guard returns false and both exchanges flip. At runtime `CoalescedRDD`'s one task drains reduce partition p0 to completion (it needs `numMaps` `EndOfStream` markers) before touching p1, while the single-threaded writer fills p1's bounded queue (64 batches) and parks in `putUnlessAbandoned` before ever emitting p0's markers. No abandon mark appears (the reader task never completes), the per-partition width-1 `require` never trips, and there is no timeout escape — the job hangs indefinitely. The AQE rule is safe (`isStatsSensitive` blocks at `CoalesceExec`); only this non-AQE guard has the gap, and the existing test covers only a unary chain (`groupBy.count().coalesce(2)`).
Fix direction: recurse into all children of any non-exchange node (`case p => p.children.exists(reachesShuffle)`), stopping only at exchanges.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1101,32 +1131,99 @@ private[spark] class DAGScheduler(
}
/**
- * Classifies the shuffle boundaries in the RDD graph rooted at `finalRDD` by kind, walking the
- * RDD dependency graph directly (before any stages are created). Returns whether the graph
- * contains any [[PipelinedShuffleDependency]] and whether it contains any regular (non-pipelined)
- * `ShuffleDependency`. Narrow dependencies are not boundaries and are ignored.
+ * Shape of the shuffle boundaries in the RDD graph rooted at a job's final RDD, computed by
+ * [[classifyJobShuffleShape]] before any stage is created (so an unsupported job is rejected
+ * fail-fast with no partial scheduler state).
*
- * A job must be either ALL-regular or ALL-pipelined: a job whose shuffle graph mixes a pipelined
- * shuffle with a regular one is rejected fail-fast (see `handleJobSubmitted`). Restricting to a
- * single kind makes the whole job one pipelined group when pipelined, so gang admission can be
- * decided up front before any stage is submitted.
+ * The supported shapes are:
+ * - all-regular (`hasPipelined` false, nothing pipelined anywhere);
+ * - all-pipelined (`hasPipelined` true, no regular boundary anywhere);
+ * - MATERIALIZED-PREFIX MIXED: pipelined shuffles in the region reachable from the final RDD
+ * without crossing a regular shuffle, where every regular boundary at the edge of that
+ * region is FULLY MATERIALIZED (all map outputs registered with the MapOutputTracker) and
+ * no pipelined shuffle sits below any regular boundary. The materialized prefix never
+ * re-runs, so the job executes exactly like an all-pipelined job whose leaves read
+ * already-materialized shuffle data; gang admission demand (final stage + suffix producers)
+ * is unchanged. This is the shape adaptive execution produces: prior map-stage jobs
+ * materialize the prefix stages, and the final job runs the pipelined tail.
+ *
+ * An UNMATERIALIZED regular boundary in a pipelined job stays rejected: its stage would have to
+ * run while gang-admitted producers already hold slots (blocked on transport backpressure
+ * waiting for consumers), and admission does not account for the prefix's slots -- the prefix
+ * could be starved and deadlock the group. Sequencing the prefix before the gang is future
+ * work. A pipelined shuffle BELOW a regular boundary is also rejected: it is not part of the
+ * suffix group, and (if the boundary were unmaterialized) would have to run under a regime the
+ * group machinery does not cover.
*/
- private def classifyJobShuffleKinds(finalRDD: RDD[_]): (Boolean, Boolean) = {
+ private case class JobShuffleShape(
+ hasPipelined: Boolean,
+ hasUnmaterializedRegularBoundary: Boolean,
+ hasPipelinedBelowRegular: Boolean) {
+ /** Mixed in a way the scheduler does not support (see class doc). */
+ def isUnsupportedMix: Boolean =
+ hasPipelinedBelowRegular || (hasPipelined && hasUnmaterializedRegularBoundary)
+ }
+
+ /** Classify `finalRDD`'s shuffle graph; see [[JobShuffleShape]] for the shape semantics. */
+ private def classifyJobShuffleShape(finalRDD: RDD[_]): JobShuffleShape = {
var hasPipelined = false
- var hasRegular = false
- traverseRDDGraph(finalRDD) { (rdd, enqueue) =>
- rdd.dependencies.foreach { dep =>
- dep match {
- case _: PipelinedShuffleDependency[_, _, _] => hasPipelined = true
- case _: ShuffleDependency[_, _, _] => hasRegular = true
- case _ => // narrow dependency: not a boundary
+ var pipelinedBelow = false
+ // Frontier regular shuffle boundaries: those reachable from the final RDD WITHOUT crossing
+ // another regular boundary, deduped by shuffle ID. Only these matter for the materialization
+ // check (a regular boundary below another one is never a runnable suffix member).
+ val regularBoundaries = new HashMap[Int, ShuffleDependency[_, _, _]]
+
+ // ONE walk, carrying `belowRegular` (true once the path from the final RDD has crossed a
+ // regular boundary), computes hasPipelined and pipelinedBelow together -- replacing the old
+ // per-boundary rddGraphHasPipelinedDependency re-walks (O(K x graph) on shared ancestors).
+ // A node reachable BOTH above and below a regular boundary must be explored in BOTH contexts:
+ // a pipelined dep under it counts as pipelinedBelow on the below path but not on the above
+ // path. So the visited set is keyed on (RDD, belowRegular), NOT on the RDD alone -- keying on
+ // the RDD alone would let the first-reached context win and drop the other, missing a
+ // pipelined-below-regular dep (a wrongly-accepted job). A node is thus visited at most twice,
+ // keeping the cost O(graph) rather than O(K x graph). hasPipelined is set only above a regular
+ // boundary, matching the old walk (which stopped at boundaries): a below-boundary pipelined
+ // dep is the pipelinedBelow reject case, never a runnable group member.
+ val visited = new HashSet[(RDD[_], Boolean)]
+ val stack = new ListBuffer[(RDD[_], Boolean)]
+ stack += ((finalRDD, false))
+ while (stack.nonEmpty) {
+ val entry = stack.remove(0)
+ val rdd = entry._1
+ val belowRegular = entry._2
+ if (visited.add(entry)) {
+ rdd.dependencies.foreach {
+ case pd: PipelinedShuffleDependency[_, _, _] =>
+ if (belowRegular) pipelinedBelow = true else hasPipelined = true
+ stack.prepend((pd.rdd, belowRegular))
+ case sd: ShuffleDependency[_, _, _] =>
+ // A frontier boundary only when not already below one; descend with belowRegular set.
+ if (!belowRegular) regularBoundaries.getOrElseUpdate(sd.shuffleId, sd)
+ stack.prepend((sd.rdd, true))
+ case narrowDep =>
+ stack.prepend((narrowDep.rdd, belowRegular))
}
- // Descend through every edge (shuffle and narrow) so a pipelined boundary behind a regular
- // one -- or vice versa -- anywhere in the graph is still detected. traverseRDDGraph dedups.
- enqueue(dep.rdd)
}
}
- (hasPipelined, hasRegular)
+
+ var hasUnmaterialized = false
+ regularBoundaries.values.foreach { sd =>
+ // Materialized means every MAP partition has a registered output: the tracker counts map
+ // outputs, so compare against the producer RDD's partition count (matching how
+ // ShuffleMapStage.isAvailable derives completeness), not the reducer-side partitioner.
+ // This is a point-in-time check at job submission. If a materialized prefix's output were
+ // LOST after this classification but before the pipelined suffix finished (executor loss),
+ // the prefix would need to re-run while the gang holds all slots -- the very deadlock this
+ // shape check forbids. That is safe here for two reasons: (1) the only supported deployment
+ // is single-executor local mode, where executor loss does not occur in normal operation;
+ // and (2) if a FetchFailed did strip the prefix, handleTaskCompletion routes it to a
+ // WHOLE-GROUP abort (the failing stage is a pipelined group member), not a lone-stage
+ // resubmit into the held slots -- the job reruns from scratch rather than deadlocking.
+ if (mapOutputTracker.getNumAvailableOutputs(sd.shuffleId) != sd.rdd.partitions.length) {
Review Comment:
**[efficiency / CONFIRMED]** This materialization loop — and the `regularBoundaries` map feeding it — runs unconditionally on every job submission, but `hasUnmaterializedRegularBoundary` is only ever consumed when `hasPipelined` is true (`isUnsupportedMix` short-circuits otherwise). Every all-regular job on every Spark deployment (feature off, no pipelined manager configured) now pays K `getNumAvailableOutputs` lookups (shuffleStatuses access + read-locked count) plus the boundary HashMap and per-node `(RDD, Boolean)` tuple allocations, inside single-threaded `handleJobSubmitted`.
Fix direction: guard the boundary collection and this loop with `hasPipelined` (skip both entirely when the walk saw no pipelined dependency).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,141 @@
+/*
+ * 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.SparkEnv
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, SparkPlan, UnaryExecNode}
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together. This is the minimal SQL entry point that lets a batch query exercise the
+ * pipelined channel execution path; a production version would be a targeted,
+ * cost/shape-aware replacement rather than a blanket rewrite.
+ *
+ * Enabled only when `spark.sql.pipelinedShuffle.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * Two shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - coalesce over a shuffle: a `CoalesceExec` (user `.coalesce(n)`, a narrow no-shuffle
+ * partition reduction) reading from a shuffle makes ONE reduce task drain SEVERAL reduce
+ * partitions sequentially (a core `CoalescedRDD` over the `ShuffledRowRDD`). The channel
+ * transport cannot serve that -- the map-side writer interleaves all partitions on one
+ * thread and parks on a full bounded queue, so a reader draining partition `start` to
+ * completion before touching `start + 1` deadlocks with the parked writer, and there is no
+ * "combine adjacent partitions" narrow read the transport could substitute without giving
+ * up the bounded-queue backpressure the transport relies on. `coalesce`'s API contract is a
+ * narrow dependency that merges adjacent partitions, so we cannot honor it by re-hashing to
+ * `n` partitions either. So if any shuffle in the plan is read by a coalesce (directly or
+ * through a narrow chain) the rule leaves the whole plan regular; the query runs correctly,
+ * just not pipelined. (Leaving only that one exchange regular would put a pipelined exchange
+ * below a regular boundary, which the scheduler rejects -- so it must be all-or-nothing.)
+ */
+case class EnablePipelinedShuffle() extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
Review Comment:
**[reuse / CONFIRMED]** This three-gate preamble (flag check, local-mode check, `PipelinedChannelShuffleManager` instance check) is duplicated verbatim in `AQEEnablePipelinedShuffle.apply`, differing only in log wording. It is a correctness gate, not cosmetics: flipping under the RPC streaming manager would skip the row copy (`requiresDetachedRecords = false`) and silently corrupt rows — so the two copies must never drift, yet nothing ties them together. When eligibility changes (a second validated transport, relaxed locality), one rule gets updated and the other silently splits AQE vs non-AQE behavior.
Fix direction: extract one shared predicate (e.g. a small helper object in `execution.exchange`) called by both `apply` methods; each rule keeps only its own plan-shape logic.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,324 @@
+/*
+ * 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
+
+ // Per-run epoch (the jobId), read from the job-level local property the DAGScheduler set for a
+ // pipelined job. The reader of this gang reads the SAME value, so both address the same
+ // per-run queues in the process-wide rendezvous; a re-run of this shuffleId is a different job
+ // and gets a different epoch, keeping its queues physically separate. Absent (a core-RDD test
+ // path that never sets it, and never re-runs a shuffleId concurrently) means epoch 0.
+ private val runEpoch = ChannelShuffleRendezvous.epochOf(TaskContext.get())
+
+ // 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, runEpoch, pid)
+ val start = System.nanoTime()
+ while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, runEpoch, pid)) {
+ // Wake on a task kill even without thread interruption. `isAbandoned` is set only by a
+ // reduce task that actually STARTED (its completion listener); if the reader for pid never
+ // started -- the group aborts while this producer is already filling queues (an
+ // unserializable consumer task, an early failure of another member, a job cancel) -- the
+ // mark never appears and this offer loop would park forever, pinning the executor slot
+ // (spark.job.interruptOnCancel defaults to false, so the kill does not interrupt the
+ // thread). Checking the TaskContext interrupt flag each cycle is the symmetric escape to
+ // the reader's takeItem.
+ Option(TaskContext.get()).foreach(_.killTaskIfInterrupted())
+ 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, runEpoch, 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 = {
+ // No stale-state reset is needed here: this run's queues and abandoned marks are keyed by
+ // runEpoch (the jobId), so an EARLIER run of this shuffleId (a RangePartitioner sampling job
+ // then the main job; executeTake batches; a re-executed classic plan) used a different epoch
+ // and its leftovers are physically separate -- this run starts against empty per-epoch state.
+
+ // One in-progress batch per reduce partition, plus its fill count.
+ val batches = Array.fill(numPartitions)(new Array[AnyRef](batchSize))
Review Comment:
**[efficiency / CONFIRMED]** This eagerly allocates `numPartitions x batchSize` object slots per map task, including partitions `liveMask` guarantees are never touched. A 2000-partition repartition allocates ~16MB of empty arrays per map task before the first record; a LIMIT partial read with 1 live partition of 200 allocates 199 arrays the `liveMask(pid)` gate guarantees are never written.
Fix direction: allocate a partition's batch lazily on its first record (one null-check branch on the hot path), or at minimum only for `liveMask`-true partitions.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3815166740
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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)
+
+ // A partition is worth writing only while a consumer still wants its data: it must be in
+ // the job's live set (Half 1: no-reader partitions), AND its reader must not have departed
+ // (Half 2: a live reader that stopped early, e.g. LIMIT). `wants` is re-checked as data
+ // flows because abandonment happens at runtime when the reader task completes.
+ private def wants(pid: Int): Boolean =
+ liveReducePartitions.forall(_.contains(pid)) &&
+ !ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)
+
+ // 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 = {
+ // This is a fresh producer attempt. Clear any abandoned marks left on our partitions by
+ // an EARLIER job that reused this shuffleId (RangePartitioner sampling job -> main job;
+ // executeTake batches) so we do not mistake a prior reader's departure for our own
+ // partitions being dead. Only abandonment happening during THIS attempt then counts.
+ var pc = 0
+ while (pc < numPartitions) {
+ ChannelShuffleRendezvous.clearAbandoned(shuffleId, pc)
Review Comment:
This clears the abandoned mark for every partition of the shuffle, but the marks are process-wide and shared by **all concurrently running map tasks of the same stage** -- so a late-starting map task erases a departure recorded while an earlier sibling was already writing.
LIMIT over a pipelined shuffle (the `PipelinedLimitHangSuite` shape):
1. `m0..m3` and the reduce task for partition 0 are gang-admitted together.
2. `m0` enters `write()` and pushes a full batch.
3. The reduce task satisfies `limit(10)`, completes, and its completion listener calls `abandon(shuffleId, 0)`.
4. `m3` is still deserializing; when it finally enters `write()`, this loop removes the mark for partition 0.
5. `wants(0)` is now permanently true for `m3`: it fills partition 0's 64-batch queue and then spins forever in `putUnlessAbandoned`'s `while (!isAbandoned) q.offer(..., 100ms)` with no drainer.
The map stage never finishes, and since the consumer's completion is deferred until its pipelined producers finish, the job hangs -- the exact failure the abandon mechanism exists to prevent, now timing-dependent rather than fixed. The window is between task launch and entering `write()` (task deserialization, broadcast fetch), which is far longer than the time it takes a reader to satisfy a small LIMIT from one already-enqueued batch.
The reset likely needs to be scoped to a job / stage-attempt generation (e.g. tag marks with the producing stage attempt id and ignore stale ones) rather than "every writer clears every partition".
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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)
+
+ // A partition is worth writing only while a consumer still wants its data: it must be in
+ // the job's live set (Half 1: no-reader partitions), AND its reader must not have departed
+ // (Half 2: a live reader that stopped early, e.g. LIMIT). `wants` is re-checked as data
+ // flows because abandonment happens at runtime when the reader task completes.
+ private def wants(pid: Int): Boolean =
Review Comment:
`wants(pid)` runs once per input record and both halves allocate: `abandoned.contains((shuffleId, pid))` builds a `Tuple2` and boxes `pid`, and `liveReducePartitions.forall(_.contains(pid))` boxes `pid` again for the `Set[Int]` lookup.
On the 20M-row repartition this transport is benchmarked against that is roughly 40M short-lived allocations plus two hash lookups per row, on the hot loop of a design whose entire justification is amortizing per-row cost via batching.
Neither fact needs to be re-read per record. An `Array[Boolean]` of per-partition liveness, seeded from `liveReducePartitions` and refreshed for a partition only when that partition's batch is flushed, reduces the per-record path to a single array load.
##########
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:
This assumes the `ResultStage`'s partition ids **are** the shuffle's reduce partition ids, but that only holds when the result RDD is the `ShuffledRowRDD` itself. Any narrow operator in between that changes the partition count breaks it.
`spark.range(1000).groupBy($"id" % 10).count().coalesce(2).collect()` (feature on, local mode):
1. `getShuffleDependenciesAndResourceProfiles(rs.rdd)` walks the `CoalescedRDD`'s narrow dep down to the pipelined dep, so `resultReadsThisShuffle` is true.
2. The property is set to `"0,1"` -- the `CoalescedRDD`'s partition ids, not reduce partition ids.
3. `ChannelShuffleWriter.wants()` then drops every record with `pid >= 2` and emits no `EndOfStream` for those partitions.
4. But `CoalescedRDD.compute` still pulls the parent iterator for reduce partitions `0..199`, so `ChannelShuffleReader` for `2..199` blocks forever in `q.take()` and the job hangs.
The same class of failure applies to a partial-read `take` over a `UnionExec` whose shuffle branch is not the first child, where result ids are offset from reduce ids.
Neither `EnablePipelinedShuffle` nor `AQEEnablePipelinedShuffle` rejects these shapes, so the check here probably needs to establish that the result stage's RDD *is* this shuffle's `ShuffledRowRDD` (identity), not merely that the dependency is reachable from it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.SparkFunSuite
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+
+/**
+ * End-to-end SQL coverage of the pipelined channel path: a batch query whose hash
+ * exchange is rewritten to a pipelined shuffle (EnablePipelinedShuffle) and served by the
+ * in-process channel manager (PipelinedChannelShuffleManager), run through the
+ * concurrent-stage scheduler on a single executor. Self-manages its SparkSession because the
+ * shuffle manager and AQE-off gate are start-up configs.
+ */
+class PipelinedShuffleSqlSuite extends SparkFunSuite with AdaptiveSparkPlanHelper {
+
+ private def withPipelinedSession(body: SparkSession => Unit): Unit = {
+ val spark = SparkSession.builder()
+ // High task-concurrency cap so the pipelined group's whole-group slot demand (the sum
+ // of every concurrent stage's partitions) is admitted. This is a correctness harness,
+ // not a perf one: on a smaller physical machine these logical slots oversubscribe the
+ // cores, which is fine for verifying results but meaningless for timing.
+ .master("local[16]")
+ .appName("pipelined-shuffle-sql")
+ .config("spark.shuffle.manager.incremental",
+ "org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager")
+ .config("spark.sql.adaptive.enabled", "false") // rule only sees exchanges with AQE off
+ .config("spark.sql.pipelinedShuffle.enabled", "true")
+ .config("spark.speculation", "false")
+ .config("spark.sql.shuffle.partitions", "4")
+ .getOrCreate()
Review Comment:
Building the session with `SparkSession.builder()...getOrCreate()` inside a plain `SparkFunSuite`, without first clearing any pre-existing session, is fragile. Same pattern in `AQEPipelinedShuffleSuite` and `PipelinedLimitHangSuite`.
sql/core suites share a JVM. If an earlier suite in the same fork leaves an active or default `SparkSession` behind, `getOrCreate` returns **that** session and silently ignores every `.config(...)` here -- including `spark.shuffle.manager.incremental` and `spark.sql.pipelinedShuffle.enabled`. The exchanges are then never flipped and `assert(pipelinedExchanges.nonEmpty)` fails with a message that points nowhere near the real cause. Worse, the `finally` block's `spark.stop()` would tear down a session this suite does not own.
Extending `LocalSparkSession` (it does `clearActiveSession()` / `clearDefaultSession()` in `beforeEach` and stops in `afterEach`), or at minimum stopping any live `SparkContext` and clearing both session slots before the builder, removes the ordering dependency.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5418335545
@dongjoon-hyun Thanks for your careful reviews! Could you take another look? Thanks!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5447762696
> This is a major code change with 4000 lines of code diff - I'm not sure we could have a sufficient time duration to bake this in rather than failing the RC every time we realize an issue. What is the impact of deferring this by 3 months?
I'm okay to have this in 4.4 instead of 4.3.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3972727861
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2604,6 +2919,67 @@ private[spark] class DAGScheduler(
.getOrElse(new Properties())
addPySparkConfigsToProperties(stage, properties)
+ // For a pipelined PRODUCER stage, tell its tasks which of its reduce partitions the job
+ // actually reads. 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.
+ //
+ // The live set is per-SHUFFLE-EDGE, not per-job: it is the reduce partitions the consumer of
+ // THIS shuffle reads. liveReduceSet computes it by walking the narrow chain from the result
+ // RDD down to this shuffle, threading the read partition subset through each dependency's
+ // getParents. 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, and
+ // its shuffle is not narrow-reachable from the result RDD (an intervening shuffle blocks the
+ // walk), so liveReduceSet returns None and the property is left unset -- fully live -- which
+ // is correct. See the None handling below for the fail-fast case.
+ stage match {
+ case sms: ShuffleMapStage
+ if isPipelinedProducer(stage) && pipelinedManagerWantsLiveReduceHints =>
+ val resultStage = jobIdToActiveJob.get(jobId).map(_.finalStage)
+ .collect { case rs: ResultStage => rs }
+ resultStage.foreach { rs =>
+ // Tell this pipelined producer's tasks which of ITS reduce partitions the job's readers
+ // will actually drain, so the writer can drop records routed to partitions no consumer
+ // reads (a partial read -- LIMIT / executeTake -- runs only a subset of the result
+ // stage's partitions; feeding the rest fills their bounded queues and deadlocks the
+ // writer). This is the reduce-partition set the result stage's partition subset maps to
+ // through the narrow chain down to this shuffle (see liveReduceSet).
+ liveReduceSet(rs.rdd, rs.partitions.toSet, sms.shuffleDep.shuffleId) match {
Review Comment:
**Hang: the live set ignores cached partitions.**
`liveReduceSet` is derived from the partitions the result stage will *run*, but a result task whose RDD block is cached returns from `RDD.getOrCompute` without ever constructing a `ChannelShuffleReader`, so its queue is never `abandon`ed.
```scala
val df = spark.range(0, 4000000, 1, 4).repartition(8)
df.persist(MEMORY_ONLY); df.collect()
// memory pressure evicts partitions 0-5 only
df.collect()
```
`getMissingParentStages` sees the cache misses and resubmits the producer as a fresh stage (empty `pipelinedCompletedPartitions`) with all 8 partitions live; tasks 6 and 7 are served from the block store and never open a reader; once partition 6/7's queue reaches capacity every map task parks in the `offer` loop forever and no consumer receives `EndOfStream`.
None of `classifyJobShuffleShape`, `pipelinedJobConcurrentTaskDemand`, `liveReduceSet` or `rddReachesShuffle` consult `getCacheLocs`, neither SQL rule has a persist/cache guard, and no test covers persistence. A fully cached plan does not hang but is still counted as phantom admission demand. `InMemoryRelation` caches `executedPlan`, so a `persist()` above a flipped exchange is a normal SQL shape, not an RDD corner case.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together.
+ *
+ * The rewrite is deliberately unconditional rather than cost-based: the scheduler requires a job to
+ * be all-pipelined (or a materialized prefix below a pipelined suffix), so choosing per exchange
+ * would produce exactly the mixed shapes it rejects. What the rule does decide is ELIGIBILITY --
+ * the environment gates in [[PipelinedShuffleEligibility]] and the plan shapes below that force the
+ * whole plan back to regular.
+ *
+ * Enabled only when `spark.sql.shuffle.localPipelined.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * These shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - an UNSUPPORTED CONSUMER reading a shuffle (see [[readsShuffleThroughUnsupportedConsumer]]):
+ * an operator that would drain a shuffle in a way the channel transport cannot serve, or that
+ * builds its own hidden regular shuffle. If such an operator sits above any shuffle the rule
+ * leaves the WHOLE plan regular (leaving only that one exchange regular would put a pipelined
+ * exchange below a regular boundary, which the scheduler rejects -- so it is all-or-nothing).
+ * The query runs correctly, just not pipelined. The unsupported consumers are:
+ * - `CoalesceExec` (user `.coalesce(n)`): its `CoalescedRDD` makes ONE reduce task drain
+ * SEVERAL reduce partitions sequentially. The single-threaded writer parks on a full
+ * bounded queue filling a later partition before emitting an earlier one's markers, so a
+ * reader draining partitions in order deadlocks the writer with no timeout escape;
+ * `coalesce`'s narrow-merge contract also cannot be honored by re-hashing to `n`.
+ * - `CartesianProductExec`: its `UnsafeCartesianRDD` reads each left (child) partition once
+ * per right partition, so N reduce tasks mint N readers on the SAME rendezvous queue for
+ * one `(shuffleId, epoch, pid)` -- rows and end-of-stream markers split
+ * nondeterministically (wrong results), a reader short of `numMaps` markers hangs, and
+ * the first to finish abandons the queue and discards the others' data. The fan-out check
+ * does not catch it (one consumer RDD, computed many times), nor the width-1 require.
+ * - `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec`: each builds a
+ * hidden regular (`pipelined = false`) shuffle inside `doExecute` via
+ * `prepareShuffleDependency`, invisible to this plan walk. A flipped exchange below one of
+ * them would sit under that unmaterialized regular boundary and the job would hard-fail at
+ * submission (`classifyJobShuffleShape`'s pipelined-below-regular rejection). (`.collect()`
+ * on a limit takes `executeTake` and never hits `doExecute`; `.write` / `.toLocalIterator`
+ * / a non-root position do.)
+ */
+object EnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the AQE rule's -- see PipelinedShuffleEligibility for why it is a correctness
+ // gate. It also requires AQE off implicitly: under AQE this rule sees no exchanges.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
+ if (shuffles.isEmpty) return plan
+
+ // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so
+ // leave the whole plan regular rather than produce a rejected job. Check subquery plans
+ // too (plan.exists walks the operator tree only): today no SQL shape can place a reused
+ // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse
+ // never fires because the subquery's own preparation pass (PlanSubqueries ->
+ // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE
+ // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery
+ // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the
+ // second mechanism is an accident of rule ordering and the third is optimizer behavior,
+ // so this gate does not rely on either.
+ if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) {
+ // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output,
+ // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has
+ // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined")
+ // rather than WARN, which would fire on every reuse-bearing query and read as a fault.
+ logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " +
+ "avoid a fan-out pipelined job.")
+ return plan
+ }
+
+ // An operator that would read a shuffle in a way the channel transport cannot serve, or that
+ // builds its own hidden regular shuffle, forces the whole plan regular (see class doc). Like
+ // the reuse fallback this is a normal, expected outcome, so log at DEBUG rather than WARN.
+ if (readsShuffleThroughUnsupportedConsumer(plan)) {
+ logDebug("EnablePipelinedShuffle: a shuffle is read through an operator the channel " +
+ "transport cannot serve (coalesce / cartesian product / a limit operator that builds a " +
+ "hidden shuffle); leaving the plan regular.")
+ return plan
+ }
+
+ plan.transformUp {
+ case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true)
+ }
+ }
+
+ /**
+ * True if any [[ShuffleExchangeExec]] in `plan` is read by an operator the channel transport
+ * cannot serve. The unsupported operators (see class doc for why each is fatal) are
+ * `CoalesceExec`, `CartesianProductExec`, and the limit operators `CollectLimitExec` /
+ * `CollectTailExec` / `TakeOrderedAndProjectExec`. For each such operator anywhere in the plan,
+ * check whether a shuffle is reachable below it.
+ *
+ * The reachability walk descends through EVERY child of a non-exchange node -- not only unary
+ * children -- so a shuffle behind a `UnionExec`/join (a `BinaryExecNode`) beneath the operator
+ * is still found. It stops at the FIRST [[ShuffleExchangeExec]] on each path: a shuffle deeper
+ * than that first one is not read by this operator (the intervening exchange's own reader reads
+ * one reduce partition per task), so it is not this operator's concern.
+ */
+ private def readsShuffleThroughUnsupportedConsumer(plan: SparkPlan): Boolean = {
Review Comment:
**Deadlock: RDD-API narrow consumers over `Dataset.rdd` bypass this guard.**
This only inspects `SparkPlan` nodes, but `Dataset.rdd` / `toRdd` hands the pipelined `ShuffledRowRDD` to user code. An RDD-level `coalesce` (narrow `CoalescedRDD`) drains several reduce partitions per task, and `r.union(r)` / `r.zip(r)` open two readers on the same `(shuffleId, epoch, pid)` queue. Nothing in `DAGScheduler` or the transport rejects either:
- `CoalescedRDD.compute` opens width-1 readers sequentially, so the reader's `require(endPartition - startPartition == 1)` passes.
- `liveReduceSet` maps `{0,1} -> {0..7}` via `NarrowDependency.getParents`, and the fan-out check counts consumer RDD ids (one `ShuffledRowRDD`), so both pass.
Repro (flag on, channel manager, `local[16]`):
```scala
spark.range(0, 2000000, 1, 4).repartition(8).rdd.coalesce(2).count()
```
Task 0 drains partition 0 waiting for `numMaps` end-of-stream markers, while every map task parks forever in `putUnlessAbandoned` on partition 1's full queue (64 x 1024 rows) before it can emit partition 0's marker. Permanent hang, no kill to escape. For `r.union(r)`, the `numMaps` markers are split between two readers so at least one never completes (or one abandons and clears the other's data).
The PR's own `Dataset.rdd` test only applies a narrow `mapPartitionsWithIndex`. Either bail out of the rule when the plan root is `DeserializeToObjectExec` (the `.rdd` boundary), or enforce the "exactly one reader per (shuffle, epoch, partition)" invariant in the rendezvous itself (fail loud on a second `getReader` / on fan-in > 1 in `liveReduceSet`), which would also cover the RDD path.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1129,6 +1189,93 @@ private[spark] class DAGScheduler(
(hasPipelined, hasRegular)
}
+ /** Classify `finalRDD`'s shuffle graph; see [[JobShuffleShape]] for the shape semantics. */
+ private[scheduler] def classifyJobShuffleShape(finalRDD: RDD[_]): JobShuffleShape = {
+ // Cheap pre-pass first: which KINDS of boundary the graph has, over the shared
+ // `traverseRDDGraph` (a HashSet[RDD] visited set, no per-visit allocation). Only a job with
+ // BOTH kinds can be an unsupported mix, and only then are the two below-regular facts
+ // meaningful:
+ // - all-regular (no pipelined dep) => nothing can be pipelined-below-regular;
+ // - all-pipelined (no regular dep) => no regular boundary to be below, or to materialize.
+ // So every job that is not mixed -- which is EVERY job on a deployment that never enables the
+ // feature -- costs exactly what it costs without this feature, instead of paying for the
+ // (RDD, Boolean)-keyed two-context walk and the boundary map below.
+ val (hasPipelinedKind, hasRegularKind) = classifyJobShuffleKinds(finalRDD)
+ if (!hasPipelinedKind || !hasRegularKind) {
+ return JobShuffleShape(
+ hasPipelined = hasPipelinedKind,
+ hasUnmaterializedRegularBoundary = false,
+ hasPipelinedBelowRegular = false)
+ }
+ var hasPipelined = false
+ var pipelinedBelow = false
+ // Frontier regular shuffle boundaries: those reachable from the final RDD WITHOUT crossing
+ // another regular boundary, deduped by shuffle ID. Only these matter for the materialization
+ // check (a regular boundary below another one is never a runnable suffix member).
+ val regularBoundaries = new HashMap[Int, ShuffleDependency[_, _, _]]
+
+ // ONE walk, carrying `belowRegular` (true once the path from the final RDD has crossed a
+ // regular boundary), computes hasPipelined and pipelinedBelow together -- replacing the old
+ // per-boundary rddGraphHasPipelinedDependency re-walks (O(K x graph) on shared ancestors).
+ // A node reachable BOTH above and below a regular boundary must be explored in BOTH contexts:
+ // a pipelined dep under it counts as pipelinedBelow on the below path but not on the above
+ // path. So the visited set is keyed on (RDD, belowRegular), NOT on the RDD alone -- keying on
+ // the RDD alone would let the first-reached context win and drop the other, missing a
+ // pipelined-below-regular dep (a wrongly-accepted job). A node is thus visited at most twice,
+ // keeping the cost O(graph) rather than O(K x graph). hasPipelined is set only above a regular
+ // boundary, matching the old walk (which stopped at boundaries): a below-boundary pipelined
+ // dep is the pipelinedBelow reject case, never a runnable group member.
+ val visited = new HashSet[(RDD[_], Boolean)]
+ val stack = new ListBuffer[(RDD[_], Boolean)]
+ stack += ((finalRDD, false))
+ while (stack.nonEmpty) {
+ val entry = stack.remove(0)
+ val rdd = entry._1
+ val belowRegular = entry._2
+ if (visited.add(entry)) {
+ rdd.dependencies.foreach {
+ case pd: PipelinedShuffleDependency[_, _, _] =>
+ if (belowRegular) pipelinedBelow = true else hasPipelined = true
+ stack.prepend((pd.rdd, belowRegular))
+ case sd: ShuffleDependency[_, _, _] =>
+ // A frontier boundary only when not already below one; descend with belowRegular set.
+ if (!belowRegular) regularBoundaries.getOrElseUpdate(sd.shuffleId, sd)
+ stack.prepend((sd.rdd, true))
+ case narrowDep =>
+ stack.prepend((narrowDep.rdd, belowRegular))
+ }
+ }
+ }
+
+ // The materialization check only matters for a pipelined job: `isUnsupportedMix` consumes
+ // `hasUnmaterializedRegularBoundary` only when `hasPipelined` is true (a pipelined shuffle
+ // below an unmaterialized regular boundary is the rejected shape). A job with no pipelined
+ // dependency -- every job on a feature-off deployment -- would otherwise pay K
+ // getNumAvailableOutputs lookups (a read-locked shuffleStatuses count) for a value never read,
+ // on the single-threaded event loop. So skip the loop entirely unless the walk saw a pipelined
+ // dependency; a non-pipelined job reports hasUnmaterialized = false (unused).
+ var hasUnmaterialized = false
+ if (hasPipelined) {
+ regularBoundaries.values.foreach { sd =>
+ // Materialized means every MAP partition has a registered output: the tracker counts map
+ // outputs, so compare against the producer RDD's partition count (matching how
+ // ShuffleMapStage.isAvailable derives completeness), not the reducer-side partitioner.
+ // This is a point-in-time check at job submission. If a materialized prefix's output were
+ // LOST after this classification but before the pipelined suffix finished (executor loss),
+ // the prefix would need to re-run while the gang holds all slots -- the very deadlock this
+ // shape check forbids. That is safe here for two reasons: (1) the only supported deployment
+ // is single-executor local mode, where executor loss does not occur in normal operation;
+ // and (2) if a FetchFailed did strip the prefix, handleTaskCompletion routes it to a
+ // WHOLE-GROUP abort (the failing stage is a pipelined group member), not a lone-stage
+ // resubmit into the held slots -- the job reruns from scratch rather than deadlocking.
+ if (mapOutputTracker.getNumAvailableOutputs(sd.shuffleId) != sd.rdd.partitions.length) {
Review Comment:
**"the job reruns from scratch" does not hold; a FetchFailed on the prefix poisons the plan.**
A `FetchFailed` raised by pipelined producer P reading materialized prefix R takes the group-atomic branch in `handleTaskCompletion` (`isPipelinedGroupMember(failedStage)`), which calls `unregisterOutputsOnFetchFailedExecutor` -- in local mode that is `removeOutputsOnExecutor("driver")`, stripping *every* regular map output including R's -- and then `abortStage`s without ever resubmitting R.
On the next submission of the same `Dataset` (AQE reuses the materialized `ShuffleQueryStageExec` for R, and the pipelined exchange is kept inline by `createNonResultQueryStages`), this check now sees `getNumAvailableOutputs(R) = 0`, `hasUnmaterialized = true`, and the job is rejected fail-fast with "only supported when every regular shuffle is a fully-materialized prefix". The Dataset cannot run again until re-planned.
This is reachable in local mode: `spark.sql.classic.shuffleDependency.fileCleanup.enabled=true` (`RemoveShuffleFiles`) deletes files but deliberately keeps tracker entries "so that stage retries would be triggered", and it is the default under `Utils.isTesting`. Run a cached mixed-shape Dataset twice: run 2 aborts instead of the contractual R recompute, run 3+ is rejected outright. The new test "losing a materialized prefix's output mid-group aborts the whole group" only asserts the abort. The prefix should be resubmitted (or its outputs left intact) when the group is aborted.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala:
##########
@@ -0,0 +1,190 @@
+/*
+ * 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}
+
+import org.apache.spark.{SparkContext, TaskContext}
+
+/**
+ * Process-wide rendezvous between the map (writer) and reduce (reader) sides of an
+ * in-process pipelined shuffle. One bounded queue exists per
+ * `(shuffleId, epoch, 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.
+ *
+ * The `epoch` is the per-run id (the jobId, propagated to both the writer and the reader of
+ * one gang via a job-level local property -- see `SparkContext.SPARK_PIPELINED_RUN_EPOCH`). A
+ * shuffleId is RE-RUN within one query (a RangePartitioner sample job then the main job;
+ * executeTake's per-batch jobs; a classic Dataset re-executing a reused plan), and each run is
+ * a different job, hence a different epoch. Keying by epoch makes each run's queues and marks
+ * PHYSICALLY separate: the new run never sees a partition whose reader never started in the old
+ * run (whose queue still holds stale batches + end-of-stream markers), and a straggler writer
+ * from an aborted run -- still looping because a task kill need not interrupt the thread -- can
+ * only touch its OWN (old) epoch's queue, never the new run's. This replaces an earlier design
+ * that reused one key per shuffleId and tried to reset shared marks between runs, which left
+ * stale queues and raced stragglers.
+ *
+ * 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
+
+ /**
+ * The per-run epoch for the current task, read from the job-level local property the
+ * DAGScheduler stamps on a pipelined job (`SparkContext.SPARK_PIPELINED_RUN_EPOCH` = jobId).
+ * Both the writer and the reader of one gang read this, so they address the same per-run
+ * queues. Absent (a core-RDD path that never sets it) defaults to 0; that is fine because such
+ * a shuffleId is never re-run into a colliding second live run.
+ */
+ def epochOf(tc: TaskContext): Int =
+ Option(tc)
+ .flatMap(t => Option(t.getLocalProperty(SparkContext.SPARK_PIPELINED_RUN_EPOCH)))
+ .map(_.toInt)
+ .getOrElse(0)
+
+ // State is nested by shuffleId FIRST, then keyed by (epoch, reducePartitionId) within it. The
+ // outer level exists so the two per-shuffle operations the ContextCleaner drives -- holdsShuffle
+ // and removeShuffle -- are O(1) map lookups instead of a scan of every live entry: holdsShuffle
+ // now runs for EVERY shuffle cleaned in a feature-on session, including the regular prefix
+ // shuffles this feature itself produces, so a flat key made that O(live entries) per cleanup.
+ //
+ // Queue values are AnyRef because a queue carries both record batches (Array[AnyRef]) and the
+ // EndOfStream marker.
+ private val queues =
+ new ConcurrentHashMap[Int, ConcurrentHashMap[(Int, Int), LinkedBlockingQueue[AnyRef]]]()
+
+ // (epoch, reducePartitionId) keys, per shuffleId, 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 =
+ new ConcurrentHashMap[Int, java.util.Set[(Int, Int)]]()
+
+ /**
+ * Per-queue capacity in BATCHES (not rows), the backpressure bound and the heap-residency
+ * knob (see spark.shuffle.channel.queueCapacity). Set once by the channel manager at
+ * construction from that conf; defaults to 64 (with the default 1024-row batch, ~64K rows per
+ * reduce partition in flight) until a manager sets it. `@volatile` because the manager sets it
+ * on the driver while writer/reader threads read it.
+ */
+ @volatile private var capacity = 64
+
+ /** Set the per-queue capacity in batches. Called by the channel manager from its conf. */
+ private[pipelined] def setCapacity(batches: Int): Unit = { capacity = batches }
+
+ /** The queue for one `(shuffleId, epoch, reducePartitionId)`, created on first access. */
+ def queue(shuffleId: Int, epoch: Int, reducePartitionId: Int): LinkedBlockingQueue[AnyRef] = {
Review Comment:
**Leak: per-epoch queues and abandon marks are never released while the `Dataset` is alive.**
Queues are created lazily per `(shuffleId, epoch = jobId, pid)`, the reader's completion listener calls `abandon` even on normal completion (adding a mark; `abandon` clears but does not remove the queue), and the only release path is `removeShuffle(shuffleId)`, which `ContextCleaner` drives only when the `ShuffleDependency` is garbage-collected.
`ShuffleExchangeExec.shuffleDependency` is a `@transient lazy val` reused through the cached `executeRDD` (also under AQE's `withFinalPlanUpdate`), so a held `df` keeps the same shuffleId reachable and every `df.collect()` / `df.toLocalIterator()` / `df.rdd.count()` adds `numPartitions` fresh `LinkedBlockingQueue`s, marks and tuple keys (~80-100 KB per run at 200 partitions), unbounded, with `holdsShuffle` true forever. `RemoveShuffleFiles` masks this in tests but the production default is `DoNotCleanup`. (`df.count()` is unaffected since it builds a new `QueryExecution` per call.)
Consider dropping the epoch's entries when the job finishes (e.g. from the job-end cleanup that removes the producer stage), or having the reader's completion listener remove the queue instead of only clearing it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together.
+ *
+ * The rewrite is deliberately unconditional rather than cost-based: the scheduler requires a job to
+ * be all-pipelined (or a materialized prefix below a pipelined suffix), so choosing per exchange
+ * would produce exactly the mixed shapes it rejects. What the rule does decide is ELIGIBILITY --
+ * the environment gates in [[PipelinedShuffleEligibility]] and the plan shapes below that force the
+ * whole plan back to regular.
+ *
+ * Enabled only when `spark.sql.shuffle.localPipelined.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * These shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - an UNSUPPORTED CONSUMER reading a shuffle (see [[readsShuffleThroughUnsupportedConsumer]]):
+ * an operator that would drain a shuffle in a way the channel transport cannot serve, or that
+ * builds its own hidden regular shuffle. If such an operator sits above any shuffle the rule
+ * leaves the WHOLE plan regular (leaving only that one exchange regular would put a pipelined
+ * exchange below a regular boundary, which the scheduler rejects -- so it is all-or-nothing).
+ * The query runs correctly, just not pipelined. The unsupported consumers are:
+ * - `CoalesceExec` (user `.coalesce(n)`): its `CoalescedRDD` makes ONE reduce task drain
+ * SEVERAL reduce partitions sequentially. The single-threaded writer parks on a full
+ * bounded queue filling a later partition before emitting an earlier one's markers, so a
+ * reader draining partitions in order deadlocks the writer with no timeout escape;
+ * `coalesce`'s narrow-merge contract also cannot be honored by re-hashing to `n`.
+ * - `CartesianProductExec`: its `UnsafeCartesianRDD` reads each left (child) partition once
+ * per right partition, so N reduce tasks mint N readers on the SAME rendezvous queue for
+ * one `(shuffleId, epoch, pid)` -- rows and end-of-stream markers split
+ * nondeterministically (wrong results), a reader short of `numMaps` markers hangs, and
+ * the first to finish abandons the queue and discards the others' data. The fan-out check
+ * does not catch it (one consumer RDD, computed many times), nor the width-1 require.
+ * - `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec`: each builds a
+ * hidden regular (`pipelined = false`) shuffle inside `doExecute` via
+ * `prepareShuffleDependency`, invisible to this plan walk. A flipped exchange below one of
+ * them would sit under that unmaterialized regular boundary and the job would hard-fail at
+ * submission (`classifyJobShuffleShape`'s pipelined-below-regular rejection). (`.collect()`
+ * on a limit takes `executeTake` and never hits `doExecute`; `.write` / `.toLocalIterator`
+ * / a non-root position do.)
+ */
+object EnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the AQE rule's -- see PipelinedShuffleEligibility for why it is a correctness
+ // gate. It also requires AQE off implicitly: under AQE this rule sees no exchanges.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
+ if (shuffles.isEmpty) return plan
+
+ // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so
+ // leave the whole plan regular rather than produce a rejected job. Check subquery plans
+ // too (plan.exists walks the operator tree only): today no SQL shape can place a reused
+ // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse
+ // never fires because the subquery's own preparation pass (PlanSubqueries ->
+ // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE
+ // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery
+ // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the
+ // second mechanism is an accident of rule ordering and the third is optimizer behavior,
+ // so this gate does not rely on either.
+ if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) {
+ // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output,
+ // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has
+ // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined")
+ // rather than WARN, which would fire on every reuse-bearing query and read as a fault.
+ logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " +
+ "avoid a fan-out pipelined job.")
+ return plan
+ }
+
+ // An operator that would read a shuffle in a way the channel transport cannot serve, or that
+ // builds its own hidden regular shuffle, forces the whole plan regular (see class doc). Like
+ // the reuse fallback this is a normal, expected outcome, so log at DEBUG rather than WARN.
+ if (readsShuffleThroughUnsupportedConsumer(plan)) {
+ logDebug("EnablePipelinedShuffle: a shuffle is read through an operator the channel " +
+ "transport cannot serve (coalesce / cartesian product / a limit operator that builds a " +
+ "hidden shuffle); leaving the plan regular.")
+ return plan
+ }
+
+ plan.transformUp {
Review Comment:
**Hard failure: a regular RDD shuffle on top of `Dataset.rdd` is rejected with no fallback.**
The rule flips every exchange it can see, but it cannot see consumers layered on `toRdd`. Any RDD-API shuffle above the pipelined `ShuffledRowRDD` yields a pipelined-below-regular job:
```scala
df.groupBy("k").count().rdd.map(r => (r.getLong(0) % 2, 1L)).reduceByKey(_ + _).collect()
```
(likewise `df.rdd.repartition(4)`, `sortByKey`, MLlib `treeAggregate` -> `foldByKey`). `classifyJobShuffleShape` sets `hasPipelinedBelowRegular` and `handleJobSubmitted` fails the job with "A job that mixes a pipelined shuffle dependency with a regular shuffle dependency is only supported when ...". `AQEEnablePipelinedShuffle` flips the same exchange as a free candidate.
The scaladoc's all-or-nothing rationale ("leaving only that one exchange regular would put a pipelined exchange below a regular boundary") is exactly what an RDD-API shuffle violates, and neither the rules nor the scheduler can honor it for that consumer. Suggest treating a `DeserializeToObjectExec` root as an unsupported consumer (leave the plan regular), since that is the `Dataset.rdd` / `toJavaRDD` boundary. Not covered by any test.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -570,7 +569,15 @@ 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. And
+ // skip it when rddWithPartitionIds already copied: needToCopyObjectsBeforeShuffle makes
+ // that RDD emit (pid, row.copy()), so a second copy here would be redundant.
+ shuffleWriterProcessor = createShuffleWriteProcessor(
+ writeMetrics,
+ copyRows = SparkEnv.get.pipelinedShuffleManager.requiresDetachedRecords &&
Review Comment:
**Nit (reuse): second row-copy site, and the new comment inside `needToCopyObjectsBeforeShuffle` is now false.**
Rows are already copied in `rddWithPartitionIds` (L535) when `needToCopyObjectsBeforeShuffle(part)` is true; this adds a second copy site inside the write processor keyed on `requiresDetachedRecords && !needToCopyObjectsBeforeShuffle(part)`. The channel path's copy decision thus depends on the *blocking* manager's sort/bypass thresholds, an unrelated proxy for "did the earlier map already copy". The comment added at L315-317 ("A pipelined shuffle ... does not go through here") is contradicted by the two calls on this path.
Giving `needToCopyObjectsBeforeShuffle` a `pipelined: Boolean` parameter that returns `true` when `pipelinedShuffleManager.requiresDetachedRecords` would keep a single copy site and let `copyRows`, the `write` override and the comment go.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleManager.scala:
##########
@@ -0,0 +1,154 @@
+/*
+ * 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)
+
+ // Per-queue depth in batches (backpressure bound + heap-residency knob). Set the process-wide
+ // rendezvous from the conf at construction, before any writer/reader creates a queue.
+ ChannelShuffleRendezvous.setCapacity(conf.get(config.SHUFFLE_PIPELINED_CHANNEL_QUEUE_CAPACITY))
+
+ override def usesStreamingShuffleOutputTracker: Boolean = false
Review Comment:
**Multi-job consumers re-run the entire producer once per job.**
Because no tracker retains a pipelined shuffle's outputs and availability lives only in the stage-local `pipelinedCompletedPartitions`, a re-created producer stage always starts from zero, whereas a regular shuffle re-creates its stage already available and skips it. Every job over the same pipelined shuffle therefore re-runs the whole producer subtree:
- `Dataset.toLocalIterator` (`executeToIterator` issues one job per partition, and the rule's scaladoc explicitly lists `.toLocalIterator` as a supported path): with 200 shuffle partitions the scan/join/partial-aggregate producer runs 200 times, side-effecting UDFs fire 200 times.
- `df.rdd.take(n)` (growing batches), `df.rdd.toLocalIterator`, the `RangePartitioner` sample job + main job for `orderBy` over an upstream pipelined shuffle, `BroadcastNestedLoopJoin`'s `executeTake(1)` probe.
The live-reduce hint restricts what the producer *emits*, not what it *computes*. Only the limit-operator consumers are excluded by the rules; nothing documents the N-x cost and no test exercises these paths. Worth either excluding `toLocalIterator`-style consumers or documenting the cost in the config description.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,229 @@
+/*
+ * 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.adaptive
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{BinaryExecNode, CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.exchange.{PipelinedShuffleEligibility, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.ShuffledJoin
+
+/**
+ * Opt-in (SPARK-57399). Flips eligible [[ShuffleExchangeExec]]
+ * nodes to `pipelined = true` under AQE, the adaptive counterpart of the non-AQE
+ * `EnablePipelinedShuffle` preparation rule (which is a no-op once the plan is wrapped in
+ * `AdaptiveSparkPlanExec`). Runs in `AdaptiveSparkPlanExec.queryStagePreparationRules`, so
+ * it is re-applied on every replanning round; the decision is deterministic on plan shape,
+ * and already-flipped exchanges are left alone.
+ *
+ * Placement policy. A flipped exchange has no
+ * map output statistics (it never materializes as a query stage: the DAGScheduler
+ * gang-runs it inline with its consumer in the final job -- see the pipelined case in
+ * `AdaptiveSparkPlanExec.createNonResultQueryStages`), so only exchanges whose statistics
+ * no AQE decision consumes are flipped:
+ *
+ * - "free" candidate: the path from the candidate to the plan root crosses no
+ * stats-sensitive node ([[BinaryExecNode]], another [[ShuffleExchangeExec]], or a
+ * query stage). Its own coalescing/skew handling is given up; nothing above needed its
+ * stats.
+ * - "join-paired" candidate: the immediate shuffle inputs of a [[ShuffledJoin]] whose
+ * path to the root is otherwise free, flipped only as a symmetric pair (an asymmetric
+ * flip would leave one side participating in AQE coalesce/skew and the other fixed).
+ * - everything else stays regular and materializes as usual -- those stages form the
+ * fully-materialized prefix the scheduler's mixed-job shape requires.
+ *
+ * A pipelined exchange supports every
+ * partitioning, so a SinglePartition exchange in a free position is simply a candidate
+ * itself. The walk stops below a flipped candidate: exchanges underneath stay regular and
+ * keep full AQE treatment. Candidates whose canonicalized form occurs more than once in
+ * the plan (including inside materialized stages and subqueries) are skipped: flipping
+ * them would trade AQE's stage reuse for duplicate recomputation, and a pipelined producer
+ * cannot be consumed twice.
+ */
+object AQEEnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the non-AQE rule's -- see PipelinedShuffleEligibility for why it is a
+ // correctness gate that must not drift between the two rules.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ flipEligibleExchanges(plan)
+ }
+
+ /**
+ * The plan-shape core of the rule, factored out of [[apply]]'s environment guards (opt-in flag,
+ * local mode, channel manager) so it can be unit-tested on a hand-built plan directly. Collects
+ * the eligible exchanges and returns the plan with each flipped to `pipelined = true`.
+ */
+ private[adaptive] def flipEligibleExchanges(plan: SparkPlan): SparkPlan = {
+ val shared = if (conf.exchangeReuseEnabled) duplicatedShuffleForms(plan) else Set.empty[Any]
+ // Collect the exchanges to flip BY IDENTITY (SparkPlan.id, unique per instance), not by the
+ // node itself: TreeNode overrides hashCode but not equals, so a HashSet[ShuffleExchangeExec]
+ // matches structurally, and the transformDown below would then flip EVERY exchange
+ // structurally equal to a collected one -- including a twin the collector deliberately left
+ // regular on a blocked path. That twin, if it sits below a regular boundary, makes
+ // classifyJobShuffleShape reject the whole job. Keying on the instance id flips exactly the
+ // nodes the collector chose, regardless of spark.sql.exchange.reuse (duplicatedShuffleForms,
+ // the only other guard, is empty when reuse is off). transformDown matches each ORIGINAL node
+ // before rebuilding it, so its id is the same instance id the collector recorded.
+ val toFlip = mutable.HashSet.empty[Int]
+ collectCandidates(plan, blocked = false, shared, toFlip)
+ if (toFlip.isEmpty) return plan
+
+ // transformDown, NOT transformUp: candidates can be nested (a SinglePartition candidate
+ // above a hash candidate). transformUp rebuilds children first, so by the time it
+ // reaches the upper candidate that node is a NEW instance whose (already flipped) child
+ // no longer matches the collected original structurally, and the upper flip is silently
+ // dropped -- leaving a regular exchange above a pipelined one, which the scheduler then
+ // rejects. transformDown hands each candidate to the pattern before its subtree is
+ // rebuilt, so both nested flips apply.
+ plan.transformDown {
+ case s: ShuffleExchangeExec if toFlip.contains(s.id) => s.copy(pipelined = true)
+ }
+ }
+
+ private def isCandidate(s: ShuffleExchangeExec, shared: Set[Any]): Boolean =
Review Comment:
**Usability: with default `spark.sql.shuffle.partitions=200`, most aggregations fail at admission.**
A pipelined exchange is never promoted to a `ShuffleQueryStageExec`, so `CoalesceShufflePartitions` no longer applies, and neither rule compares the exchange width against `sc.maxNumConcurrentTasks`. With both opt-in configs set and everything else default (`local[8]`, 200 partitions):
```sql
SELECT k, count(*) FROM t GROUP BY k
```
has gang demand 200 + N > 8 and fails with `CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT`, whereas with the flag off AQE coalesces 200 down to a few and the query succeeds. I understand fail-loud is a deliberate decision, but the config doc does not tell users they must set `spark.sql.shuffle.partitions` below the local core count, and every suite hides this by pinning `local[16]` with 4 partitions. As written there is no configuration where enabling the feature works without retuning. A width check here (skip the flip and log, like the other shape fallbacks) or an explicit note in the `SQLConf` doc would help.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,229 @@
+/*
+ * 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.adaptive
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{BinaryExecNode, CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.exchange.{PipelinedShuffleEligibility, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.ShuffledJoin
+
+/**
+ * Opt-in (SPARK-57399). Flips eligible [[ShuffleExchangeExec]]
+ * nodes to `pipelined = true` under AQE, the adaptive counterpart of the non-AQE
+ * `EnablePipelinedShuffle` preparation rule (which is a no-op once the plan is wrapped in
+ * `AdaptiveSparkPlanExec`). Runs in `AdaptiveSparkPlanExec.queryStagePreparationRules`, so
+ * it is re-applied on every replanning round; the decision is deterministic on plan shape,
+ * and already-flipped exchanges are left alone.
+ *
+ * Placement policy. A flipped exchange has no
+ * map output statistics (it never materializes as a query stage: the DAGScheduler
+ * gang-runs it inline with its consumer in the final job -- see the pipelined case in
+ * `AdaptiveSparkPlanExec.createNonResultQueryStages`), so only exchanges whose statistics
+ * no AQE decision consumes are flipped:
+ *
+ * - "free" candidate: the path from the candidate to the plan root crosses no
+ * stats-sensitive node ([[BinaryExecNode]], another [[ShuffleExchangeExec]], or a
+ * query stage). Its own coalescing/skew handling is given up; nothing above needed its
+ * stats.
+ * - "join-paired" candidate: the immediate shuffle inputs of a [[ShuffledJoin]] whose
+ * path to the root is otherwise free, flipped only as a symmetric pair (an asymmetric
+ * flip would leave one side participating in AQE coalesce/skew and the other fixed).
+ * - everything else stays regular and materializes as usual -- those stages form the
+ * fully-materialized prefix the scheduler's mixed-job shape requires.
+ *
+ * A pipelined exchange supports every
+ * partitioning, so a SinglePartition exchange in a free position is simply a candidate
+ * itself. The walk stops below a flipped candidate: exchanges underneath stay regular and
+ * keep full AQE treatment. Candidates whose canonicalized form occurs more than once in
+ * the plan (including inside materialized stages and subqueries) are skipped: flipping
+ * them would trade AQE's stage reuse for duplicate recomputation, and a pipelined producer
+ * cannot be consumed twice.
+ */
+object AQEEnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the non-AQE rule's -- see PipelinedShuffleEligibility for why it is a
+ // correctness gate that must not drift between the two rules.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ flipEligibleExchanges(plan)
+ }
+
+ /**
+ * The plan-shape core of the rule, factored out of [[apply]]'s environment guards (opt-in flag,
+ * local mode, channel manager) so it can be unit-tested on a hand-built plan directly. Collects
+ * the eligible exchanges and returns the plan with each flipped to `pipelined = true`.
+ */
+ private[adaptive] def flipEligibleExchanges(plan: SparkPlan): SparkPlan = {
+ val shared = if (conf.exchangeReuseEnabled) duplicatedShuffleForms(plan) else Set.empty[Any]
+ // Collect the exchanges to flip BY IDENTITY (SparkPlan.id, unique per instance), not by the
+ // node itself: TreeNode overrides hashCode but not equals, so a HashSet[ShuffleExchangeExec]
+ // matches structurally, and the transformDown below would then flip EVERY exchange
+ // structurally equal to a collected one -- including a twin the collector deliberately left
+ // regular on a blocked path. That twin, if it sits below a regular boundary, makes
+ // classifyJobShuffleShape reject the whole job. Keying on the instance id flips exactly the
+ // nodes the collector chose, regardless of spark.sql.exchange.reuse (duplicatedShuffleForms,
+ // the only other guard, is empty when reuse is off). transformDown matches each ORIGINAL node
+ // before rebuilding it, so its id is the same instance id the collector recorded.
+ val toFlip = mutable.HashSet.empty[Int]
+ collectCandidates(plan, blocked = false, shared, toFlip)
+ if (toFlip.isEmpty) return plan
+
+ // transformDown, NOT transformUp: candidates can be nested (a SinglePartition candidate
+ // above a hash candidate). transformUp rebuilds children first, so by the time it
+ // reaches the upper candidate that node is a NEW instance whose (already flipped) child
+ // no longer matches the collected original structurally, and the upper flip is silently
+ // dropped -- leaving a regular exchange above a pipelined one, which the scheduler then
+ // rejects. transformDown hands each candidate to the pattern before its subtree is
+ // rebuilt, so both nested flips apply.
+ plan.transformDown {
+ case s: ShuffleExchangeExec if toFlip.contains(s.id) => s.copy(pipelined = true)
+ }
+ }
+
+ private def isCandidate(s: ShuffleExchangeExec, shared: Set[Any]): Boolean =
+ !s.pipelined && !shared.contains(s.canonicalized)
+
+ /**
+ * Top-down walk collecting exchanges to flip. `blocked` is true once the path from the
+ * root has crossed a stats-sensitive node.
+ */
+ private def collectCandidates(
+ plan: SparkPlan,
+ blocked: Boolean,
+ shared: Set[Any],
+ out: mutable.HashSet[Int]): Unit = plan match {
+ case s: ShuffleExchangeExec =>
+ val flipped = !blocked && isCandidate(s, shared)
+ if (flipped) {
+ out += s.id
+ }
+ // A flipped SinglePartition exchange keeps the walk going: AQE makes no decision at
+ // it (it cannot be coalesced or skew-split), so free candidates BELOW it flip too,
+ // forming a pipelined chain: all exchanges in such a chain must flip together (a
+ // SinglePartition exchange cannot be coalesced or skew-split, so AQE makes no decision
+ // at it and free candidates below it flip too). Below any OTHER exchange (flipped or
+ // not) the walk stops: what is underneath either materializes as the prefix or feeds
+ // a regular exchange whose stats AQE uses, and keeps full AQE treatment either way.
+ if (flipped &&
+ s.outputPartitioning == org.apache.spark.sql.catalyst.plans.physical.SinglePartition) {
+ collectCandidates(s.child, blocked = false, shared, out)
+ }
+
+ case _: QueryStageExec => // already materialized; a leaf here
+
+ case j: ShuffledJoin if !blocked =>
+ // Flip the join's immediate shuffle inputs only as a symmetric pair.
+ val leftCandidate = immediateShuffleInput(j.left, shared)
+ val rightCandidate = immediateShuffleInput(j.right, shared)
+ (leftCandidate, rightCandidate) match {
+ case (Some(l), Some(r)) =>
+ out += l.id
+ out += r.id
+ case _ => // asymmetric (a broadcast side, a materialized stage, no clean input): skip
+ }
+ // Anything deeper is below a join input; blocked either way.
+
+ case p if isStatsSensitive(p) =>
+ p.children.foreach(collectCandidates(_, blocked = true, shared, out))
+
+ case p =>
+ p.children.foreach(collectCandidates(_, blocked, shared, out))
+ }
+
+ /**
+ * The single eligible [[ShuffleExchangeExec]] at the top of one join input, looking
+ * through unary non-stats-sensitive forwarders. None if the input is anything else.
+ */
+ private def immediateShuffleInput(
+ plan: SparkPlan,
+ shared: Set[Any]): Option[ShuffleExchangeExec] = plan match {
+ case s: ShuffleExchangeExec => Some(s).filter(isCandidate(_, shared))
+ case _: QueryStageExec => None
+ case p if isStatsSensitive(p) => None
+ case p if p.children.size == 1 => immediateShuffleInput(p.children.head, shared)
+ case _ => None
+ }
+
+ /**
+ * Nodes below which a candidate exchange must NOT be flipped. Reasons a node lands here:
+ * - AQE consumes map output statistics from the stages below it ([[BinaryExecNode]], another
+ * [[ShuffleExchangeExec]]); flipping below it would give up stats a decision needs.
+ * - [[CoalesceExec]] reads its child shuffle multi-partition-per-task (a `CoalescedRDD` over
+ * the `ShuffledRowRDD`), which the channel transport cannot serve; the shuffle it reads
+ * must stay regular.
+ * - the limit operators [[CollectLimitExec]] / [[CollectTailExec]] /
+ * [[TakeOrderedAndProjectExec]] each build a hidden regular (`pipelined = false`) shuffle
+ * inside `doExecute` (via `prepareShuffleDependency`) that no plan walk can see; a flipped
+ * exchange below one of them would sit under that unmaterialized regular boundary and the
+ * job would hard-fail at submission (pipelined-below-regular).
+ * Blocking here keeps the shuffle below -- and everything deeper -- regular, so no pipelined
+ * exchange ends up below the operator's regular boundary. See the non-AQE
+ * `EnablePipelinedShuffle` for the full rationale on each operator.
+ */
+ private def isStatsSensitive(plan: SparkPlan): Boolean = plan match {
Review Comment:
**Nit (reuse): the transport-unsupported operator set is duplicated across the two rules.**
`EnablePipelinedShuffle.isUnsupportedConsumer` lists `CoalesceExec | CartesianProductExec | CollectLimitExec | CollectTailExec | TakeOrderedAndProjectExec`; this method repeats `CoalesceExec` and the three limit operators and only catches `CartesianProductExec` incidentally through the `BinaryExecNode` stats case. `PipelinedShuffleEligibility` was introduced precisely so the two rules cannot drift, so this set belongs there as `isUnsupportedConsumer(plan)`, with `isStatsSensitive` becoming `BinaryExecNode || ShuffleExchangeExec || PipelinedShuffleEligibility.isUnsupportedConsumer(p)`.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2129,8 +2428,24 @@ private[spark] class DAGScheduler(
// Job submitted, clear internal data.
barrierJobIdToNumTasksCheckFailures.remove(jobId)
+ // For a pipelined job, stamp the per-run epoch (the jobId) into the job's properties BEFORE
+ // creating the ActiveJob, so submitMissingTasks -- which clones jobIdToActiveJob(jobId)
+ // .properties per stage -- carries the SAME epoch to every stage's tasks. Both the producer
+ // (writer) and the consumer (reader) of the one gang belong to this job, so both read one
+ // value; a different run is a different job, hence a different epoch, which keys the
+ // in-process channel rendezvous per run (see SPARK_PIPELINED_RUN_EPOCH). Copy the caller's
+ // Properties rather than mutating it. Inert for a non-pipelined job (property never set,
+ // never read).
+ val jobProperties =
+ if (hasPipelined && pipelinedManagerWantsLiveReduceHints) {
+ val p = Utils.cloneProperties(if (properties == null) new Properties() else properties)
+ p.setProperty(SparkContext.SPARK_PIPELINED_RUN_EPOCH, jobId.toString)
Review Comment:
**Behavior change: concurrent actions on one `Dataset` fail with `PIPELINED_SHUFFLE_CROSS_JOB_REUSE`.**
Two concurrent actions on the same `Dataset` share one lazy `shuffleDependency`. The second job's `createResultStage` -> `getOrCreateShuffleMapStage` finds the producer still bound to the first job (`stage.jobIds = {job1}`) and throws `PIPELINED_SHUFFLE_CROSS_JOB_REUSE` (0A000):
```scala
val df = spark.table("t").groupBy("k").count()
Seq(1, 2).par.foreach(_ => df.collect()) // second call fails while the first completes
```
The guard itself predates this PR, but it was only reachable through Real-Time Mode (a fresh plan per micro-batch). Extending pipelined shuffles to batch SQL exposes it to shared DataFrames in notebooks, Thrift Server sessions and thread pools, where cross-job stage sharing is otherwise normal. Sequential re-execution is fine (job cleanup removes the stage before the `JobWaiter` wakes). AQE hits the same path via `ResultQueryStageExec` re-creation over the same `optimizedPlan`.
At minimum this deserves a SQL-level test and a note in the config doc; a cleaner fix would be to key the run epoch by stage attempt rather than job so a second job can wait for or share the running producer.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5513411546
Thanks for the thorough pass -- all ten are addressed. Pushed as `9c40252`, `52b560d`, `51e5396`.
**1. Streaming plans.** `PipelinedShuffleEligibility` now bails out when any node's logical link `isStreaming`, using the same idiom `InsertAdaptiveSparkPlan` uses to keep AQE off streaming plans, so the blanket rule never reaches an `IncrementalExecution` plan and `markStreamingPath` keeps deciding the static side. Pinned by a test that a streaming query in a feature-on session comes out with no pipelined exchange.
**2. Capability gate for the RTM path.** Added `PipelinedShuffleManager.supportsLiveReducePartitionHints` (default `false`), overridden `true` only by the channel manager. It gates both the epoch stamping and the live-set computation, so an RTM job on the streaming manager takes the master path with no cloned `Properties`, no `liveReduceSet`/`rddReachesShuffle`, and no abort. A `DAGSchedulerSuite` test asserts a manager without the capability gets nothing stamped.
**3. The early-out.** You are right that the branch is unreachable -- `initializeShuffleManager` always instantiates it -- so I restored `classifyJobShuffleKinds` as the cheap pre-pass and run the two-context walk only when a job has both kinds. The misleading comment is gone.
One correction to my own claim, since it affects how you read the fix: for all-regular and all-pipelined jobs this is a pure cost restoration with no observable behavioral difference, so no behavior test can prove it. I pinned what is testable instead -- the pre-pass contract for all three non-mixed shapes, plus mix detection, where breaking the fix does make the test fail.
**4. Transport-agnostic cleanup.** Restored `holdsShuffle` on the trait with the channel manager delegating to the rendezvous; `ContextCleaner` no longer names a transport and reads `sc.env.pipelinedShuffleManager`. Answering from the rendezvous rather than a registry is deliberate -- a shuffle unregistered before its job runs recreates its queues -- and that case has a regression test. `outputTrackerMaster` now uses `sc.env` too (`51e5396`).
**5. Rendezvous keying.** Nested by `shuffleId` first, so `holdsShuffle` and `removeShuffle` are O(1) map lookups and `removeShuffle` is one atomic `remove` per map.
**6. The identity assumption.** Replaced with a `ShuffleReducePartitionMapping` trait that the reader RDD implements: `ShuffledRDD` reports identity, `ShuffledRowRDD` answers from the partition spec (width-1 `CoalescedPartitionSpec`, `PartialReducerPartitionSpec`), and anything else returns `None`, which marks the read unmappable and falls back to all-partitions-live. So the scheduler no longer infers a mapping from a count, and the invariant no longer rests on a SQL rule's behavior. Also added the dropped-record debug counter you suggested, as the second line of defense.
**7. Fetch-wait.** Now times the whole wait once and converts at the end, as `withFetchWaitTimeTracked` does.
**8. Benchmark.** Shapes are now derived from the core count so the workflow's 4-core runner can regenerate the file, and the `local[48]` case is dropped.
Worth flagging what this exposed: the demand is not `inputParts + shufflePartitions + 1`. `pipelinedJobConcurrentTaskDemand` adds the map-side width of every distinct pipelined shuffle, and in the two-exchange `groupBy(k).count+orderBy(k)` chain the middle stage's width counts twice, so that shape needs `inputParts + 2*shufflePartitions + 1`. My first attempt at shrinking used the wrong formula and failed admission at 20 slots vs 16 free. Sizing off that worst case, 4 cores fit `1+1+1+1`. A side effect is that the regenerated file now contains the AQE-on chain, which the old shape had been crashing on.
The committed file is a local run, so I have kicked off the benchmark workflow on this branch for JDK 21 and 25 with `create-commit=true`; it will overwrite that file with runner-generated numbers and add the missing `-jdk25-` variant.
**9. Config naming and surface.** Renamed to `spark.shuffle.channel.batchSize` / `spark.shuffle.channel.queueCapacity`. Kept the SQL flag public and made it self-sufficient in the sense you described: it now warns once when the manager does not match, instead of only `logDebug`.
**10. Small things.** The scaladoc sentence is replaced with the actual reason the rule is blanket (the scheduler requires an all-pipelined job, so a per-exchange choice would just create rejected mixes); both rules are `object`s; the `stringArgs` comment is corrected; and the three suites' harness is now one `PipelinedShuffleTestSession` trait. On tearing down a session the suite did not create: the feature is gated on configs that must be present at session creation, so `SharedSparkSession` + `withSQLConf` cannot express it -- I documented that in the trait rather than leaving it implicit, but if you know a precedent I would rather follow it. Comment density is trimmed and the review-history narration moved out.
One item I looked at and did not do, so it does not look like an oversight: making the limit operators' hidden shuffle pipelined-aware (option A) is viable -- the reads are width-1 and `prepareShuffleDependency` already installs the detaching write processor -- but it touches shared `limit.scala` and adds an execution-time mechanism, so I would rather do it as a follow-up JIRA than widen this PR. The current change blocks those operators instead.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "HeartSaVioR (via GitHub)" <gi...@apache.org>.
HeartSaVioR commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5447533734
This is a major code change with 4000 lines of code diff - I'm not sure we could have a good moment to bake this in rather than failing the RC every time we realize an issue. What is the impact of deferring this by 3 months?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5628701893
Pushed `bbb1cefaa94` to address this round.
Registered the Python suite, added RDD-scan exclusions and regular execution for recursive CTE iteration RDDs, memoized the fallback QueryExecution, and strengthened the cache-eviction test. I retained the conservative fallback for unknown file-scan widths and documented that limitation in both SQLConf and the PR description.
Local validation: 143 SQL/cursor tests and the PySpark regression suite passed. The dangling-Python-test check, Scala style checks, and Python lint checks also passed. These results do not imply that the new commit's CI has completed.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5614533527
A clarification on the intended scope of this PR and a change from my earlier response:
**Configuration visibility:** I previously chose to keep `spark.sql.shuffle.localPipelined.enabled` public, with a one-time warning when the configured shuffle manager is incompatible. After reviewing the remaining restrictions, I have reconsidered that decision and marked the flag, along with the channel batch-size and queue-capacity settings, as internal. The feature remains opt-in and disabled by default. The mismatch warning is retained as a diagnostic, but we should treat this as experimental functionality while its supported execution paths are still limited.
**The latest changes address two of the reported limitations:** both AQE and non-AQE planning now retain regular shuffles when the estimated whole-group task demand exceeds local capacity, without changing the requested partition count. `toLocalIterator` now uses a separate regular-shuffle plan so its per-partition jobs can reuse materialized shuffle output. Runtime slot admission still applies, and this does not provide general output reuse for every multi-job consumer.
**Two broader capabilities remain deferred:**
- Concurrent actions sharing the same pipelined exchange are still unsupported. The earlier lifecycle fix protects the owning action when another action is rejected; it does not enable concurrent sharing.
- Recovery from a lost regular-shuffle prefix within a mixed pipelined job is not implemented. The SQL rewrite conservatively keeps such plans regular. Although the lower-level implementation can execute a mixed job with an already materialized prefix, that does not establish a complete recovery contract.
The PR description still presents materialized-prefix mixed execution too broadly. I will update it to distinguish the lower-level capability from the supported SQL behavior and explicitly document these deferred items. My earlier replies should also have distinguished these fallbacks and safeguards more clearly from full support.
The latest changes passed 101 tests and the relevant Scala style checks. Does this narrower, internal-only scope seem reasonable for this PR, with concurrent sharing and mixed-prefix recovery addressed separately?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3817080506
##########
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:
Fixed. The comment now says the transport serves exactly one reduce partition per reader task (matching `ChannelShuffleReader` width-1 `require`) and notes the `require` makes any future wider range fail loud rather than deadlock.
##########
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:
Fixed. The reader now reports `incFetchWaitTime` for the time parked polling the channel -- the read-side backpressure signal -- so "records + read/write time" holds on both sides now.
##########
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:
Fixed. `inputParts` is now derived from `cores` (`min(6, cores - 9)`) so the gang demand stays under the core count, and `runBenchmarkSuite` skips with a clear message below 11 cores (where even the minimum shape does not fit), so checked-in results are only generated where the fair comparison holds.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3832936449
##########
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:
Follow-up: the manager-capability gate I described wasn't quite enough. With the feature on, a session can still run regular shuffles (e.g. a plan the rules leave regular, including the coalesce fallback above); an already-cleaned one reaching this arm would still fire a duplicate `RemoveShuffle`. I've scoped the arm to shuffles the manager actually holds: `PipelinedShuffleManager.holdsShuffle(id)` (the channel manager tracks its registered ids), gated as `!usesStreamingShuffleOutputTracker && holdsShuffle(id)`. A regular shuffle is never held, so it stays a no-op.
##########
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:
Follow-up: this is now subsumed by the `getParents`-based `liveReduceSet` (see the reply above) -- the index mapping is derived from the narrow chain rather than assumed, so an offset spec maps correctly instead of degrading to all-live.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5348578745
Thanks again, @dongjoon-hyun -- another very sharp pass. All ten are addressed in a new commit; the two you flagged as real (the uninterruptible reader and the unscoped `RemoveShuffle`) plus the identity-guard and the cross-context rendezvous are the substantive ones, the rest are the allocation/config/comment/coverage cleanups. Replies inline.
On the not-behind-the-flag one specifically: you are right, thank you -- that was mine to catch and I missed it. The cleaner arm is now gated on a tracker-less pipelined manager being active, so a feature-off deployment sees no extra `RemoveShuffle` or `shuffleCleaned`.
CI is running on the new commit; I will confirm green when it finishes.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya closed pull request #58097: [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor
URL: https://github.com/apache/spark/pull/58097
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5455815109
Thanks -- this was a very useful pass. All ten are addressed; four commits pushed. Notes where my fix differs from the suggestion, or where the cost is worth calling out.
**1. Cartesian N-to-1 read.** Confirmed and fixed. `UnsafeCartesianRDD` computes each child partition once per right partition, so N reduce tasks mint N readers on one `(shuffleId, epoch, pid)` queue -- rows and markers split, a reader short of `numMaps` hangs, and the first to finish abandons the queue under the others. As you note neither the fan-out check (one consumer RDD, computed many times) nor the width-1 `require` catches it. A shuffle read through a `CartesianProductExec` now leaves the plan regular.
**2. coalesce over a binary node.** Confirmed and fixed -- the walk now recurses through every child and stops at the first exchange on each path. Worth recording how the regression test went: my first attempt used two identical `repartition($"k")` branches and passed even with the old guard, because exchange reuse collapsed them and the reuse gate bailed out first, and ~25K rows per reduce partition never filled the 64-batch queue. The committed test uses two structurally different shuffles (`groupBy` vs `repartition`) and enough rows to park the writer; it deadlocks with the old guard and passes with the new one.
**3. Hidden shuffles in the limit operators.** Confirmed and fixed by blocking all three in both rules -- but the block is broader than the hazard, and that costs real coverage: `limit(n).collect()` shapes are no longer pipelined even though `executeCollect` never builds the hidden shuffle. I tried to narrow it to "not at the plan root" and measured that this does not work: for `groupBy.count().orderBy(c).limit(5)`, `TakeOrderedAndProject` sits at the root both for `.collect()` (safe) and for `.write.parquet(...)` (hazardous), and the write path re-plans to the same root shape -- the rule sees an identical plan and cannot know the action. So the broad block is deliberate.
Your alternative (emit a pipelined-aware dependency) does look feasible: `prepareShuffleDependency`'s `pipelined` branch already installs the detaching write processor, the hidden shuffle is `SinglePartition` so the width-1 read is satisfied trivially, and `doExecute` already holds `child.execute()` so it could inspect the child RDD graph for a `PipelinedShuffleDependency` rather than guessing from the conf. I left it out of this PR because it changes three shared operators in `limit.scala` and introduces execution-time creation of pipelined dependencies -- a mechanism worth its own review. Happy to open a follow-up JIRA if you agree.
One consequence I should flag: `PipelinedLimitHangSuite`'s shapes are now regular, so it could no longer exercise the early-stop path and would have passed vacuously. It now asserts the fallback instead, and the live-reduce-set/early-stop machinery moved to three RDD-level partial-read tests in `PipelinedChannelShuffleSuite` (identity, union offsets, zip fan-in). I verified each fails -- hangs to the deadline -- when `liveReduceSet` is deliberately broken.
**4. Unmemoized reachability.** Fixed. Reachability is computed once by a memoized, iterative (two-phase post-order) `rddReachesShuffle` that returns its memo, and `liveReduceSet` is now an explicit worklist keyed on `(rdd -> live subset)`. Merging live sets at a shared node is equivalent to the old per-branch recursion plus union because `getParents` distributes over union; a node is re-enqueued only when its set actually grew (set inequality, not size), which terminates in a finite index domain.
**5. Cleanup scoped to the manager, not the shuffle.** Fixed, and thank you -- the registry I had was worse than you describe: it is dropped at `unregisterShuffle` and never re-added, so after the unregister-before-run sequence the recreated queues leaked permanently. The arm now keys off `ChannelShuffleRendezvous.holdsShuffle`, i.e. the thing that actually holds the state to free. The manager registry and the `holdsShuffle` trait method are gone.
**6. Per-boundary materialization loop.** Fixed -- it runs only when the walk saw a pipelined dependency. I also added the early-out you would probably have asked for next: `classifyJobShuffleShape` returns immediately when no pipelined manager is configured, so a default cluster no longer pays for the graph walk itself on every job submission.
**7. Config versions.** Fixed: the three new configs move to 4.4.0. Which also answers your RC1 question -- with branch-4.3 cut this cannot land in 4.3.0, so I have set them to the release it would actually first ship in.
**8. Eager batch allocation.** Fixed -- allocated per partition on first record, and a handed-off batch leaves `null` behind rather than eagerly allocating a successor.
**9. Duplicated eligibility gate.** Fixed -- extracted to `PipelinedShuffleEligibility`, used by both rules.
**10. `MultiShuffleManager` + channel manager.** Confirmed; the `Some` arm now consults `usesStreamingShuffleOutputTracker` rather than mere tracker presence, so the documented invariant and the routing-suite assertion now match the behavior. My earlier "unreachable by construction" comment missed the `blockingIsMulti` clause -- removed.
Tests: core 264 and the SQL pipelined suites 26, all green.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5443344474
BTW, is this new feature still aiming Apache Spark 4.3.0 whose RC1 is scheduled next week by @HeartSaVioR ?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5612379825
Thanks for the review. Pushed the changes as `e569cd39841`.
The main change is a more conservative eligibility policy: RDD boundaries, cache construction/reads, and mixed AQE shuffle plans retain regular execution. In particular, this gives up the materialized-prefix/pipelined-tail optimization for now, so repeated actions can use the existing regular-shuffle recovery path.
I also added job-scoped epoch cleanup, including protection against late tasks recreating completed state and against one action’s cleanup removing another action’s active queues.
Concurrent sharing of a pipelined exchange remains unsupported, and multi-job consumers still recompute their pipelined producers. Those limitations, along with the slot requirements, are now documented and covered by tests.
Validation: 269 core/scheduler tests and 34 SQL tests passed. Scalastyle passed for the affected modules. Details below in the individual threads.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5614992231
> I also prefer the narrower scope. BTW, could you resolve the conflicts, @viirya ?
Yes, let me resolve the conflicts. Thanks!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3816668638
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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)
+
+ // A partition is worth writing only while a consumer still wants its data: it must be in
+ // the job's live set (Half 1: no-reader partitions), AND its reader must not have departed
+ // (Half 2: a live reader that stopped early, e.g. LIMIT). `wants` is re-checked as data
+ // flows because abandonment happens at runtime when the reader task completes.
+ private def wants(pid: Int): Boolean =
+ liveReducePartitions.forall(_.contains(pid)) &&
+ !ChannelShuffleRendezvous.isAbandoned(shuffleId, pid)
+
+ // 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 = {
+ // This is a fresh producer attempt. Clear any abandoned marks left on our partitions by
+ // an EARLIER job that reused this shuffleId (RangePartitioner sampling job -> main job;
+ // executeTake batches) so we do not mistake a prior reader's departure for our own
+ // partitions being dead. Only abandonment happening during THIS attempt then counts.
+ var pc = 0
+ while (pc < numPartitions) {
+ ChannelShuffleRendezvous.clearAbandoned(shuffleId, pc)
Review Comment:
Fixed, exactly your diagnosis -- a late map task erasing a departure a sibling reader recorded. The per-`write()` reset is gone; the reset now happens once, in the DAGScheduler when it submits the producer stage (before any map task of that run starts), via a new `PipelinedShuffleManager.onPipelinedProducerStageSubmit` hook that the channel manager implements. That point has no live task of the run, so it cannot race the run own writers or readers, while still resetting a prior run marks (the sample-job/main-job and executeTake re-runs). Verified `PipelinedLimitHangSuite` completes in both AQE modes.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,272 @@
+/*
+ * 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)
+
+ // A partition is worth writing only while a consumer still wants its data: it must be in
+ // the job's live set (Half 1: no-reader partitions), AND its reader must not have departed
+ // (Half 2: a live reader that stopped early, e.g. LIMIT). `wants` is re-checked as data
+ // flows because abandonment happens at runtime when the reader task completes.
+ private def wants(pid: Int): Boolean =
Review Comment:
Fixed. The static live set is now a precomputed `Array[Boolean]` seeded once from the property, so the per-record path is a single array load with no boxing or `Tuple2`. Abandonment (the dynamic half) is not re-read per record -- it is checked at hand-off in `putUnlessAbandoned`, which is correct because accumulating a few more rows into an in-memory batch for a since-abandoned partition is harmless (the batch is never put).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3817074134
##########
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:
Fixed exactly as you suggested. `advance()` now polls with a 100ms timeout and calls `killTaskIfInterrupted()` each cycle, so a reader whose producers died wakes on the group abort interrupt flag instead of parking forever and pinning the slot -- the symmetric escape to the writer `putUnlessAbandoned`. Added a regression test: a producer map task that throws now fails the job promptly (under a deadline) rather than hanging.
##########
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:
Fixed. `PipelinedChannelShuffleManager.stop()` (reached from `SparkEnv.stop()`) now clears `ChannelShuffleRendezvous`, so a new `SparkContext` in the same JVM cannot drain rows or end-of-stream markers a previous context left under a reused shuffle id. `clearForTesting` is promoted to `clear()`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3817077723
##########
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:
Good catch on the offset spec. `readsShuffleByIdentity` now also requires `rdd.partitions.length == sd.partitioner.numPartitions` at the shuffle node, so an offset `CoalescedPartitionSpec` (which passes the width-1 `require` but breaks index identity) degrades to "all partitions live" instead of dropping a partition and hanging. As you note nothing produces that today, so this is defense-in-depth.
##########
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:
Fixed. Added the positive case to `PipelinedShuffleRoutingSuite`: a pipelined dep with a no-tracker manager routes to that manager, creates no `StreamingShuffleOutputTracker`, and is absent from the `MapOutputTracker` -- pinning the relaxed successor to the deleted invariant, and exercising the `None` arm of `outputTrackerMaster`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5361867213
Thanks, @dongjoon-hyun. Both fixed in a new commit.
The writer now checks `killTaskIfInterrupted` each cycle, same as the reader `takeItem`. And the rendezvous is now keyed by a per-run epoch (the jobId), so a re-run of the same shuffleId uses physically separate queues -- a straggler writer from an aborted run cannot cross into the next run, and a re-run never drains leftover batches or counts leftover end-of-stream markers. That let me delete the between-run mark reset (and the `onPipelinedProducerStageSubmit` hook) entirely, since there is nothing to clear once runs are isolated. Replies inline.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5363201382
Thanks for the detailed write-up in the PR description -- the inline rationale (especially in `ChannelShuffleRendezvous` and the `epoch` design) made this much easier to follow. I reviewed the full diff and ran the new suites locally. One confirmed deadlock plus three smaller items below.
## 1. Deadlock: partial-read job whose result RDD reaches the pipelined shuffle through a non-identity chain
`submitMissingTasks` only stamps `SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS` when `readsShuffleByIdentity` holds (`DAGScheduler.scala:2736`), i.e. the result RDD reaches the shuffle through a chain of *single* `OneToOneDependency` hops. When it does not hold, the property is left unset and the writer reads that as "every reduce partition is live".
The comment there argues this is safe because "that operator drains every reduce partition". That is true for a **full** read, but not for a **partial** one. `executeTake`/LIMIT runs its first job on partition 0 only, so the remaining reduce partitions have no reduce task at all -- not a reader that stopped early (which `abandon` covers), but one that never started, so the abandoned mark never appears and `putUnlessAbandoned`'s `q.offer(..., 100ms)` loop has no exit.
A join is the easy case to hit: the result RDD is a `ZippedPartitionsRDD2` with two dependencies, so the match falls through to `case _ => false`.
```scala
// local mode, spark.sql.pipelinedShuffle.enabled=true, channel manager, shuffle.partitions=4
left.join(right, $"k" === $"k2").limit(10).collect() // hangs
```
I added this to `PipelinedLimitHangSuite` (purely additive; existing helpers untouched):
```scala
private def joinLimitOverPipelinedCompletesWithin(seconds: Int, aqe: Boolean): Boolean = {
val pool = Executors.newSingleThreadExecutor()
val fut = pool.submit(new Runnable {
override def run(): Unit = withSession(aqe) { spark =>
import spark.implicits._
// Force a shuffled join: a broadcast side would leave only one exchange, and the
// AQE rule flips a join's inputs only as a symmetric pair.
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
spark.conf.set("spark.sql.adaptive.autoBroadcastJoinThreshold", "-1")
// Distinct row counts and column names so the two exchanges cannot canonicalize
// alike: a ReusedExchangeExec would make EnablePipelinedShuffle leave the whole plan
// regular and the query would pass for the wrong reason.
val left = spark.range(0, 4000000L, 1, 4).withColumn("k", ($"id" % 100))
val right = spark.range(0, 100000L, 1, 4).withColumn("k2", ($"id" % 100))
val rows = left.join(right, $"k" === $"k2").limit(10).collect().length
require(rows == 10, s"expected 10 rows, got $rows")
}
})
try {
fut.get(seconds.toLong, TimeUnit.SECONDS)
true
} catch {
case _: java.util.concurrent.TimeoutException =>
fut.cancel(true)
false
} finally {
pool.shutdownNow()
}
}
test("LIMIT over a pipelined shuffle read through a join completes (AQE off)") {
assert(joinLimitOverPipelinedCompletesWithin(90, aqe = false),
"LIMIT over a joined pipelined shuffle should complete, but the writer hung feeding " +
"reduce partitions that executeTake never reads")
}
test("LIMIT over a pipelined shuffle read through a join completes (AQE on)") {
assert(joinLimitOverPipelinedCompletesWithin(90, aqe = true),
"LIMIT over a joined pipelined shuffle should complete, but the writer hung feeding " +
"reduce partitions that executeTake never reads")
}
```
Result -- the two existing tests pass, the two new ones burn the full 90s deadline in both AQE modes:
```
[info] - LIMIT over a pipelined shuffle completes (AQE off) (3 seconds, 136 milliseconds)
[info] - LIMIT over a pipelined shuffle completes (AQE on) (287 milliseconds)
[info] - LIMIT over a pipelined shuffle read through a join completes (AQE off) *** FAILED *** (1 minute, 30 seconds)
[info] - LIMIT over a pipelined shuffle read through a join completes (AQE on) *** FAILED *** (1 minute, 30 seconds)
```
Thread dump of the forked test JVM at the 45s mark confirms it is a deadlock, not slowness -- every live executor thread is parked, each having burned ~0.1s of CPU over 44.8s elapsed:
```
"Executor task launch worker for task 0.0 in stage 0.0 (TID 0)" TIMED_WAITING (parking)
at java.util.concurrent.LinkedBlockingQueue.offer(...)
at ...ChannelShuffleWriter.putUnlessAbandoned(ChannelShuffleWriterReader.scala:116)
at ...ChannelShuffleWriter.write(ChannelShuffleWriterReader.scala:164)
"Executor task launch worker for task 0.0 in stage 2.0 (TID 8)" TIMED_WAITING (parking)
at java.util.concurrent.LinkedBlockingQueue.poll(...)
at ...ChannelShuffleReader$$anon$1.takeItem(ChannelShuffleWriterReader.scala:281)
at ...ChannelShuffleReader$$anon$1.advance(ChannelShuffleWriterReader.scala:298)
```
| thread | state | cpu / elapsed |
|---|---|---|
| stage 0.0 TID 0..3 (producer map tasks) | `putUnlessAbandoned` -> `offer` | 104-116ms / 44.81s |
| stage 2.0 TID 8 (result stage, partition 0) | `takeItem` -> `poll` | 128ms / 44.81s |
Note there is exactly **one** reduce task for 4 shuffle partitions, which is the crux: partitions 1-3 never get a reader, so they never get abandoned, so the writers never reach their end-of-stream loop, so the partition-0 reader never counts `numMaps` markers.
The same reasoning applies to `coalesce(...).limit(n)` and `union(...).limit(n)`.
On the fix: extending `readsShuffleByIdentity` to cover join/union/coalesce means re-deriving each operator's partition mapping, which seems fragile. It may be more robust to invert the fallback. Today "cannot determine the live set" degrades to "everything is live", but the safe degradation is the opposite: if the result stage does not read all of its partitions (`rs.partitions.length != rs.rdd.partitions.length`) and the live set cannot be determined exactly, reject the job fail-fast (or leave the plan regular) rather than risk a hang. Happy to be wrong here if you see a cheaper invariant.
## 2. `AQEEnablePipelinedShuffle`: `toFlip` matches structurally, not by identity
`toFlip` is a `mutable.HashSet[ShuffleExchangeExec]` (`AQEEnablePipelinedShuffle.scala:79`). `TreeNode` overrides `hashCode` but not `equals`, so the case-class structural `equals` applies and `toFlip.contains(s)` at line 91 flips *every* structurally equal exchange, not just the node the collector selected.
`duplicatedShuffleForms` is the only guard, and it is skipped entirely when `spark.sql.exchange.reuse=false` (`val shared = if (conf.exchangeReuseEnabled) ... else Set.empty`). In that configuration a twin on a `blocked` path -- one `collectCandidates` deliberately left regular -- gets flipped too. If it sits below a regular boundary, `classifyJobShuffleShape` then rejects the whole job with `SparkException`.
Keying on `SparkPlan.id`, or an `IdentityHashMap`-backed set, removes the ambiguity regardless of the reuse flag.
## 3. `ContextCleaner`: the new arm is gated on the manager, not on the shuffle
The `else if (tracklessPipelinedManagerActive)` arm (`ContextCleaner.scala:264`) fires for *any* shuffle found in neither tracker once the channel manager is configured -- including regular ones. The comment at 265-277 anticipates the double-clean problem and gates on the manager to avoid it for default deployments, but the gate does not distinguish a regular shuffle from a pipelined one.
That matters here precisely because this feature produces regular shuffles: the materialized prefix of a mixed job. A prefix shuffle unregistered early (`RDD.cleanShuffleDependencies`, or `Dataset.rdd` with `spark.sql.classic.shuffleDependency.fileCleanup.enabled`) will, on later GC, take the new arm and fire `shuffleCleaned` a second time plus an extra `RemoveShuffle` RPC. Tracking channel shuffle ids in the manager, or checking that the dependency was pipelined, would scope it correctly.
## 4. `classifyJobShuffleShape` adds per-boundary graph walks to every job submission
`rddGraphHasPipelinedDependency(sd.rdd)` (`DAGScheduler.scala:1202`) is called once per frontier regular boundary, and each call builds its own `visited` set. The previous `classifyJobShuffleKinds` visited each RDD exactly once.
For a job whose frontier boundaries share a large ancestor subgraph (a wide join or union over a common base), this is O(K x |graph|) on the single-threaded DAGScheduler event loop, before any stage is created -- and it runs for everyone, including users who never enable this feature, where the answer is always `false`. Carrying a `belowRegular` flag through the single existing traversal would restore the original cost.
## Minor notes
- The PR description says "Two registered configs" but there are three -- `spark.shuffle.pipelined.channel.queueCapacity` is missing from the user-facing-change section. The description also mentions a `clearAbandoned` test helper that does not exist in the code.
- `PipelinedChannelShuffleManager.stop()`'s comment says the rendezvous is "keyed only by (shuffleId, reducePartitionId)", which is stale now that `epoch` is part of the key.
Nothing above touches the core design, which I think is sound -- routing by dependency type, the `requiresDetachedRecords` gate combined with the `needToCopyObjectsBeforeShuffle` de-duplication, the per-run `epoch` keying, and the cooperative `killTaskIfInterrupted` escapes on both sides all look right to me.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3982694255
##########
python/pyspark/sql/tests/test_pipelined_shuffle.py:
##########
@@ -0,0 +1,105 @@
+#
+# 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.
+#
+
+import unittest
+
+from pyspark.sql.functions import udf
+from pyspark.sql.types import LongType
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+
+
+class PipelinedShuffleTests(ReusedSQLTestCase):
Review Comment:
**CI blocker: this suite is not registered in `dev/sparktestsupport/modules.py`.**
`dev/run-tests.py` runs `determine_dangling_python_tests` before any test, so every one of the 21 module jobs on the fork (core, sql, hive, all pyspark shards, yarn, Docker integration) stops at the start of "Run tests" with:
```
[error] Found the following dangling Python tests python/pyspark/sql/tests/test_pipelined_shuffle.py
[error] Please add the tests to the appropriate module.
```
Nothing in this PR (Scala or Python) has executed in CI for `022f26af749`. Please add `"pyspark.sql.tests.test_pipelined_shuffle"` to the `pyspark_sql` module's `python_test_goals`.
(The suite itself looks right: `ReusedSQLTestCase` honors the `master()` / `conf()` overrides at `SparkContext` creation, so the manager and flag are in place before the session exists.)
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicBoolean
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, DeserializeToObjectExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.QueryStageExec
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Shared environment gate for the two pipelined-shuffle enabling rules
+ * ([[EnablePipelinedShuffle]] non-AQE and `AQEEnablePipelinedShuffle` under AQE). This is a
+ * CORRECTNESS gate, not cosmetics: flipping an exchange to pipelined while the incremental manager
+ * is still the RPC streaming one would route it to an untested transport that reports
+ * `requiresDetachedRecords = false`, so the SQL layer would skip the row copy and silently corrupt
+ * rows shared across the writer/reader threads. Both rules must apply the identical gate, so it
+ * lives here rather than being copy-pasted into each `apply` (where the two could drift and split
+ * AQE vs non-AQE behavior). Each rule keeps only its own plan-shape logic.
+ */
+private[sql] object PipelinedShuffleEligibility extends Logging {
+
+ // The flag/manager mismatch is a start-up misconfiguration, so warn once per JVM rather than on
+ // every query planned in the session.
+ private val mismatchWarned = new AtomicBoolean(false)
+
+ /** Operators whose consumers cannot safely drain a bounded, single-reader channel. */
+ def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match {
+ case _: CoalesceExec | _: CartesianProductExec | _: DeserializeToObjectExec |
+ _: CollectLimitExec | _: CollectTailExec | _: TakeOrderedAndProjectExec => true
+ case _ => false
+ }
+
+ /**
+ * Estimate the complete group's task demand without executing or materializing the plan.
+ * Each exchange contributes its producer width; the root contributes the consumer width.
+ * Unknown widths conservatively retain regular execution. The scheduler still checks actual
+ * RDD widths and free slots at submission, including contention from other jobs.
+ */
+ def fitsLocalCapacity(plan: SparkPlan, candidates: Set[Int]): Boolean = {
+ val exchanges = plan.collect {
+ case s: ShuffleExchangeExec if s.pipelined || candidates.contains(s.id) => s
+ }.groupBy(_.id).values.map(_.head).toSeq
+ val widths = plan.outputPartitioning.numPartitions +:
Review Comment:
**Practical scope: `outputPartitioning.numPartitions` is 0 for every non-bucketed file scan, so file-backed queries never pipeline.**
`SparkPlan.outputPartitioning` defaults to `UnknownPartitioning(0)`, and non-bucketed `FileSourceScanExec` and non-key-grouped DSv2 scans return exactly that, so `widths.forall(_ > 0)` fails and both rules leave the plan regular for any Parquet/ORC/JSON/CSV/Hive input. Only `spark.range`, `LocalTableScan`, bucketed and key-grouped inputs can ever flip, which is what every suite and the benchmark use, so nothing notices.
```scala
spark.read.parquet(path).repartition(4).collect() // stays regular, DEBUG log only
```
This is safe (0 never underestimates), but the "avoid shuffle-file I/O for local batch workloads" motivation does not hold for real tables in this iteration. Either derive the map width from the scan (`FileSourceScanExec.selectedPartitions` / `DataSourceV2ScanExecBase.partitions` are available at planning time) or state in the config doc and PR description that only in-memory inputs are eligible for now, and add one file-backed test so the behavior is explicit.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala:
##########
@@ -0,0 +1,132 @@
+/*
+ * 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 java.util.concurrent.atomic.AtomicBoolean
+
+import org.apache.spark.SparkEnv
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.config
+import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, DeserializeToObjectExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.adaptive.QueryStageExec
+import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Shared environment gate for the two pipelined-shuffle enabling rules
+ * ([[EnablePipelinedShuffle]] non-AQE and `AQEEnablePipelinedShuffle` under AQE). This is a
+ * CORRECTNESS gate, not cosmetics: flipping an exchange to pipelined while the incremental manager
+ * is still the RPC streaming one would route it to an untested transport that reports
+ * `requiresDetachedRecords = false`, so the SQL layer would skip the row copy and silently corrupt
+ * rows shared across the writer/reader threads. Both rules must apply the identical gate, so it
+ * lives here rather than being copy-pasted into each `apply` (where the two could drift and split
+ * AQE vs non-AQE behavior). Each rule keeps only its own plan-shape logic.
+ */
+private[sql] object PipelinedShuffleEligibility extends Logging {
+
+ // The flag/manager mismatch is a start-up misconfiguration, so warn once per JVM rather than on
+ // every query planned in the session.
+ private val mismatchWarned = new AtomicBoolean(false)
+
+ /** Operators whose consumers cannot safely drain a bounded, single-reader channel. */
+ def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match {
+ case _: CoalesceExec | _: CartesianProductExec | _: DeserializeToObjectExec |
+ _: CollectLimitExec | _: CollectTailExec | _: TakeOrderedAndProjectExec => true
+ case _ => false
+ }
+
+ /**
+ * Estimate the complete group's task demand without executing or materializing the plan.
+ * Each exchange contributes its producer width; the root contributes the consumer width.
+ * Unknown widths conservatively retain regular execution. The scheduler still checks actual
+ * RDD widths and free slots at submission, including contention from other jobs.
+ */
+ def fitsLocalCapacity(plan: SparkPlan, candidates: Set[Int]): Boolean = {
+ val exchanges = plan.collect {
+ case s: ShuffleExchangeExec if s.pipelined || candidates.contains(s.id) => s
+ }.groupBy(_.id).values.map(_.head).toSeq
+ val widths = plan.outputPartitioning.numPartitions +:
+ exchanges.map(_.child.outputPartitioning.numPartitions)
+ val sc = plan.session.sparkContext
+ val slots = sc.maxNumConcurrentTasks(sc.resourceProfileManager.defaultResourceProfile)
+ val demand = widths.map(_.toLong).sum
+ val fits = widths.forall(_ > 0) && demand <= slots
+ if (!fits) {
+ logDebug(s"Pipelined shuffle: estimated stage widths ${widths.mkString(",")} cannot " +
+ s"fit $slots local task slots; retaining regular shuffles.")
+ }
+ fits
+ }
+
+ /**
+ * Whether the pipelined channel transport may be used for `plan` at all, independent of plan
+ * shape. Requires: the opt-in flag on; single-executor local mode (the in-process channel
+ * transport needs producer and consumer in one JVM); and the configured incremental manager
+ * actually being the in-process channel manager. Returns false (with a DEBUG diagnostic) when
+ * any gate fails, so the caller leaves the plan regular.
+ */
+ def enabled(plan: SparkPlan, conf: SQLConf): Boolean = {
+ if (!conf.localPipelinedShuffleEnabled) {
+ return false
+ }
+ if (plan.session == null || !plan.session.sparkContext.isLocal) {
+ return false
+ }
+ // Batch only. `IncrementalExecution.preparations` inherits QueryExecution's list, so without
+ // this gate a streaming plan would be rewritten here: every micro-batch exchange (the
+ // state-store shuffles, the static side of a stream-static join) would be flipped to pipelined
+ // BEFORE `MarkPipelinedShuffleForRealTimeMode` runs. That contradicts what the Real-Time Mode
+ // rule deliberately does -- it leaves the static side regular, because pulling it into the gang
+ // would demand slots for stages that must instead finish first, failing admission. Streaming
+ // marks its own pipelined boundaries; this opt-in batch path must not pre-empt that decision.
+ // (`logicalLink.exists(_.isStreaming)` is the same signal InsertAdaptiveSparkPlan uses to keep
+ // AQE off streaming plans.)
+ // Dataset.rdd exposes arbitrary consumers beyond the SQL plan, including RDD shuffles,
+ // multi-partition reads and repeated reads of the same partition. Cached inputs also hide
+ // shuffle lineage: cache hits skip those readers, while misses may require regular stages.
+ def hasUnsupportedBoundary(p: SparkPlan): Boolean = p match {
+ case _: DeserializeToObjectExec | _: InMemoryTableScanExec => true
Review Comment:
**Remaining gap: `LogicalRDD` inputs and `UnionLoopExec` are still not covered.**
`withRegularShuffle` closes the API exits, but two internal consumers of `queryExecution.toRdd` bypass both it and this node check:
- `RDDScanExec` / `ExternalRDDScanExec` are not listed here. `spark.createDataFrame(rdd.reduceByKey(...)).groupBy("k").count().collect()` (no `.rdd`, so the exit helper is not on the path) flips the SQL exchange above an RDD-level regular shuffle and the scheduler rejects the mixed job at submission. The new test only covers `createDataFrame(rdd).repartition(4).rdd`, which does go through the helper.
- `UnionLoopExec` (recursive CTE) uses `queryExecution.toRdd` directly: it wraps each iteration's RDD in a `LogicalRDD` for the next iteration, references iteration *i*'s RDD both in the final `Union` and through iteration *i+1*'s lineage, and finishes with `df.coalesce(n)`. With the iteration's `REPARTITION` exchange flipped, the final job opens two readers on the same `(shuffleId, epoch, pid)` queues and the `CoalescedRDD` drains several reduce partitions per task, which is the wrong-results / hang shape from the first review. The `CoalesceExec` sits above `RDDScanExec`, not above the exchange, so the consumer check does not see it.
Adding `RDDScanExec | ExternalRDDScanExec` here and applying `withRegularShuffle` in `UnionLoopExec` would close both; a recursive-CTE test with the feature on would pin it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala:
##########
@@ -409,6 +409,23 @@ class QueryExecution(
def assertExecutedPlanPrepared(): Unit = executedPlan
+ /**
+ * RDD consumers and partition-at-a-time iterators can outlive one SQL job. Give them a
+ * separate plan whose shuffles retain output, without changing this execution's plan.
+ * Keep the cloned session for lazy planning and AQE, not just for this method's call.
+ */
+ private[sql] def withRegularShuffle: QueryExecution = {
+ val session = SparkSession.getOrCloneSessionWithConfigsOff(
Review Comment:
**Cost: this clones a `SparkSession` on every call and never closes it.**
With the flag on, `getOrCloneSessionWithConfigsOff` always clones (SessionState plus `ArtifactManager.clone`, which copies the artifact directory when one exists and bumps cached-block refcounts), then a full re-analyze/optimize/plan follows. This now runs on every `javaToPython`, `toPythonIterator`, `toLocalIterator`, checkpoint and cursor open, regardless of whether the plan has a shuffle at all, and the clone is dropped without `close()`.
```python
for _ in range(1000):
df.foreachPartition(f) # one session clone + re-plan per iteration
```
Two cheap mitigations: return `this` when `executedPlan` contains no `ShuffleExchangeExec` with `pipelined = true` (or when `PipelinedShuffleEligibility.enabled` is false for this plan), and memoize the fallback `QueryExecution` in a lazy val so repeated exits on one `QueryExecution` share it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -0,0 +1,728 @@
+/*
+ * 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 java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit}
+
+import org.apache.spark.{PipelinedShuffleDependency, SparkEnv, SparkFunSuite}
+import org.apache.spark.rdd.RDD
+import org.apache.spark.shuffle.local.pipelined.ChannelShuffleRendezvous
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.storage.{RDDBlockId, StorageLevel}
+
+/**
+ * End-to-end SQL coverage of the pipelined channel path: a batch query whose hash
+ * exchange is rewritten to a pipelined shuffle (EnablePipelinedShuffle) and served by the
+ * in-process channel manager (PipelinedChannelShuffleManager), run through the
+ * concurrent-stage scheduler on a single executor. Self-manages its SparkSession because the
+ * shuffle manager and AQE-off gate are start-up configs.
+ */
+class PipelinedShuffleSqlSuite extends SparkFunSuite
+ with AdaptiveSparkPlanHelper with PipelinedShuffleTestSession {
+
+ private def withPipelinedSession(body: SparkSession => Unit): Unit =
+ withPipelinedSession("pipelined-shuffle-sql", aqe = false)(body)
+
+ test("batch repartition($k) runs end-to-end through the pipelined channel shuffle") {
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 10)).repartition($"k")
+
+ // Single action only: a pipelined shuffle is single-shot, so collect exactly once and
+ // derive everything from that one result.
+ val rows = df.select($"id").as[Long].collect()
+ val ids = rows.toSet
+
+ // The rule fired and the exchange is pipelined.
+ val pipelinedExchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec if s.pipelined => s
+ }
+ assert(pipelinedExchanges.nonEmpty,
+ s"expected a pipelined ShuffleExchangeExec; plan was:\n${df.queryExecution.executedPlan}")
+
+ // Correctness: the same 1000 ids, repartitioned, all present exactly once.
+ assert(rows.length === 1000, s"expected 1000 rows, got ${rows.length}")
+ assert(ids === (0L until 1000L).toSet)
+ }
+ }
+
+ test("a single keyed groupBy runs end-to-end through the pipelined channel shuffle") {
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7))
+ .groupBy($"k").count()
+
+ // Single action only (pipelined shuffle is single-shot).
+ val counts = df.as[(Long, Long)].collect().toMap
+ val pipelined = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec if s.pipelined => s
+ }
+ assert(pipelined.nonEmpty,
+ s"expected pipelined exchange; plan:\n${df.queryExecution.executedPlan}")
+ // Each residue class 0..6 of 0..999.
+ val expected = (0L until 1000L).groupBy(_ % 7).map { case (k, vs) => (k, vs.size.toLong) }
+ assert(counts === expected)
+ }
+ }
+
+ test("groupBy with ORDER BY (hash + range exchanges) is all-pipelined") {
+ // A trailing ORDER BY adds a RANGE exchange (global sort with 4 shuffle partitions ->
+ // RangePartitioning) on top of the groupBy's hash exchange. The relaxed rule pipelines
+ // BOTH (a mixed pipelined/regular job would be rejected). Range is the interesting case:
+ // RangePartitioner construction runs a SAMPLE job over the exchange's child -- which here
+ // reads the pipelined hash shuffle -- before the main job runs, so this also exercises
+ // two successive jobs over the same single-shot pipelined producer.
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7))
+ .groupBy($"k").count().orderBy($"k")
+
+ val rows = df.as[(Long, Long)].collect()
+ val exchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.nonEmpty && exchanges.forall(_.pipelined),
+ s"every exchange should be pipelined; plan:\n${df.queryExecution.executedPlan}")
+ // Pin the partitioning shapes so this test can't silently stop covering range.
+ val partitionings = exchanges.map(_.outputPartitioning.getClass.getSimpleName).sorted
+ assert(exchanges.exists(_.outputPartitioning.isInstanceOf[
+ org.apache.spark.sql.catalyst.plans.physical.RangePartitioning]),
+ s"expected a RangePartitioning exchange, got: $partitionings; " +
+ s"plan:\n${df.queryExecution.executedPlan}")
+ // Result is correct AND globally ordered by k.
+ val expected = (0L until 1000L).groupBy(_ % 7).map { case (k, vs) => (k, vs.size.toLong) }
+ .toSeq.sortBy(_._1)
+ assert(rows.toSeq === expected)
+ }
+ }
+
+ test("repartitionByRange (pure range exchange) runs through the pipelined channel shuffle") {
+ // A range exchange directly over the scan: RangePartitioner samples the scan (a job with
+ // no shuffle at all), then the main job runs the pipelined range shuffle. Verifies the
+ // channel transport is agnostic to the partitioner kind, and rows land range-partitioned.
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 100))
+ .repartitionByRange($"k")
+ // spark_partition_id() records which output partition each row landed in without
+ // leaving the DataFrame API (Dataset.rdd would execute a separate QueryExecution).
+ .select($"k", org.apache.spark.sql.functions.spark_partition_id().as("p"))
+
+ val partitioned = df.as[(Long, Int)].collect().map { case (k, p) => (p, k) }
+ val exchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.nonEmpty && exchanges.forall(_.pipelined),
+ s"expected a pipelined exchange; plan:\n${df.queryExecution.executedPlan}")
+ assert(exchanges.exists(_.outputPartitioning.isInstanceOf[
+ org.apache.spark.sql.catalyst.plans.physical.RangePartitioning]),
+ s"expected RangePartitioning; plan:\n${df.queryExecution.executedPlan}")
+
+ // No rows lost, and the partitioning is a genuine range split: key ranges of distinct
+ // partitions must not overlap.
+ assert(partitioned.length === 1000)
+ val ranges = partitioned.groupBy(_._1).map { case (p, rows) =>
+ (p, rows.map(_._2).min, rows.map(_._2).max)
+ }.toSeq.sortBy(_._2)
+ ranges.sliding(2).foreach {
+ case Seq((p1, _, max1), (p2, min2, _)) =>
+ assert(max1 <= min2, s"partitions $p1 and $p2 overlap: max($p1)=$max1 > min($p2)=$min2")
+ case _ =>
+ }
+ }
+ }
+
+ test("sort-merge join (both sides hash-exchanged) is all-pipelined and correct") {
+ // A shuffled join is the last TPC-DS transport shape not yet covered: both join inputs
+ // get a hash ShuffleExchangeExec. Disable broadcast so the join is a SortMergeJoin with
+ // two real shuffles; the relaxed rule pipelines both, and the concurrent-stage group
+ // (two producers + the join stage) runs together.
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
+ // Two structurally DIFFERENT inputs so exchange reuse does not collapse them into one
+ // ReusedExchange (which the rule would skip). Different ranges + key expressions.
+ val left = spark.range(0, 200, 1, 2).withColumn("k", ($"id" % 10))
+ .select($"k", $"id".as("lv"))
+ val right = spark.range(0, 120, 1, 2).withColumn("k", ($"id" % 6))
+ .select($"k", $"id".as("rv"))
+ val joined = left.join(right, "k")
+
+ val rows = joined.select($"k", $"lv", $"rv").as[(Long, Long, Long)].collect()
+ val exchanges = collect(joined.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.length >= 2 && exchanges.forall(_.pipelined),
+ s"both join inputs should be pipelined; plan:\n${joined.queryExecution.executedPlan}")
+
+ // Ground truth: an equi-join on k over the two relations.
+ val l = (0L until 200L).map(i => (i % 10, i))
+ val r = (0L until 120L).map(i => (i % 6, i))
+ val expected = (for ((lk, lv) <- l; (rk, rv) <- r if lk == rk) yield (lk, lv, rv)).toSet
+ assert(rows.toSet === expected)
+ }
+ }
+
+ test("global aggregate (single-partition exchange) runs through the pipelined channel") {
+ // An ungrouped aggregate requires AllTuples, planned as a SinglePartition exchange: the
+ // channel's numPartitions == 1 degenerate case (everything routes to queue 0).
+ withPipelinedSession { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 1000, 1, 2).agg(org.apache.spark.sql.functions.sum($"id"))
+
+ val result = df.as[Long].collect()
+ val exchanges = collect(df.queryExecution.executedPlan) {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.nonEmpty && exchanges.forall(_.pipelined),
+ s"expected a pipelined exchange; plan:\n${df.queryExecution.executedPlan}")
+ assert(exchanges.exists(_.outputPartitioning ==
+ org.apache.spark.sql.catalyst.plans.physical.SinglePartition),
+ s"expected a SinglePartition exchange; plan:\n${df.queryExecution.executedPlan}")
+ assert(result.toSeq === Seq((0L until 1000L).sum))
+ }
+ }
+
+ private def assertRegularRDD(root: RDD[_]): Unit = {
+ val visited = scala.collection.mutable.Set.empty[Int]
+ val pending = scala.collection.mutable.Stack[RDD[_]](root)
+ while (pending.nonEmpty) {
+ val rdd = pending.pop()
+ if (visited.add(rdd.id)) {
+ rdd.dependencies.foreach { dep =>
+ assert(!dep.isInstanceOf[PipelinedShuffleDependency[_, _, _]])
+ pending.push(dep.rdd)
+ }
+ }
+ }
+ }
+
+ for (aqe <- Seq(false, true)) {
+ test(s"typed and Python RDD exports retain regular shuffle lineage with AQE=$aqe") {
+ withPipelinedSession("pipelined-rdd-exports", aqe) { spark =>
+ import spark.implicits._
+ val df = spark.range(0, 4000, 1, 2).repartition(4)
+ assert(df.collect().length === 4000)
+ val typed = df.groupByKey(_ % 4).mapGroups { (key, rows) =>
+ (key, rows.size.toLong)
+ }
+ val rdd = typed.rdd
+ assertRegularRDD(rdd)
+ assert(rdd.coalesce(1).collect().map(_._2).sum === 4000L)
+ val pythonRDD = df.asInstanceOf[org.apache.spark.sql.classic.Dataset[Long]]
+ .javaToPython.rdd
+ assertRegularRDD(pythonRDD)
+ assert(pythonRDD.coalesce(1).count() > 0)
+ val fromRDD = spark.createDataFrame(rdd).repartition(4).rdd
+ assertRegularRDD(fromRDD)
+ assert(fromRDD.coalesce(1).count() === 4L)
+ }
+ }
+
+ test(s"lazy checkpoints retain regular shuffle lineage with AQE=$aqe") {
+ withPipelinedSession("pipelined-checkpoint", aqe) { spark =>
+ withTempDir { dir =>
+ spark.sparkContext.setCheckpointDir(dir.getCanonicalPath)
+ for (reliable <- Seq(false, true)) {
+ val df = spark.range(0, 4000, 1, 2).repartition(4)
+ assert(df.collect().length === 4000)
+ val checkpointed = if (reliable) {
+ df.checkpoint(eager = false)
+ } else {
+ df.localCheckpoint(eager = false)
+ }
+ val rdd = checkpointed.rdd
+ assertRegularRDD(rdd)
+ assert(rdd.coalesce(1).count() === 4000L)
+ assert(checkpointed.collect().sorted === (0L until 4000L).toArray)
+ }
+ }
+ }
+ }
+
+ test(s"SQL cursor computes its producer once with AQE=$aqe") {
+ withPipelinedSession("pipelined-cursor", aqe) { spark =>
+ spark.conf.set("spark.sql.scripting.enabled", "true")
+ spark.conf.set("spark.sql.scripting.cursorEnabled", "true")
+ spark.conf.set("spark.sql.classic.shuffleDependency.fileCleanup.enabled", "false")
+ val evaluated = spark.sparkContext.longAccumulator("cursor producer rows")
+ spark.udf.register("record_cursor_row", (id: Long) => {
+ evaluated.add(1)
+ id
+ })
+ val result = spark.sql(
+ """BEGIN
+ | DECLARE v BIGINT;
+ | DECLARE total BIGINT DEFAULT 0;
+ | DECLARE i INT DEFAULT 0;
+ | DECLARE c CURSOR FOR
+ | SELECT /*+ REPARTITION(4) */ record_cursor_row(id) FROM range(100);
+ | OPEN c;
+ | WHILE i < 100 DO
+ | FETCH c INTO v;
+ | SET total = total + v;
+ | SET i = i + 1;
+ | END WHILE;
+ | CLOSE c;
+ | VALUES (total);
+ |END""".stripMargin)
+ assert(result.collect().head.getLong(0) === (0L until 100L).sum)
+ assert(evaluated.value === 100L)
+ }
+ }
+
+ test(s"Dataset.rdd supports narrow and shuffle consumers with AQE=$aqe") {
+ withPipelinedSession("pipelined-rdd-boundary", aqe) { spark =>
+ val rdd = spark.range(0, 4000, 1, 2).repartition(4).toDF().rdd
+ // Assert eligibility before executing a shape that would hang with the channel.
+ assertRegularRDD(rdd)
+ assert(rdd.coalesce(2).count() === 4000L)
+ assert(rdd.union(rdd).count() === 8000L)
+ assert(rdd.zip(rdd).count() === 4000L)
+ val grouped = rdd.map(r => (r.getLong(0) % 2, 1L)).reduceByKey(_ + _)
+ assert(grouped.collect().toMap === Map(0L -> 2000L, 1L -> 2000L))
+ assert(rdd.repartition(2).count() === 4000L)
+ }
+ }
+
+ test(s"cache construction and partial eviction use regular shuffles with AQE=$aqe") {
+ withPipelinedSession("pipelined-cache", aqe) { spark =>
+ val df = spark.range(0, 4000, 1, 2).repartition(4).persist(StorageLevel.MEMORY_ONLY)
+ try {
+ val classicDf = df.asInstanceOf[org.apache.spark.sql.classic.Dataset[_]]
+ val cached = spark.sharedState.cacheManager.lookupCachedData(classicDf).get
+ .cachedRepresentation.cacheBuilder
+ assert(collect(cached.cachedPlan) {
+ case s: ShuffleExchangeExec if s.pipelined => s
+ }.isEmpty)
+ val expected = (0L until 4000L).toArray
+ assert(df.collect().sorted === expected)
+ val buffers = cached.cachedColumnBuffers
+ assert(buffers.getNumPartitions > 1)
+ SparkEnv.get.blockManager.removeBlock(RDDBlockId(buffers.id, 0))
Review Comment:
**Test strength: this cannot reproduce the original hang.**
4000 rows over 2 map and 4 reduce partitions is about 4 batches per queue against `queueCapacity = 64`, so the writer's `offer` loop never blocks; the test would have passed before the cache fix too. It pins the policy (cached plan has no pipelined exchange) but not the backpressure hang, and it checks only result values, not that partition 0 was actually recomputed.
To make it a real regression test: set `spark.shuffle.channel.queueCapacity` / `batchSize` low (or raise the row count) so the unread partitions would exceed capacity, and assert recomputation with an accumulator on the scan. Same note for the 3x re-execution test in `AQEPipelinedShuffleSuite`: its comment says "A later fetch must recover the prefix", but in local mode `RemoveShuffle` unregisters the tracker entry, so the rerun takes the unavailable-stage recompute path rather than a `FetchFailed`.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala:
##########
@@ -0,0 +1,359 @@
+/*
+ * 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] with org.apache.spark.internal.Logging {
+
+ 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
+
+ // Per-run epoch (the jobId), read from the job-level local property the DAGScheduler set for a
+ // pipelined job. The reader of this gang reads the SAME value, so both address the same
+ // per-run queues in the process-wide rendezvous; a re-run of this shuffleId is a different job
+ // and gets a different epoch, keeping its queues physically separate. Absent (a core-RDD test
+ // path that never sets it, and never re-runs a shuffleId concurrently) means epoch 0.
+ private val runEpoch = ChannelShuffleRendezvous.epochOf(TaskContext.get())
+
+ // 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, runEpoch, pid)
+ val start = System.nanoTime()
+ while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, runEpoch, pid)) {
+ // Wake on a task kill even without thread interruption. `isAbandoned` is set only by a
+ // reduce task that actually STARTED (its completion listener); if the reader for pid never
+ // started -- the group aborts while this producer is already filling queues (an
+ // unserializable consumer task, an early failure of another member, a job cancel) -- the
+ // mark never appears and this offer loop would park forever, pinning the executor slot
+ // (spark.job.interruptOnCancel defaults to false, so the kill does not interrupt the
+ // thread). Checking the TaskContext interrupt flag each cycle is the symmetric escape to
+ // the reader's takeItem.
+ Option(TaskContext.get()).foreach(_.killTaskIfInterrupted())
+ 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, runEpoch, 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 = {
+ // No stale-state reset is needed here: this run's queues and abandoned marks are keyed by
+ // runEpoch (the jobId), so an EARLIER run of this shuffleId (a RangePartitioner sampling job
+ // then the main job; executeTake batches; a re-executed classic plan) used a different epoch
+ // and its leftovers are physically separate -- this run starts against empty per-epoch state.
+
+ // One in-progress batch per reduce partition, plus its fill count. A partition's batch array
+ // is allocated LAZILY, on its first record (batches(pid) starts null), so a map task pays for
+ // only the partitions it actually writes -- a wide shuffle (thousands of partitions) or a
+ // partial read (liveMask leaves most partitions dead) does not eagerly allocate
+ // numPartitions * batchSize empty slots up front.
+ val batches = new Array[Array[AnyRef]](numPartitions)
+ val sizes = new Array[Int](numPartitions)
+ // Records skipped because their reduce partition has no consumer (see liveMask).
+ // Reported at the end of write().
+ var droppedRecords = 0L
+
+ 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)) {
+ // This record is routed to a reduce partition the driver said no consumer reads, so it is
+ // dropped -- see liveMask. Count it: dropping is CORRECT only if the live set was computed
+ // correctly, and everything else here fails loudly (a wrong width fails a require, a
+ // reader-less live partition hangs the writer), while an under-approximated live set would
+ // instead lose rows quietly. A non-zero count at the end of a job whose result looks wrong
+ // is the thread to pull.
+ droppedRecords += 1
+ } else {
+ if (batches(pid) == null) batches(pid) = new Array[AnyRef](batchSize)
+ // 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
Review Comment:
Nit: stale comment. The `ShuffleWriteProcessor` override this refers to was removed in `6df77facce7`; the copy now happens in `ShuffleExchangeExec.prepareShuffleDependency` via `needToCopyObjectsBeforeShuffle(part, pipelined = true)`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5628746999
> +1, LGTM (Pending CIs). Thank you for the thorough follow-ups on every round, @viirya.
Thank you @dongjoon-hyun for the careful review!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3976797848
##########
sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala:
##########
@@ -1570,7 +1570,18 @@ class Dataset[T] private[sql](
/** @inheritdoc */
def toLocalIterator(): java.util.Iterator[T] = {
- withAction("toLocalIterator", queryExecution) { plan =>
+ // This action submits one job per output partition. A channel cannot retain its output
+ // between those jobs, so use a separate regular plan even if queryExecution already ran.
+ // Clone the session to keep the setting off during lazy planning and AQE replanning without
+ // changing the caller's session or its other actions.
+ val iteratorSession = SparkSession.getOrCloneSessionWithConfigsOff(
Review Comment:
**PySpark `DataFrame.rdd` / `DataFrame.toLocalIterator()` bypass both new guards.**
This re-plan and the `DeserializeToObjectExec` gate in `PipelinedShuffleEligibility.hasUnsupportedBoundary` only cover the JVM API. PySpark's `DataFrame.rdd` and `DataFrame.toLocalIterator()` (`python/pyspark/sql/classic/dataframe.py`) go through `Dataset.javaToPython` / `Dataset.toPythonIterator`, which call `queryExecution.toRdd` on the *original* plan: no `DeserializeToObjectExec` is inserted, and this method is never entered, so the pipelined `ShuffledRowRDD` is handed straight to the RDD API. With both opt-in configs on a local PySpark session:
```python
df = spark.range(0, 2000000, 1, 4).repartition(8)
df.rdd.coalesce(2).count() # deadlock: one task drains several reduce partitions
df.rdd.map(lambda r: (r.id % 2, 1)).reduceByKey(lambda a, b: a + b).collect() # pipelined-below-regular rejection
list(df.toLocalIterator()) # producer re-runs once per partition
list(df.toLocalIterator(prefetchPartitions=True)) # job i+1 is submitted while job i owns the producer -> PIPELINED_SHUFFLE_CROSS_JOB_REUSE
```
`PythonRDD.toLocalIteratorAndServe` submits one job per partition and, with prefetch, submits the next one before the current finishes, so the last case fails on the second partition rather than merely recomputing. SQL scripting's `OpenCursorExec` (`executedPlan.executeToIterator()`) has the same exposure.
Since every RDD-side consumer ultimately goes through `QueryExecution.toRdd` / `executedPlan.execute()`, gating there (or in `javaToPython` / `toPythonIterator`, which are the two Python entry points) would close the Python paths and the other JVM exposures (`mapGroups(...).rdd`, `checkpoint(eager = false)`, `LogicalRDD` inputs) at once, instead of enumerating plan nodes. At minimum the PR description's "`Dataset.rdd` ... use regular shuffles" and "`toLocalIterator` uses a separate regular-shuffle plan" should say JVM-only, and a PySpark test for `df.rdd` / `df.toLocalIterator()` would be worth adding.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3980713436
##########
sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala:
##########
@@ -1570,7 +1570,18 @@ class Dataset[T] private[sql](
/** @inheritdoc */
def toLocalIterator(): java.util.Iterator[T] = {
- withAction("toLocalIterator", queryExecution) { plan =>
+ // This action submits one job per output partition. A channel cannot retain its output
+ // between those jobs, so use a separate regular plan even if queryExecution already ran.
+ // Clone the session to keep the setting off during lazy planning and AQE replanning without
+ // changing the caller's session or its other actions.
+ val iteratorSession = SparkSession.getOrCloneSessionWithConfigsOff(
Review Comment:
You are right—the previous guards missed the Python entry points, and `DeserializeToObjectExec` is not a reliable API boundary because optimization can eliminate it.
Fixed in `022f26af749` with a shared `QueryExecution.withRegularShuffle` helper applied at the actual exits:
- JVM Dataset RDD and `toLocalIterator`
- `javaToPython` and `toPythonIterator`
- checkpoint / localCheckpoint
- `OpenCursorExec`
When the feature is enabled, these paths use a separate regular-shuffle plan, including after the original Dataset has already executed. The cloned session keeps pipelining disabled during lazy planning and AQE without changing the caller's setting or original plan. Python iterator execution tracking also uses the fallback QueryExecution.
I did not change generic `SparkPlan.execute()` or `QueryExecution.toRdd`; supported SQL execution also uses those paths. The existing plan-node guard remains conservative, but these API exits no longer depend on it.
Added AQE-on/off regressions for typed RDD exports, RDD-backed DataFrames, lazy checkpoints, and cursors. The PySpark tests exercise coalesce, reduceByKey, and iterator execution with and without prefetch, including after `collect`. The iterator tests confirm that four partition jobs evaluate 1,000 producer rows once per action. All pass.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5660485687
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/6232f2191437d30f2077cdaab06398e789d74e0b
- merged into branch-4.x https://github.com/apache/spark/commit/51ac6a9f0f72d3d28b2bcd3295efd6f15b4237d8
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5443258732
Sorry for missing your ping here, @viirya . Let me review now.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5453968770
Thank you. +1 for the decision to retarget for Apache Spark 4.4.0. Please update the config versions accordingly.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3901749163
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -1269,6 +1256,18 @@ object SQLConf {
.booleanConf
.createWithDefault(true)
+ val PIPELINED_SHUFFLE_ENABLED = buildConf("spark.sql.pipelinedShuffle.enabled")
Review Comment:
Good call on both counts -- renamed to `spark.sql.shuffle.localPipelined.enabled`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3816671413
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.SparkFunSuite
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+
+/**
+ * End-to-end SQL coverage of the pipelined channel path: a batch query whose hash
+ * exchange is rewritten to a pipelined shuffle (EnablePipelinedShuffle) and served by the
+ * in-process channel manager (PipelinedChannelShuffleManager), run through the
+ * concurrent-stage scheduler on a single executor. Self-manages its SparkSession because the
+ * shuffle manager and AQE-off gate are start-up configs.
+ */
+class PipelinedShuffleSqlSuite extends SparkFunSuite with AdaptiveSparkPlanHelper {
+
+ private def withPipelinedSession(body: SparkSession => Unit): Unit = {
+ val spark = SparkSession.builder()
+ // High task-concurrency cap so the pipelined group's whole-group slot demand (the sum
+ // of every concurrent stage's partitions) is admitted. This is a correctness harness,
+ // not a perf one: on a smaller physical machine these logical slots oversubscribe the
+ // cores, which is fine for verifying results but meaningless for timing.
+ .master("local[16]")
+ .appName("pipelined-shuffle-sql")
+ .config("spark.shuffle.manager.incremental",
+ "org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager")
+ .config("spark.sql.adaptive.enabled", "false") // rule only sees exchanges with AQE off
+ .config("spark.sql.pipelinedShuffle.enabled", "true")
+ .config("spark.speculation", "false")
+ .config("spark.sql.shuffle.partitions", "4")
+ .getOrCreate()
Review Comment:
Fixed in all three suites (`PipelinedShuffleSqlSuite`, `AQEPipelinedShuffleSuite`, `PipelinedLimitHangSuite`): each `with...Session` helper now stops any active/default session and clears both slots before the builder, so `getOrCreate()` cannot return a session another suite left in the JVM.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5376850523
Thanks for the thorough pass -- all four plus the minors are addressed. Pushed on top of the branch.
**1. Partial-read deadlock through a non-identity chain.** Your invert-the-fallback suggestion is what I went with, plus a bit more. `submitMissingTasks` now maps the live reduce set through `NarrowDependency.getParents` (`liveReduceSet`), threading the read-partition subset through OneToOne (identity), RangeDependency (union's offset), and join fan-in -- so `join`/`union` + LIMIT map correctly 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 (leaves it regular) rather than risking a hang, exactly as you proposed. Your `coalesce(...).limit(n)` case turned out to deadlock regardless of the live set (one reduce task drains several reduce partitions sequentially against a bounded queue), so a `coalesce` reading a shuffle is now left regular altogether. Added union/join + LIMIT and coalesce-fallback tests.
**2. `toFlip` matches structurally, not by identity.** Fixed -- `toFlip` is now keyed on `SparkPlan.id`. `transformDown` matches each original node before rebuilding it, so its id is the same instance id the collector recorded; a structural twin the collector left on a blocked path is no longer flipped, regardless of `spark.sql.exchange.reuse`. Added a unit test that builds a free exchange plus a structurally identical blocked twin and asserts only the collected one flips (it fails on the old structural key, passes on the id key).
**3. Cleaner arm gated on the manager, not the shuffle.** Fixed -- the channel manager now tracks the shuffle ids it holds, and the tracker-less arm is gated on `!usesStreamingShuffleOutputTracker && holdsShuffle(id)`. A regular shuffle in a feature-on session (including the materialized prefix of a mixed job) is never held, so it stays a no-op there; a channel shuffle's own eager-then-GC re-clean is now a clean no-op too.
**4. Per-boundary graph walks.** Folded into the single existing traversal with a `belowRegular` flag, so the per-boundary re-walks are gone. One caveat worth flagging: a single `belowRegular` bool with RDD-only dedup is unsound under DAG sharing -- a node reachable both above and below a regular boundary would take the first-reached context and could miss a pipelined dep only reachable below (a wrongly-accepted job). So the walk keys its visited set on `(RDD, belowRegular)`: a node is visited at most twice, keeping it O(graph) rather than O(K x graph) while staying correct.
**Minor notes.** The PR description now lists all three configs and drops the stale `clearAbandoned` reference; the two stale comments in the manager are corrected.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5443515612
> BTW, is this new feature still aiming Apache Spark 4.3.0 whose RC1 is scheduled next week by @HeartSaVioR ?
Let's see if we can finish the review before that? If this can land in 4.3.0, it would be good.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3817076228
##########
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:
Fixed, and thank you -- this contradicting "No behavior change by default" was the right call. The else arm is now gated on `!SparkEnv.get.pipelinedShuffleManager.usesStreamingShuffleOutputTracker` (a tracker-less pipelined manager being active), so an already-cleaned regular shuffle reaching it with the feature off is a no-op again -- no extra cluster-wide `RemoveShuffle` RPC, no extra `shuffleCleaned` callback. Kept it a capability check rather than a concrete-type check so the cleaner stays transport-agnostic.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5347919762
Thanks for the very thorough review, @dongjoon-hyun -- especially catching the two hangs; both were reachable on the opt-in path and I have fixed all four. Pushed as a new commit.
- **Live-reduce-partition mapping**: you are right that "reachable" was wrong. Tightening it to a *direct* dependency then broke `LIMIT` (`executeTake`/`collect` wrap the `ShuffledRowRDD` in a `MapPartitionsRDD`), so the discriminator is now an **identity-preserving** narrow chain: `readsShuffleByIdentity` walks `OneToOneDependency` hops only. That admits the collect/take wrapper (1:1, same index) and rejects `coalesce` (a custom `NarrowDependency`) and `union` (a `RangeDependency` offset), which stay fully live and are drained by that operator. Added your `coalesce` case as a regression test.
- **`clearAbandoned` reset**: moved out of `write()` entirely. The DAGScheduler now calls a new `PipelinedShuffleManager.onPipelinedProducerStageSubmit` hook when it submits the producer stage -- one point before any map task of the run starts, so it cannot race the run own writers/readers. The scheduler stays transport-agnostic (it calls the trait, not the concrete transport).
- **`wants` per-record allocation**: the static live set is now a precomputed `Array[Boolean]` (one array load per record); abandonment stays checked at hand-off, not per record.
- **Test session isolation**: the SQL suites now stop and clear any active/default session before `getOrCreate()`.
CI is re-triggered by the push. I will follow up once it is green.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #58097:
URL: https://github.com/apache/spark/pull/58097#issuecomment-5345074525
Could you re-trigger the CI, @viirya ?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3975214459
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together.
+ *
+ * The rewrite is deliberately unconditional rather than cost-based: the scheduler requires a job to
+ * be all-pipelined (or a materialized prefix below a pipelined suffix), so choosing per exchange
+ * would produce exactly the mixed shapes it rejects. What the rule does decide is ELIGIBILITY --
+ * the environment gates in [[PipelinedShuffleEligibility]] and the plan shapes below that force the
+ * whole plan back to regular.
+ *
+ * Enabled only when `spark.sql.shuffle.localPipelined.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * These shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - an UNSUPPORTED CONSUMER reading a shuffle (see [[readsShuffleThroughUnsupportedConsumer]]):
+ * an operator that would drain a shuffle in a way the channel transport cannot serve, or that
+ * builds its own hidden regular shuffle. If such an operator sits above any shuffle the rule
+ * leaves the WHOLE plan regular (leaving only that one exchange regular would put a pipelined
+ * exchange below a regular boundary, which the scheduler rejects -- so it is all-or-nothing).
+ * The query runs correctly, just not pipelined. The unsupported consumers are:
+ * - `CoalesceExec` (user `.coalesce(n)`): its `CoalescedRDD` makes ONE reduce task drain
+ * SEVERAL reduce partitions sequentially. The single-threaded writer parks on a full
+ * bounded queue filling a later partition before emitting an earlier one's markers, so a
+ * reader draining partitions in order deadlocks the writer with no timeout escape;
+ * `coalesce`'s narrow-merge contract also cannot be honored by re-hashing to `n`.
+ * - `CartesianProductExec`: its `UnsafeCartesianRDD` reads each left (child) partition once
+ * per right partition, so N reduce tasks mint N readers on the SAME rendezvous queue for
+ * one `(shuffleId, epoch, pid)` -- rows and end-of-stream markers split
+ * nondeterministically (wrong results), a reader short of `numMaps` markers hangs, and
+ * the first to finish abandons the queue and discards the others' data. The fan-out check
+ * does not catch it (one consumer RDD, computed many times), nor the width-1 require.
+ * - `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec`: each builds a
+ * hidden regular (`pipelined = false`) shuffle inside `doExecute` via
+ * `prepareShuffleDependency`, invisible to this plan walk. A flipped exchange below one of
+ * them would sit under that unmaterialized regular boundary and the job would hard-fail at
+ * submission (`classifyJobShuffleShape`'s pipelined-below-regular rejection). (`.collect()`
+ * on a limit takes `executeTake` and never hits `doExecute`; `.write` / `.toLocalIterator`
+ * / a non-root position do.)
+ */
+object EnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the AQE rule's -- see PipelinedShuffleEligibility for why it is a correctness
+ // gate. It also requires AQE off implicitly: under AQE this rule sees no exchanges.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
+ if (shuffles.isEmpty) return plan
+
+ // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so
+ // leave the whole plan regular rather than produce a rejected job. Check subquery plans
+ // too (plan.exists walks the operator tree only): today no SQL shape can place a reused
+ // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse
+ // never fires because the subquery's own preparation pass (PlanSubqueries ->
+ // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE
+ // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery
+ // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the
+ // second mechanism is an accident of rule ordering and the third is optimizer behavior,
+ // so this gate does not rely on either.
+ if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) {
+ // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output,
+ // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has
+ // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined")
+ // rather than WARN, which would fire on every reuse-bearing query and read as a fault.
+ logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " +
+ "avoid a fan-out pipelined job.")
+ return plan
+ }
+
+ // An operator that would read a shuffle in a way the channel transport cannot serve, or that
+ // builds its own hidden regular shuffle, forces the whole plan regular (see class doc). Like
+ // the reuse fallback this is a normal, expected outcome, so log at DEBUG rather than WARN.
+ if (readsShuffleThroughUnsupportedConsumer(plan)) {
+ logDebug("EnablePipelinedShuffle: a shuffle is read through an operator the channel " +
+ "transport cannot serve (coalesce / cartesian product / a limit operator that builds a " +
+ "hidden shuffle); leaving the plan regular.")
+ return plan
+ }
+
+ plan.transformUp {
+ case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true)
+ }
+ }
+
+ /**
+ * True if any [[ShuffleExchangeExec]] in `plan` is read by an operator the channel transport
+ * cannot serve. The unsupported operators (see class doc for why each is fatal) are
+ * `CoalesceExec`, `CartesianProductExec`, and the limit operators `CollectLimitExec` /
+ * `CollectTailExec` / `TakeOrderedAndProjectExec`. For each such operator anywhere in the plan,
+ * check whether a shuffle is reachable below it.
+ *
+ * The reachability walk descends through EVERY child of a non-exchange node -- not only unary
+ * children -- so a shuffle behind a `UnionExec`/join (a `BinaryExecNode`) beneath the operator
+ * is still found. It stops at the FIRST [[ShuffleExchangeExec]] on each path: a shuffle deeper
+ * than that first one is not read by this operator (the intervening exchange's own reader reads
+ * one reduce partition per task), so it is not this operator's concern.
+ */
+ private def readsShuffleThroughUnsupportedConsumer(plan: SparkPlan): Boolean = {
Review Comment:
Addressed by keeping plans containing `DeserializeToObjectExec` regular in both SQL rules. This covers the `Dataset.rdd` boundary, after which SQL planning cannot see or constrain the RDD consumers.
The guard is deliberately conservative: it can also exclude other typed Dataset operations containing that node. Added AQE-on/off tests that assert the RDD lineage has no pipelined dependency before exercising `coalesce`, self-`union`, and self-`zip`.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2604,6 +2919,67 @@ private[spark] class DAGScheduler(
.getOrElse(new Properties())
addPySparkConfigsToProperties(stage, properties)
+ // For a pipelined PRODUCER stage, tell its tasks which of its reduce partitions the job
+ // actually reads. 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.
+ //
+ // The live set is per-SHUFFLE-EDGE, not per-job: it is the reduce partitions the consumer of
+ // THIS shuffle reads. liveReduceSet computes it by walking the narrow chain from the result
+ // RDD down to this shuffle, threading the read partition subset through each dependency's
+ // getParents. 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, and
+ // its shuffle is not narrow-reachable from the result RDD (an intervening shuffle blocks the
+ // walk), so liveReduceSet returns None and the property is left unset -- fully live -- which
+ // is correct. See the None handling below for the fail-fast case.
+ stage match {
+ case sms: ShuffleMapStage
+ if isPipelinedProducer(stage) && pipelinedManagerWantsLiveReduceHints =>
+ val resultStage = jobIdToActiveJob.get(jobId).map(_.finalStage)
+ .collect { case rs: ResultStage => rs }
+ resultStage.foreach { rs =>
+ // Tell this pipelined producer's tasks which of ITS reduce partitions the job's readers
+ // will actually drain, so the writer can drop records routed to partitions no consumer
+ // reads (a partial read -- LIMIT / executeTake -- runs only a subset of the result
+ // stage's partitions; feeding the rest fills their bounded queues and deadlocks the
+ // writer). This is the reduce-partition set the result stage's partition subset maps to
+ // through the narrow chain down to this shuffle (see liveReduceSet).
+ liveReduceSet(rs.rdd, rs.partitions.toSet, sms.shuffleDep.shuffleId) match {
Review Comment:
Addressed at both cache construction and consumption. CacheManager disables the channel rewrite when building or rebuilding a cache, and plans reading cached inputs also stay regular.
The guard looks inside AQE query stages; the regression test caught `TableCacheQueryStage` hiding the cache scan during replanning. Tests remove one cached partition, verify recomputation and repeated actions, and check that a downstream shuffle remains regular.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -1129,6 +1189,93 @@ private[spark] class DAGScheduler(
(hasPipelined, hasRegular)
}
+ /** Classify `finalRDD`'s shuffle graph; see [[JobShuffleShape]] for the shape semantics. */
+ private[scheduler] def classifyJobShuffleShape(finalRDD: RDD[_]): JobShuffleShape = {
+ // Cheap pre-pass first: which KINDS of boundary the graph has, over the shared
+ // `traverseRDDGraph` (a HashSet[RDD] visited set, no per-visit allocation). Only a job with
+ // BOTH kinds can be an unsupported mix, and only then are the two below-regular facts
+ // meaningful:
+ // - all-regular (no pipelined dep) => nothing can be pipelined-below-regular;
+ // - all-pipelined (no regular dep) => no regular boundary to be below, or to materialize.
+ // So every job that is not mixed -- which is EVERY job on a deployment that never enables the
+ // feature -- costs exactly what it costs without this feature, instead of paying for the
+ // (RDD, Boolean)-keyed two-context walk and the boundary map below.
+ val (hasPipelinedKind, hasRegularKind) = classifyJobShuffleKinds(finalRDD)
+ if (!hasPipelinedKind || !hasRegularKind) {
+ return JobShuffleShape(
+ hasPipelined = hasPipelinedKind,
+ hasUnmaterializedRegularBoundary = false,
+ hasPipelinedBelowRegular = false)
+ }
+ var hasPipelined = false
+ var pipelinedBelow = false
+ // Frontier regular shuffle boundaries: those reachable from the final RDD WITHOUT crossing
+ // another regular boundary, deduped by shuffle ID. Only these matter for the materialization
+ // check (a regular boundary below another one is never a runnable suffix member).
+ val regularBoundaries = new HashMap[Int, ShuffleDependency[_, _, _]]
+
+ // ONE walk, carrying `belowRegular` (true once the path from the final RDD has crossed a
+ // regular boundary), computes hasPipelined and pipelinedBelow together -- replacing the old
+ // per-boundary rddGraphHasPipelinedDependency re-walks (O(K x graph) on shared ancestors).
+ // A node reachable BOTH above and below a regular boundary must be explored in BOTH contexts:
+ // a pipelined dep under it counts as pipelinedBelow on the below path but not on the above
+ // path. So the visited set is keyed on (RDD, belowRegular), NOT on the RDD alone -- keying on
+ // the RDD alone would let the first-reached context win and drop the other, missing a
+ // pipelined-below-regular dep (a wrongly-accepted job). A node is thus visited at most twice,
+ // keeping the cost O(graph) rather than O(K x graph). hasPipelined is set only above a regular
+ // boundary, matching the old walk (which stopped at boundaries): a below-boundary pipelined
+ // dep is the pipelinedBelow reject case, never a runnable group member.
+ val visited = new HashSet[(RDD[_], Boolean)]
+ val stack = new ListBuffer[(RDD[_], Boolean)]
+ stack += ((finalRDD, false))
+ while (stack.nonEmpty) {
+ val entry = stack.remove(0)
+ val rdd = entry._1
+ val belowRegular = entry._2
+ if (visited.add(entry)) {
+ rdd.dependencies.foreach {
+ case pd: PipelinedShuffleDependency[_, _, _] =>
+ if (belowRegular) pipelinedBelow = true else hasPipelined = true
+ stack.prepend((pd.rdd, belowRegular))
+ case sd: ShuffleDependency[_, _, _] =>
+ // A frontier boundary only when not already below one; descend with belowRegular set.
+ if (!belowRegular) regularBoundaries.getOrElseUpdate(sd.shuffleId, sd)
+ stack.prepend((sd.rdd, true))
+ case narrowDep =>
+ stack.prepend((narrowDep.rdd, belowRegular))
+ }
+ }
+ }
+
+ // The materialization check only matters for a pipelined job: `isUnsupportedMix` consumes
+ // `hasUnmaterializedRegularBoundary` only when `hasPipelined` is true (a pipelined shuffle
+ // below an unmaterialized regular boundary is the rejected shape). A job with no pipelined
+ // dependency -- every job on a feature-off deployment -- would otherwise pay K
+ // getNumAvailableOutputs lookups (a read-locked shuffleStatuses count) for a value never read,
+ // on the single-threaded event loop. So skip the loop entirely unless the walk saw a pipelined
+ // dependency; a non-pipelined job reports hasUnmaterialized = false (unused).
+ var hasUnmaterialized = false
+ if (hasPipelined) {
+ regularBoundaries.values.foreach { sd =>
+ // Materialized means every MAP partition has a registered output: the tracker counts map
+ // outputs, so compare against the producer RDD's partition count (matching how
+ // ShuffleMapStage.isAvailable derives completeness), not the reducer-side partitioner.
+ // This is a point-in-time check at job submission. If a materialized prefix's output were
+ // LOST after this classification but before the pipelined suffix finished (executor loss),
+ // the prefix would need to re-run while the gang holds all slots -- the very deadlock this
+ // shape check forbids. That is safe here for two reasons: (1) the only supported deployment
+ // is single-executor local mode, where executor loss does not occur in normal operation;
+ // and (2) if a FetchFailed did strip the prefix, handleTaskCompletion routes it to a
+ // WHOLE-GROUP abort (the failing stage is a pipelined group member), not a lone-stage
+ // resubmit into the held slots -- the job reruns from scratch rather than deadlocking.
+ if (mapOutputTracker.getNumAvailableOutputs(sd.shuffleId) != sd.rdd.partitions.length) {
Review Comment:
You are right that “the job reruns from scratch” was not implemented here. I corrected that comment and narrowed AQE eligibility: if any shuffle would remain regular, including a materialized stage or sibling branch, the rewrite is skipped.
This gives up the materialized-prefix/pipelined-tail optimization and preserves regular-shuffle recovery; it does not add pipelined-group retry. The regression executes the same Dataset three times with shuffle-file cleanup enabled and verifies the results each time.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,148 @@
+/*
+ * 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.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.joins.CartesianProductExec
+
+/**
+ * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a
+ * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process
+ * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages
+ * together.
+ *
+ * The rewrite is deliberately unconditional rather than cost-based: the scheduler requires a job to
+ * be all-pipelined (or a materialized prefix below a pipelined suffix), so choosing per exchange
+ * would produce exactly the mixed shapes it rejects. What the rule does decide is ELIGIBILITY --
+ * the environment gates in [[PipelinedShuffleEligibility]] and the plan shapes below that force the
+ * whole plan back to regular.
+ *
+ * Enabled only when `spark.sql.shuffle.localPipelined.enabled=true`. It runs in the non-AQE
+ * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind
+ * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges).
+ *
+ * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which
+ * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected.
+ * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only
+ * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the
+ * id (SinglePartition is the numPartitions == 1 degenerate case).
+ *
+ * These shapes make the rule leave the whole plan regular:
+ * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any
+ * exchange in the plan is reused the rule bails out.
+ * - an UNSUPPORTED CONSUMER reading a shuffle (see [[readsShuffleThroughUnsupportedConsumer]]):
+ * an operator that would drain a shuffle in a way the channel transport cannot serve, or that
+ * builds its own hidden regular shuffle. If such an operator sits above any shuffle the rule
+ * leaves the WHOLE plan regular (leaving only that one exchange regular would put a pipelined
+ * exchange below a regular boundary, which the scheduler rejects -- so it is all-or-nothing).
+ * The query runs correctly, just not pipelined. The unsupported consumers are:
+ * - `CoalesceExec` (user `.coalesce(n)`): its `CoalescedRDD` makes ONE reduce task drain
+ * SEVERAL reduce partitions sequentially. The single-threaded writer parks on a full
+ * bounded queue filling a later partition before emitting an earlier one's markers, so a
+ * reader draining partitions in order deadlocks the writer with no timeout escape;
+ * `coalesce`'s narrow-merge contract also cannot be honored by re-hashing to `n`.
+ * - `CartesianProductExec`: its `UnsafeCartesianRDD` reads each left (child) partition once
+ * per right partition, so N reduce tasks mint N readers on the SAME rendezvous queue for
+ * one `(shuffleId, epoch, pid)` -- rows and end-of-stream markers split
+ * nondeterministically (wrong results), a reader short of `numMaps` markers hangs, and
+ * the first to finish abandons the queue and discards the others' data. The fan-out check
+ * does not catch it (one consumer RDD, computed many times), nor the width-1 require.
+ * - `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec`: each builds a
+ * hidden regular (`pipelined = false`) shuffle inside `doExecute` via
+ * `prepareShuffleDependency`, invisible to this plan walk. A flipped exchange below one of
+ * them would sit under that unmaterialized regular boundary and the job would hard-fail at
+ * submission (`classifyJobShuffleShape`'s pipelined-below-regular rejection). (`.collect()`
+ * on a limit takes `executeTake` and never hits `doExecute`; `.write` / `.toLocalIterator`
+ * / a non-root position do.)
+ */
+object EnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the AQE rule's -- see PipelinedShuffleEligibility for why it is a correctness
+ // gate. It also requires AQE off implicitly: under AQE this rule sees no exchanges.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ val shuffles = plan.collect { case s: ShuffleExchangeExec => s }
+ if (shuffles.isEmpty) return plan
+
+ // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so
+ // leave the whole plan regular rather than produce a rejected job. Check subquery plans
+ // too (plan.exists walks the operator tree only): today no SQL shape can place a reused
+ // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse
+ // never fires because the subquery's own preparation pass (PlanSubqueries ->
+ // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE
+ // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery
+ // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the
+ // second mechanism is an accident of rule ordering and the third is optimizer behavior,
+ // so this gate does not rely on either.
+ if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) {
+ // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output,
+ // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has
+ // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined")
+ // rather than WARN, which would fire on every reuse-bearing query and read as a fault.
+ logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " +
+ "avoid a fan-out pipelined job.")
+ return plan
+ }
+
+ // An operator that would read a shuffle in a way the channel transport cannot serve, or that
+ // builds its own hidden regular shuffle, forces the whole plan regular (see class doc). Like
+ // the reuse fallback this is a normal, expected outcome, so log at DEBUG rather than WARN.
+ if (readsShuffleThroughUnsupportedConsumer(plan)) {
+ logDebug("EnablePipelinedShuffle: a shuffle is read through an operator the channel " +
+ "transport cannot serve (coalesce / cartesian product / a limit operator that builds a " +
+ "hidden shuffle); leaving the plan regular.")
+ return plan
+ }
+
+ plan.transformUp {
Review Comment:
Addressed by the same `DeserializeToObjectExec` boundary guard as the narrow-consumer issue. The RDD exposed to user code now has regular shuffle dependencies, so adding an RDD shuffle does not create a pipelined-below-regular job.
Added AQE-on/off coverage for `reduceByKey` and `repartition`, including result checks.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala:
##########
@@ -0,0 +1,190 @@
+/*
+ * 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}
+
+import org.apache.spark.{SparkContext, TaskContext}
+
+/**
+ * Process-wide rendezvous between the map (writer) and reduce (reader) sides of an
+ * in-process pipelined shuffle. One bounded queue exists per
+ * `(shuffleId, epoch, 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.
+ *
+ * The `epoch` is the per-run id (the jobId, propagated to both the writer and the reader of
+ * one gang via a job-level local property -- see `SparkContext.SPARK_PIPELINED_RUN_EPOCH`). A
+ * shuffleId is RE-RUN within one query (a RangePartitioner sample job then the main job;
+ * executeTake's per-batch jobs; a classic Dataset re-executing a reused plan), and each run is
+ * a different job, hence a different epoch. Keying by epoch makes each run's queues and marks
+ * PHYSICALLY separate: the new run never sees a partition whose reader never started in the old
+ * run (whose queue still holds stale batches + end-of-stream markers), and a straggler writer
+ * from an aborted run -- still looping because a task kill need not interrupt the thread -- can
+ * only touch its OWN (old) epoch's queue, never the new run's. This replaces an earlier design
+ * that reused one key per shuffleId and tried to reset shared marks between runs, which left
+ * stale queues and raced stragglers.
+ *
+ * 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
+
+ /**
+ * The per-run epoch for the current task, read from the job-level local property the
+ * DAGScheduler stamps on a pipelined job (`SparkContext.SPARK_PIPELINED_RUN_EPOCH` = jobId).
+ * Both the writer and the reader of one gang read this, so they address the same per-run
+ * queues. Absent (a core-RDD path that never sets it) defaults to 0; that is fine because such
+ * a shuffleId is never re-run into a colliding second live run.
+ */
+ def epochOf(tc: TaskContext): Int =
+ Option(tc)
+ .flatMap(t => Option(t.getLocalProperty(SparkContext.SPARK_PIPELINED_RUN_EPOCH)))
+ .map(_.toInt)
+ .getOrElse(0)
+
+ // State is nested by shuffleId FIRST, then keyed by (epoch, reducePartitionId) within it. The
+ // outer level exists so the two per-shuffle operations the ContextCleaner drives -- holdsShuffle
+ // and removeShuffle -- are O(1) map lookups instead of a scan of every live entry: holdsShuffle
+ // now runs for EVERY shuffle cleaned in a feature-on session, including the regular prefix
+ // shuffles this feature itself produces, so a flat key made that O(live entries) per cleanup.
+ //
+ // Queue values are AnyRef because a queue carries both record batches (Array[AnyRef]) and the
+ // EndOfStream marker.
+ private val queues =
+ new ConcurrentHashMap[Int, ConcurrentHashMap[(Int, Int), LinkedBlockingQueue[AnyRef]]]()
+
+ // (epoch, reducePartitionId) keys, per shuffleId, 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 =
+ new ConcurrentHashMap[Int, java.util.Set[(Int, Int)]]()
+
+ /**
+ * Per-queue capacity in BATCHES (not rows), the backpressure bound and the heap-residency
+ * knob (see spark.shuffle.channel.queueCapacity). Set once by the channel manager at
+ * construction from that conf; defaults to 64 (with the default 1024-row batch, ~64K rows per
+ * reduce partition in flight) until a manager sets it. `@volatile` because the manager sets it
+ * on the driver while writer/reader threads read it.
+ */
+ @volatile private var capacity = 64
+
+ /** Set the per-queue capacity in batches. Called by the channel manager from its conf. */
+ private[pipelined] def setCapacity(batches: Int): Unit = { capacity = batches }
+
+ /** The queue for one `(shuffleId, epoch, reducePartitionId)`, created on first access. */
+ def queue(shuffleId: Int, epoch: Int, reducePartitionId: Int): LinkedBlockingQueue[AnyRef] = {
Review Comment:
Added manager `startRun`/`endRun` hooks tied to the job lifecycle. Job cleanup releases that epoch’s queues and abandon marks, including on failure or cancellation. Late task accesses cannot recreate state for an ended run.
Unregister also preserves active epochs until their owning jobs finish. This matters when a concurrent action fails and its SQL cleanup runs while the first action is still using the shuffle.
Added tests for repeated actions with a reachable dependency, epoch isolation during cleanup, late task access, failure cleanup, and unregister during an active run.
##########
core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala:
##########
@@ -2129,8 +2428,24 @@ private[spark] class DAGScheduler(
// Job submitted, clear internal data.
barrierJobIdToNumTasksCheckFailures.remove(jobId)
+ // For a pipelined job, stamp the per-run epoch (the jobId) into the job's properties BEFORE
+ // creating the ActiveJob, so submitMissingTasks -- which clones jobIdToActiveJob(jobId)
+ // .properties per stage -- carries the SAME epoch to every stage's tasks. Both the producer
+ // (writer) and the consumer (reader) of the one gang belong to this job, so both read one
+ // value; a different run is a different job, hence a different epoch, which keys the
+ // in-process channel rendezvous per run (see SPARK_PIPELINED_RUN_EPOCH). Copy the caller's
+ // Properties rather than mutating it. Inert for a non-pipelined job (property never set,
+ // never read).
+ val jobProperties =
+ if (hasPipelined && pipelinedManagerWantsLiveReduceHints) {
+ val p = Utils.cloneProperties(if (properties == null) new Properties() else properties)
+ p.setProperty(SparkContext.SPARK_PIPELINED_RUN_EPOCH, jobId.toString)
Review Comment:
The cross-job reuse restriction remains, and it is now documented. Added a non-AQE SQL test that verifies the second action is rejected, its cleanup does not disrupt the owning action, and a subsequent sequential action succeeds. The test runs with shuffle-file cleanup enabled.
One qualification to the AQE example: execution of the same adaptive plan is synchronized through `withFinalPlanUpdate`, so that case can serialize rather than reach the cross-job rejection. I have not added multi-job sharing of the channel.
##########
core/src/main/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleManager.scala:
##########
@@ -0,0 +1,154 @@
+/*
+ * 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)
+
+ // Per-queue depth in batches (backpressure bound + heap-residency knob). Set the process-wide
+ // rendezvous from the conf at construction, before any writer/reader creates a queue.
+ ChannelShuffleRendezvous.setCapacity(conf.get(config.SHUFFLE_PIPELINED_CHANNEL_QUEUE_CAPACITY))
+
+ override def usesStreamingShuffleOutputTracker: Boolean = false
Review Comment:
Agreed. The live-reduce hint limits emitted data, not producer computation. The config documentation now states the per-job recomputation cost, including once per output partition for `toLocalIterator`, and recommends caching first.
Added AQE-on/off tests using an accumulator to verify producer evaluation counts across repeated iterator actions, alongside result and epoch-cleanup checks. The recomputation behavior remains; this change documents and tests it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,229 @@
+/*
+ * 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.adaptive
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{BinaryExecNode, CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.exchange.{PipelinedShuffleEligibility, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.ShuffledJoin
+
+/**
+ * Opt-in (SPARK-57399). Flips eligible [[ShuffleExchangeExec]]
+ * nodes to `pipelined = true` under AQE, the adaptive counterpart of the non-AQE
+ * `EnablePipelinedShuffle` preparation rule (which is a no-op once the plan is wrapped in
+ * `AdaptiveSparkPlanExec`). Runs in `AdaptiveSparkPlanExec.queryStagePreparationRules`, so
+ * it is re-applied on every replanning round; the decision is deterministic on plan shape,
+ * and already-flipped exchanges are left alone.
+ *
+ * Placement policy. A flipped exchange has no
+ * map output statistics (it never materializes as a query stage: the DAGScheduler
+ * gang-runs it inline with its consumer in the final job -- see the pipelined case in
+ * `AdaptiveSparkPlanExec.createNonResultQueryStages`), so only exchanges whose statistics
+ * no AQE decision consumes are flipped:
+ *
+ * - "free" candidate: the path from the candidate to the plan root crosses no
+ * stats-sensitive node ([[BinaryExecNode]], another [[ShuffleExchangeExec]], or a
+ * query stage). Its own coalescing/skew handling is given up; nothing above needed its
+ * stats.
+ * - "join-paired" candidate: the immediate shuffle inputs of a [[ShuffledJoin]] whose
+ * path to the root is otherwise free, flipped only as a symmetric pair (an asymmetric
+ * flip would leave one side participating in AQE coalesce/skew and the other fixed).
+ * - everything else stays regular and materializes as usual -- those stages form the
+ * fully-materialized prefix the scheduler's mixed-job shape requires.
+ *
+ * A pipelined exchange supports every
+ * partitioning, so a SinglePartition exchange in a free position is simply a candidate
+ * itself. The walk stops below a flipped candidate: exchanges underneath stay regular and
+ * keep full AQE treatment. Candidates whose canonicalized form occurs more than once in
+ * the plan (including inside materialized stages and subqueries) are skipped: flipping
+ * them would trade AQE's stage reuse for duplicate recomputation, and a pipelined producer
+ * cannot be consumed twice.
+ */
+object AQEEnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the non-AQE rule's -- see PipelinedShuffleEligibility for why it is a
+ // correctness gate that must not drift between the two rules.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ flipEligibleExchanges(plan)
+ }
+
+ /**
+ * The plan-shape core of the rule, factored out of [[apply]]'s environment guards (opt-in flag,
+ * local mode, channel manager) so it can be unit-tested on a hand-built plan directly. Collects
+ * the eligible exchanges and returns the plan with each flipped to `pipelined = true`.
+ */
+ private[adaptive] def flipEligibleExchanges(plan: SparkPlan): SparkPlan = {
+ val shared = if (conf.exchangeReuseEnabled) duplicatedShuffleForms(plan) else Set.empty[Any]
+ // Collect the exchanges to flip BY IDENTITY (SparkPlan.id, unique per instance), not by the
+ // node itself: TreeNode overrides hashCode but not equals, so a HashSet[ShuffleExchangeExec]
+ // matches structurally, and the transformDown below would then flip EVERY exchange
+ // structurally equal to a collected one -- including a twin the collector deliberately left
+ // regular on a blocked path. That twin, if it sits below a regular boundary, makes
+ // classifyJobShuffleShape reject the whole job. Keying on the instance id flips exactly the
+ // nodes the collector chose, regardless of spark.sql.exchange.reuse (duplicatedShuffleForms,
+ // the only other guard, is empty when reuse is off). transformDown matches each ORIGINAL node
+ // before rebuilding it, so its id is the same instance id the collector recorded.
+ val toFlip = mutable.HashSet.empty[Int]
+ collectCandidates(plan, blocked = false, shared, toFlip)
+ if (toFlip.isEmpty) return plan
+
+ // transformDown, NOT transformUp: candidates can be nested (a SinglePartition candidate
+ // above a hash candidate). transformUp rebuilds children first, so by the time it
+ // reaches the upper candidate that node is a NEW instance whose (already flipped) child
+ // no longer matches the collected original structurally, and the upper flip is silently
+ // dropped -- leaving a regular exchange above a pipelined one, which the scheduler then
+ // rejects. transformDown hands each candidate to the pattern before its subtree is
+ // rebuilt, so both nested flips apply.
+ plan.transformDown {
+ case s: ShuffleExchangeExec if toFlip.contains(s.id) => s.copy(pipelined = true)
+ }
+ }
+
+ private def isCandidate(s: ShuffleExchangeExec, shared: Set[Any]): Boolean =
Review Comment:
Retained the explicit admission failure, but expanded the config documentation to explain that the default 200 partitions is usually too wide for local execution, that AQE does not coalesce pipelined exchanges, and that all producer and consumer tasks must fit concurrently.
The setting guidance is therefore based on total group demand, not just requiring `spark.sql.shuffle.partitions` to be below the core count. The existing admission-failure tests remain.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala:
##########
@@ -570,7 +569,15 @@ 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. And
+ // skip it when rddWithPartitionIds already copied: needToCopyObjectsBeforeShuffle makes
+ // that RDD emit (pid, row.copy()), so a second copy here would be redundant.
+ shuffleWriterProcessor = createShuffleWriteProcessor(
+ writeMetrics,
+ copyRows = SparkEnv.get.pipelinedShuffleManager.requiresDetachedRecords &&
Review Comment:
Applied this suggestion. `needToCopyObjectsBeforeShuffle` now takes the pipelined flag and consults `pipelinedShuffleManager.requiresDetachedRecords` for that path.
Rows are copied at the existing shuffle-input construction site. Removed `copyRows`, the write-processor override, and the incorrect comment, so the channel’s copy decision no longer depends on the blocking manager’s thresholds.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58097:
URL: https://github.com/apache/spark/pull/58097#discussion_r3975215351
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala:
##########
@@ -0,0 +1,229 @@
+/*
+ * 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.adaptive
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{BinaryExecNode, CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec}
+import org.apache.spark.sql.execution.exchange.{PipelinedShuffleEligibility, ReusedExchangeExec, ShuffleExchangeExec}
+import org.apache.spark.sql.execution.joins.ShuffledJoin
+
+/**
+ * Opt-in (SPARK-57399). Flips eligible [[ShuffleExchangeExec]]
+ * nodes to `pipelined = true` under AQE, the adaptive counterpart of the non-AQE
+ * `EnablePipelinedShuffle` preparation rule (which is a no-op once the plan is wrapped in
+ * `AdaptiveSparkPlanExec`). Runs in `AdaptiveSparkPlanExec.queryStagePreparationRules`, so
+ * it is re-applied on every replanning round; the decision is deterministic on plan shape,
+ * and already-flipped exchanges are left alone.
+ *
+ * Placement policy. A flipped exchange has no
+ * map output statistics (it never materializes as a query stage: the DAGScheduler
+ * gang-runs it inline with its consumer in the final job -- see the pipelined case in
+ * `AdaptiveSparkPlanExec.createNonResultQueryStages`), so only exchanges whose statistics
+ * no AQE decision consumes are flipped:
+ *
+ * - "free" candidate: the path from the candidate to the plan root crosses no
+ * stats-sensitive node ([[BinaryExecNode]], another [[ShuffleExchangeExec]], or a
+ * query stage). Its own coalescing/skew handling is given up; nothing above needed its
+ * stats.
+ * - "join-paired" candidate: the immediate shuffle inputs of a [[ShuffledJoin]] whose
+ * path to the root is otherwise free, flipped only as a symmetric pair (an asymmetric
+ * flip would leave one side participating in AQE coalesce/skew and the other fixed).
+ * - everything else stays regular and materializes as usual -- those stages form the
+ * fully-materialized prefix the scheduler's mixed-job shape requires.
+ *
+ * A pipelined exchange supports every
+ * partitioning, so a SinglePartition exchange in a free position is simply a candidate
+ * itself. The walk stops below a flipped candidate: exchanges underneath stay regular and
+ * keep full AQE treatment. Candidates whose canonicalized form occurs more than once in
+ * the plan (including inside materialized stages and subqueries) are skipped: flipping
+ * them would trade AQE's stage reuse for duplicate recomputation, and a pipelined producer
+ * cannot be consumed twice.
+ */
+object AQEEnablePipelinedShuffle extends Rule[SparkPlan] {
+
+ override def apply(plan: SparkPlan): SparkPlan = {
+ // Shared environment gate (opt-in flag, single-executor local mode, channel manager active),
+ // identical to the non-AQE rule's -- see PipelinedShuffleEligibility for why it is a
+ // correctness gate that must not drift between the two rules.
+ if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan
+
+ flipEligibleExchanges(plan)
+ }
+
+ /**
+ * The plan-shape core of the rule, factored out of [[apply]]'s environment guards (opt-in flag,
+ * local mode, channel manager) so it can be unit-tested on a hand-built plan directly. Collects
+ * the eligible exchanges and returns the plan with each flipped to `pipelined = true`.
+ */
+ private[adaptive] def flipEligibleExchanges(plan: SparkPlan): SparkPlan = {
+ val shared = if (conf.exchangeReuseEnabled) duplicatedShuffleForms(plan) else Set.empty[Any]
+ // Collect the exchanges to flip BY IDENTITY (SparkPlan.id, unique per instance), not by the
+ // node itself: TreeNode overrides hashCode but not equals, so a HashSet[ShuffleExchangeExec]
+ // matches structurally, and the transformDown below would then flip EVERY exchange
+ // structurally equal to a collected one -- including a twin the collector deliberately left
+ // regular on a blocked path. That twin, if it sits below a regular boundary, makes
+ // classifyJobShuffleShape reject the whole job. Keying on the instance id flips exactly the
+ // nodes the collector chose, regardless of spark.sql.exchange.reuse (duplicatedShuffleForms,
+ // the only other guard, is empty when reuse is off). transformDown matches each ORIGINAL node
+ // before rebuilding it, so its id is the same instance id the collector recorded.
+ val toFlip = mutable.HashSet.empty[Int]
+ collectCandidates(plan, blocked = false, shared, toFlip)
+ if (toFlip.isEmpty) return plan
+
+ // transformDown, NOT transformUp: candidates can be nested (a SinglePartition candidate
+ // above a hash candidate). transformUp rebuilds children first, so by the time it
+ // reaches the upper candidate that node is a NEW instance whose (already flipped) child
+ // no longer matches the collected original structurally, and the upper flip is silently
+ // dropped -- leaving a regular exchange above a pipelined one, which the scheduler then
+ // rejects. transformDown hands each candidate to the pattern before its subtree is
+ // rebuilt, so both nested flips apply.
+ plan.transformDown {
+ case s: ShuffleExchangeExec if toFlip.contains(s.id) => s.copy(pipelined = true)
+ }
+ }
+
+ private def isCandidate(s: ShuffleExchangeExec, shared: Set[Any]): Boolean =
+ !s.pipelined && !shared.contains(s.canonicalized)
+
+ /**
+ * Top-down walk collecting exchanges to flip. `blocked` is true once the path from the
+ * root has crossed a stats-sensitive node.
+ */
+ private def collectCandidates(
+ plan: SparkPlan,
+ blocked: Boolean,
+ shared: Set[Any],
+ out: mutable.HashSet[Int]): Unit = plan match {
+ case s: ShuffleExchangeExec =>
+ val flipped = !blocked && isCandidate(s, shared)
+ if (flipped) {
+ out += s.id
+ }
+ // A flipped SinglePartition exchange keeps the walk going: AQE makes no decision at
+ // it (it cannot be coalesced or skew-split), so free candidates BELOW it flip too,
+ // forming a pipelined chain: all exchanges in such a chain must flip together (a
+ // SinglePartition exchange cannot be coalesced or skew-split, so AQE makes no decision
+ // at it and free candidates below it flip too). Below any OTHER exchange (flipped or
+ // not) the walk stops: what is underneath either materializes as the prefix or feeds
+ // a regular exchange whose stats AQE uses, and keeps full AQE treatment either way.
+ if (flipped &&
+ s.outputPartitioning == org.apache.spark.sql.catalyst.plans.physical.SinglePartition) {
+ collectCandidates(s.child, blocked = false, shared, out)
+ }
+
+ case _: QueryStageExec => // already materialized; a leaf here
+
+ case j: ShuffledJoin if !blocked =>
+ // Flip the join's immediate shuffle inputs only as a symmetric pair.
+ val leftCandidate = immediateShuffleInput(j.left, shared)
+ val rightCandidate = immediateShuffleInput(j.right, shared)
+ (leftCandidate, rightCandidate) match {
+ case (Some(l), Some(r)) =>
+ out += l.id
+ out += r.id
+ case _ => // asymmetric (a broadcast side, a materialized stage, no clean input): skip
+ }
+ // Anything deeper is below a join input; blocked either way.
+
+ case p if isStatsSensitive(p) =>
+ p.children.foreach(collectCandidates(_, blocked = true, shared, out))
+
+ case p =>
+ p.children.foreach(collectCandidates(_, blocked, shared, out))
+ }
+
+ /**
+ * The single eligible [[ShuffleExchangeExec]] at the top of one join input, looking
+ * through unary non-stats-sensitive forwarders. None if the input is anything else.
+ */
+ private def immediateShuffleInput(
+ plan: SparkPlan,
+ shared: Set[Any]): Option[ShuffleExchangeExec] = plan match {
+ case s: ShuffleExchangeExec => Some(s).filter(isCandidate(_, shared))
+ case _: QueryStageExec => None
+ case p if isStatsSensitive(p) => None
+ case p if p.children.size == 1 => immediateShuffleInput(p.children.head, shared)
+ case _ => None
+ }
+
+ /**
+ * Nodes below which a candidate exchange must NOT be flipped. Reasons a node lands here:
+ * - AQE consumes map output statistics from the stages below it ([[BinaryExecNode]], another
+ * [[ShuffleExchangeExec]]); flipping below it would give up stats a decision needs.
+ * - [[CoalesceExec]] reads its child shuffle multi-partition-per-task (a `CoalescedRDD` over
+ * the `ShuffledRowRDD`), which the channel transport cannot serve; the shuffle it reads
+ * must stay regular.
+ * - the limit operators [[CollectLimitExec]] / [[CollectTailExec]] /
+ * [[TakeOrderedAndProjectExec]] each build a hidden regular (`pipelined = false`) shuffle
+ * inside `doExecute` (via `prepareShuffleDependency`) that no plan walk can see; a flipped
+ * exchange below one of them would sit under that unmaterialized regular boundary and the
+ * job would hard-fail at submission (pipelined-below-regular).
+ * Blocking here keeps the shuffle below -- and everything deeper -- regular, so no pipelined
+ * exchange ends up below the operator's regular boundary. See the non-AQE
+ * `EnablePipelinedShuffle` for the full rationale on each operator.
+ */
+ private def isStatsSensitive(plan: SparkPlan): Boolean = plan match {
Review Comment:
Moved the operator list into `PipelinedShuffleEligibility.isUnsupportedConsumer`. The non-AQE rule uses it directly, and the AQE rule combines it with its stats-sensitive checks.
The RDD/cache boundary checks also live in the shared eligibility gate so the two rules apply the same exclusions.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org