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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
