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: [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]
