viirya commented on code in PR #58097: URL: https://github.com/apache/spark/pull/58097#discussion_r3975214459
########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala: ########## @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.joins.CartesianProductExec + +/** + * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a + * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process + * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages + * together. + * + * The rewrite is deliberately unconditional rather than cost-based: the scheduler requires a job to + * be all-pipelined (or a materialized prefix below a pipelined suffix), so choosing per exchange + * would produce exactly the mixed shapes it rejects. What the rule does decide is ELIGIBILITY -- + * the environment gates in [[PipelinedShuffleEligibility]] and the plan shapes below that force the + * whole plan back to regular. + * + * Enabled only when `spark.sql.shuffle.localPipelined.enabled=true`. It runs in the non-AQE + * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind + * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges). + * + * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which + * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected. + * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only + * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the + * id (SinglePartition is the numPartitions == 1 degenerate case). + * + * These shapes make the rule leave the whole plan regular: + * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any + * exchange in the plan is reused the rule bails out. + * - an UNSUPPORTED CONSUMER reading a shuffle (see [[readsShuffleThroughUnsupportedConsumer]]): + * an operator that would drain a shuffle in a way the channel transport cannot serve, or that + * builds its own hidden regular shuffle. If such an operator sits above any shuffle the rule + * leaves the WHOLE plan regular (leaving only that one exchange regular would put a pipelined + * exchange below a regular boundary, which the scheduler rejects -- so it is all-or-nothing). + * The query runs correctly, just not pipelined. The unsupported consumers are: + * - `CoalesceExec` (user `.coalesce(n)`): its `CoalescedRDD` makes ONE reduce task drain + * SEVERAL reduce partitions sequentially. The single-threaded writer parks on a full + * bounded queue filling a later partition before emitting an earlier one's markers, so a + * reader draining partitions in order deadlocks the writer with no timeout escape; + * `coalesce`'s narrow-merge contract also cannot be honored by re-hashing to `n`. + * - `CartesianProductExec`: its `UnsafeCartesianRDD` reads each left (child) partition once + * per right partition, so N reduce tasks mint N readers on the SAME rendezvous queue for + * one `(shuffleId, epoch, pid)` -- rows and end-of-stream markers split + * nondeterministically (wrong results), a reader short of `numMaps` markers hangs, and + * the first to finish abandons the queue and discards the others' data. The fan-out check + * does not catch it (one consumer RDD, computed many times), nor the width-1 require. + * - `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec`: each builds a + * hidden regular (`pipelined = false`) shuffle inside `doExecute` via + * `prepareShuffleDependency`, invisible to this plan walk. A flipped exchange below one of + * them would sit under that unmaterialized regular boundary and the job would hard-fail at + * submission (`classifyJobShuffleShape`'s pipelined-below-regular rejection). (`.collect()` + * on a limit takes `executeTake` and never hits `doExecute`; `.write` / `.toLocalIterator` + * / a non-root position do.) + */ +object EnablePipelinedShuffle extends Rule[SparkPlan] { + + override def apply(plan: SparkPlan): SparkPlan = { + // Shared environment gate (opt-in flag, single-executor local mode, channel manager active), + // identical to the AQE rule's -- see PipelinedShuffleEligibility for why it is a correctness + // gate. It also requires AQE off implicitly: under AQE this rule sees no exchanges. + if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan + + val shuffles = plan.collect { case s: ShuffleExchangeExec => s } + if (shuffles.isEmpty) return plan + + // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so + // leave the whole plan regular rather than produce a rejected job. Check subquery plans + // too (plan.exists walks the operator tree only): today no SQL shape can place a reused + // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse + // never fires because the subquery's own preparation pass (PlanSubqueries -> + // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE + // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery + // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the + // second mechanism is an accident of rule ordering and the third is optimizer behavior, + // so this gate does not rely on either. + if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) { + // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output, + // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has + // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined") + // rather than WARN, which would fire on every reuse-bearing query and read as a fault. + logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " + + "avoid a fan-out pipelined job.") + return plan + } + + // An operator that would read a shuffle in a way the channel transport cannot serve, or that + // builds its own hidden regular shuffle, forces the whole plan regular (see class doc). Like + // the reuse fallback this is a normal, expected outcome, so log at DEBUG rather than WARN. + if (readsShuffleThroughUnsupportedConsumer(plan)) { + logDebug("EnablePipelinedShuffle: a shuffle is read through an operator the channel " + + "transport cannot serve (coalesce / cartesian product / a limit operator that builds a " + + "hidden shuffle); leaving the plan regular.") + return plan + } + + plan.transformUp { + case s: ShuffleExchangeExec if !s.pipelined => s.copy(pipelined = true) + } + } + + /** + * True if any [[ShuffleExchangeExec]] in `plan` is read by an operator the channel transport + * cannot serve. The unsupported operators (see class doc for why each is fatal) are + * `CoalesceExec`, `CartesianProductExec`, and the limit operators `CollectLimitExec` / + * `CollectTailExec` / `TakeOrderedAndProjectExec`. For each such operator anywhere in the plan, + * check whether a shuffle is reachable below it. + * + * The reachability walk descends through EVERY child of a non-exchange node -- not only unary + * children -- so a shuffle behind a `UnionExec`/join (a `BinaryExecNode`) beneath the operator + * is still found. It stops at the FIRST [[ShuffleExchangeExec]] on each path: a shuffle deeper + * than that first one is not read by this operator (the intervening exchange's own reader reads + * one reduce partition per task), so it is not this operator's concern. + */ + private def readsShuffleThroughUnsupportedConsumer(plan: SparkPlan): Boolean = { Review Comment: Addressed by keeping plans containing `DeserializeToObjectExec` regular in both SQL rules. This covers the `Dataset.rdd` boundary, after which SQL planning cannot see or constrain the RDD consumers. The guard is deliberately conservative: it can also exclude other typed Dataset operations containing that node. Added AQE-on/off tests that assert the RDD lineage has no pipelined dependency before exercising `coalesce`, self-`union`, and self-`zip`. ########## core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala: ########## @@ -2604,6 +2919,67 @@ private[spark] class DAGScheduler( .getOrElse(new Properties()) addPySparkConfigsToProperties(stage, properties) + // For a pipelined PRODUCER stage, tell its tasks which of its reduce partitions the job + // actually reads. The in-process channel writer drops records routed to partitions no + // consumer will drain -- otherwise a partial-read job (LIMIT / executeTake reads a subset) + // fills the unread partitions' bounded queues and deadlocks the writer. The result stage + // is created before submitStage, so its partitions are known here. + // + // The live set is per-SHUFFLE-EDGE, not per-job: it is the reduce partitions the consumer of + // THIS shuffle reads. liveReduceSet computes it by walking the narrow chain from the result + // RDD down to this shuffle, threading the read partition subset through each dependency's + // getParents. A MIDDLE pipelined exchange in a chain (e.g. a subquery's hash below a + // single-partition agg) is consumed by another map stage that reads ALL its partitions, and + // its shuffle is not narrow-reachable from the result RDD (an intervening shuffle blocks the + // walk), so liveReduceSet returns None and the property is left unset -- fully live -- which + // is correct. See the None handling below for the fail-fast case. + stage match { + case sms: ShuffleMapStage + if isPipelinedProducer(stage) && pipelinedManagerWantsLiveReduceHints => + val resultStage = jobIdToActiveJob.get(jobId).map(_.finalStage) + .collect { case rs: ResultStage => rs } + resultStage.foreach { rs => + // Tell this pipelined producer's tasks which of ITS reduce partitions the job's readers + // will actually drain, so the writer can drop records routed to partitions no consumer + // reads (a partial read -- LIMIT / executeTake -- runs only a subset of the result + // stage's partitions; feeding the rest fills their bounded queues and deadlocks the + // writer). This is the reduce-partition set the result stage's partition subset maps to + // through the narrow chain down to this shuffle (see liveReduceSet). + liveReduceSet(rs.rdd, rs.partitions.toSet, sms.shuffleDep.shuffleId) match { Review Comment: Addressed at both cache construction and consumption. CacheManager disables the channel rewrite when building or rebuilding a cache, and plans reading cached inputs also stay regular. The guard looks inside AQE query stages; the regression test caught `TableCacheQueryStage` hiding the cache scan during replanning. Tests remove one cached partition, verify recomputation and repeated actions, and check that a downstream shuffle remains regular. ########## core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala: ########## @@ -1129,6 +1189,93 @@ private[spark] class DAGScheduler( (hasPipelined, hasRegular) } + /** Classify `finalRDD`'s shuffle graph; see [[JobShuffleShape]] for the shape semantics. */ + private[scheduler] def classifyJobShuffleShape(finalRDD: RDD[_]): JobShuffleShape = { + // Cheap pre-pass first: which KINDS of boundary the graph has, over the shared + // `traverseRDDGraph` (a HashSet[RDD] visited set, no per-visit allocation). Only a job with + // BOTH kinds can be an unsupported mix, and only then are the two below-regular facts + // meaningful: + // - all-regular (no pipelined dep) => nothing can be pipelined-below-regular; + // - all-pipelined (no regular dep) => no regular boundary to be below, or to materialize. + // So every job that is not mixed -- which is EVERY job on a deployment that never enables the + // feature -- costs exactly what it costs without this feature, instead of paying for the + // (RDD, Boolean)-keyed two-context walk and the boundary map below. + val (hasPipelinedKind, hasRegularKind) = classifyJobShuffleKinds(finalRDD) + if (!hasPipelinedKind || !hasRegularKind) { + return JobShuffleShape( + hasPipelined = hasPipelinedKind, + hasUnmaterializedRegularBoundary = false, + hasPipelinedBelowRegular = false) + } + var hasPipelined = false + var pipelinedBelow = false + // Frontier regular shuffle boundaries: those reachable from the final RDD WITHOUT crossing + // another regular boundary, deduped by shuffle ID. Only these matter for the materialization + // check (a regular boundary below another one is never a runnable suffix member). + val regularBoundaries = new HashMap[Int, ShuffleDependency[_, _, _]] + + // ONE walk, carrying `belowRegular` (true once the path from the final RDD has crossed a + // regular boundary), computes hasPipelined and pipelinedBelow together -- replacing the old + // per-boundary rddGraphHasPipelinedDependency re-walks (O(K x graph) on shared ancestors). + // A node reachable BOTH above and below a regular boundary must be explored in BOTH contexts: + // a pipelined dep under it counts as pipelinedBelow on the below path but not on the above + // path. So the visited set is keyed on (RDD, belowRegular), NOT on the RDD alone -- keying on + // the RDD alone would let the first-reached context win and drop the other, missing a + // pipelined-below-regular dep (a wrongly-accepted job). A node is thus visited at most twice, + // keeping the cost O(graph) rather than O(K x graph). hasPipelined is set only above a regular + // boundary, matching the old walk (which stopped at boundaries): a below-boundary pipelined + // dep is the pipelinedBelow reject case, never a runnable group member. + val visited = new HashSet[(RDD[_], Boolean)] + val stack = new ListBuffer[(RDD[_], Boolean)] + stack += ((finalRDD, false)) + while (stack.nonEmpty) { + val entry = stack.remove(0) + val rdd = entry._1 + val belowRegular = entry._2 + if (visited.add(entry)) { + rdd.dependencies.foreach { + case pd: PipelinedShuffleDependency[_, _, _] => + if (belowRegular) pipelinedBelow = true else hasPipelined = true + stack.prepend((pd.rdd, belowRegular)) + case sd: ShuffleDependency[_, _, _] => + // A frontier boundary only when not already below one; descend with belowRegular set. + if (!belowRegular) regularBoundaries.getOrElseUpdate(sd.shuffleId, sd) + stack.prepend((sd.rdd, true)) + case narrowDep => + stack.prepend((narrowDep.rdd, belowRegular)) + } + } + } + + // The materialization check only matters for a pipelined job: `isUnsupportedMix` consumes + // `hasUnmaterializedRegularBoundary` only when `hasPipelined` is true (a pipelined shuffle + // below an unmaterialized regular boundary is the rejected shape). A job with no pipelined + // dependency -- every job on a feature-off deployment -- would otherwise pay K + // getNumAvailableOutputs lookups (a read-locked shuffleStatuses count) for a value never read, + // on the single-threaded event loop. So skip the loop entirely unless the walk saw a pipelined + // dependency; a non-pipelined job reports hasUnmaterialized = false (unused). + var hasUnmaterialized = false + if (hasPipelined) { + regularBoundaries.values.foreach { sd => + // Materialized means every MAP partition has a registered output: the tracker counts map + // outputs, so compare against the producer RDD's partition count (matching how + // ShuffleMapStage.isAvailable derives completeness), not the reducer-side partitioner. + // This is a point-in-time check at job submission. If a materialized prefix's output were + // LOST after this classification but before the pipelined suffix finished (executor loss), + // the prefix would need to re-run while the gang holds all slots -- the very deadlock this + // shape check forbids. That is safe here for two reasons: (1) the only supported deployment + // is single-executor local mode, where executor loss does not occur in normal operation; + // and (2) if a FetchFailed did strip the prefix, handleTaskCompletion routes it to a + // WHOLE-GROUP abort (the failing stage is a pipelined group member), not a lone-stage + // resubmit into the held slots -- the job reruns from scratch rather than deadlocking. + if (mapOutputTracker.getNumAvailableOutputs(sd.shuffleId) != sd.rdd.partitions.length) { Review Comment: You are right that “the job reruns from scratch” was not implemented here. I corrected that comment and narrowed AQE eligibility: if any shuffle would remain regular, including a materialized stage or sibling branch, the rewrite is skipped. This gives up the materialized-prefix/pipelined-tail optimization and preserves regular-shuffle recovery; it does not add pipelined-group retry. The regression executes the same Dataset three times with shuffle-file cleanup enabled and verifies the results each time. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnablePipelinedShuffle.scala: ########## @@ -0,0 +1,148 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.joins.CartesianProductExec + +/** + * Opt-in (SPARK-57399). Rewrites EVERY [[ShuffleExchangeExec]] in a + * batch physical plan to `pipelined = true`, so each shuffle is served by the in-process + * pipelined channel manager and the concurrent-stage scheduler runs the map and reduce stages + * together. + * + * The rewrite is deliberately unconditional rather than cost-based: the scheduler requires a job to + * be all-pipelined (or a materialized prefix below a pipelined suffix), so choosing per exchange + * would produce exactly the mixed shapes it rejects. What the rule does decide is ELIGIBILITY -- + * the environment gates in [[PipelinedShuffleEligibility]] and the plan shapes below that force the + * whole plan back to regular. + * + * Enabled only when `spark.sql.shuffle.localPipelined.enabled=true`. It runs in the non-AQE + * `preparations` list, so it also requires AQE to be off (under AQE the plan is hidden behind + * an opaque `AdaptiveSparkPlanExec` leaf and this rule sees no exchanges). + * + * Rewriting ALL shuffles (not just hash-partitioning ones) keeps the job all-pipelined, which + * the DAGScheduler requires: a mix of pipelined and regular shuffles in one job is rejected. + * SinglePartition and RangePartitioning exchanges pipeline fine -- the channel transport only + * routes by `partitioner.getPartition(key)` and does not care which partitioning produced the + * id (SinglePartition is the numPartitions == 1 degenerate case). + * + * These shapes make the rule leave the whole plan regular: + * - reuse: a pipelined producer with more than one consumer (fan-out) is rejected, so if any + * exchange in the plan is reused the rule bails out. + * - an UNSUPPORTED CONSUMER reading a shuffle (see [[readsShuffleThroughUnsupportedConsumer]]): + * an operator that would drain a shuffle in a way the channel transport cannot serve, or that + * builds its own hidden regular shuffle. If such an operator sits above any shuffle the rule + * leaves the WHOLE plan regular (leaving only that one exchange regular would put a pipelined + * exchange below a regular boundary, which the scheduler rejects -- so it is all-or-nothing). + * The query runs correctly, just not pipelined. The unsupported consumers are: + * - `CoalesceExec` (user `.coalesce(n)`): its `CoalescedRDD` makes ONE reduce task drain + * SEVERAL reduce partitions sequentially. The single-threaded writer parks on a full + * bounded queue filling a later partition before emitting an earlier one's markers, so a + * reader draining partitions in order deadlocks the writer with no timeout escape; + * `coalesce`'s narrow-merge contract also cannot be honored by re-hashing to `n`. + * - `CartesianProductExec`: its `UnsafeCartesianRDD` reads each left (child) partition once + * per right partition, so N reduce tasks mint N readers on the SAME rendezvous queue for + * one `(shuffleId, epoch, pid)` -- rows and end-of-stream markers split + * nondeterministically (wrong results), a reader short of `numMaps` markers hangs, and + * the first to finish abandons the queue and discards the others' data. The fan-out check + * does not catch it (one consumer RDD, computed many times), nor the width-1 require. + * - `CollectLimitExec` / `CollectTailExec` / `TakeOrderedAndProjectExec`: each builds a + * hidden regular (`pipelined = false`) shuffle inside `doExecute` via + * `prepareShuffleDependency`, invisible to this plan walk. A flipped exchange below one of + * them would sit under that unmaterialized regular boundary and the job would hard-fail at + * submission (`classifyJobShuffleShape`'s pipelined-below-regular rejection). (`.collect()` + * on a limit takes `executeTake` and never hits `doExecute`; `.write` / `.toLocalIterator` + * / a non-root position do.) + */ +object EnablePipelinedShuffle extends Rule[SparkPlan] { + + override def apply(plan: SparkPlan): SparkPlan = { + // Shared environment gate (opt-in flag, single-executor local mode, channel manager active), + // identical to the AQE rule's -- see PipelinedShuffleEligibility for why it is a correctness + // gate. It also requires AQE off implicitly: under AQE this rule sees no exchanges. + if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan + + val shuffles = plan.collect { case s: ShuffleExchangeExec => s } + if (shuffles.isEmpty) return plan + + // A reused exchange has more than one consumer; a pipelined producer cannot fan out, so + // leave the whole plan regular rather than produce a rejected job. Check subquery plans + // too (plan.exists walks the operator tree only): today no SQL shape can place a reused + // PIPELINED exchange there -- same-tree reuse is caught here, main-vs-subquery reuse + // never fires because the subquery's own preparation pass (PlanSubqueries -> + // prepareExecutedPlan, which includes this rule) flips its exchanges pipelined BEFORE + // the outer ReuseExchangeAndSubquery compares canonical forms, and subquery-vs-subquery + // duplication is collapsed by MergeScalarSubqueries / subquery reuse first -- but the + // second mechanism is an accident of rule ordering and the third is optimizer behavior, + // so this gate does not rely on either. + if (plan.collectWithSubqueries { case r: ReusedExchangeExec => r }.nonEmpty) { + // Not a warning: this is a normal, expected fallback (reuse is routine optimizer output, + // e.g. self-joins), the query still runs correctly as a regular shuffle, and the user has + // nothing to act on. Log at DEBUG as diagnostic ("why this query did not go pipelined") + // rather than WARN, which would fire on every reuse-bearing query and read as a fault. + logDebug("EnablePipelinedShuffle: plan has a reused exchange; leaving it regular to " + + "avoid a fan-out pipelined job.") + return plan + } + + // An operator that would read a shuffle in a way the channel transport cannot serve, or that + // builds its own hidden regular shuffle, forces the whole plan regular (see class doc). Like + // the reuse fallback this is a normal, expected outcome, so log at DEBUG rather than WARN. + if (readsShuffleThroughUnsupportedConsumer(plan)) { + logDebug("EnablePipelinedShuffle: a shuffle is read through an operator the channel " + + "transport cannot serve (coalesce / cartesian product / a limit operator that builds a " + + "hidden shuffle); leaving the plan regular.") + return plan + } + + plan.transformUp { Review Comment: Addressed by the same `DeserializeToObjectExec` boundary guard as the narrow-consumer issue. The RDD exposed to user code now has regular shuffle dependencies, so adding an RDD shuffle does not create a pipelined-below-regular job. Added AQE-on/off coverage for `reduceByKey` and `repartition`, including result checks. ########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleRendezvous.scala: ########## @@ -0,0 +1,190 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.local.pipelined + +import java.util.concurrent.{ConcurrentHashMap, LinkedBlockingQueue} + +import org.apache.spark.{SparkContext, TaskContext} + +/** + * Process-wide rendezvous between the map (writer) and reduce (reader) sides of an + * in-process pipelined shuffle. One bounded queue exists per + * `(shuffleId, epoch, reducePartitionId)`; every map task writing to a given reduce partition + * shares the queue with the single reduce task that drains it. Queue elements are BATCHES + * of records (`Array[AnyRef]` of pairs, see [[ChannelShuffleWriter]]) or the + * [[EndOfStream]] marker, so the queue's per-operation lock cost is paid per batch, not + * per row. + * + * The `epoch` is the per-run id (the jobId, propagated to both the writer and the reader of + * one gang via a job-level local property -- see `SparkContext.SPARK_PIPELINED_RUN_EPOCH`). A + * shuffleId is RE-RUN within one query (a RangePartitioner sample job then the main job; + * executeTake's per-batch jobs; a classic Dataset re-executing a reused plan), and each run is + * a different job, hence a different epoch. Keying by epoch makes each run's queues and marks + * PHYSICALLY separate: the new run never sees a partition whose reader never started in the old + * run (whose queue still holds stale batches + end-of-stream markers), and a straggler writer + * from an aborted run -- still looping because a task kill need not interrupt the thread -- can + * only touch its OWN (old) epoch's queue, never the new run's. This replaces an earlier design + * that reused one key per shuffleId and tried to reset shared marks between runs, which left + * stale queues and raced stragglers. + * + * This is correct only when producer and consumer tasks are co-resident in the same JVM, + * i.e. a single executor (local mode). The concurrent-stage scheduler co-schedules the two + * stages so both are running, but it does not by itself guarantee co-location; the + * [[PipelinedChannelShuffleManager]] is only intended for single-executor deployments, + * where co-location is automatic. Cross-executor pipelined shuffle is served by the RPC + * streaming shuffle instead. + * + * The queues are bounded, so a fast producer blocks on `put` when the consumer lags -- + * this is the backpressure that keeps the pipelined hand-off memory-bounded. + */ +private[spark] object ChannelShuffleRendezvous { + + /** + * Marker placed on a queue by each map task when it finishes writing to that reduce + * partition. A reader stops once it has seen one marker per map task. + */ + val EndOfStream: AnyRef = new AnyRef + + /** + * The per-run epoch for the current task, read from the job-level local property the + * DAGScheduler stamps on a pipelined job (`SparkContext.SPARK_PIPELINED_RUN_EPOCH` = jobId). + * Both the writer and the reader of one gang read this, so they address the same per-run + * queues. Absent (a core-RDD path that never sets it) defaults to 0; that is fine because such + * a shuffleId is never re-run into a colliding second live run. + */ + def epochOf(tc: TaskContext): Int = + Option(tc) + .flatMap(t => Option(t.getLocalProperty(SparkContext.SPARK_PIPELINED_RUN_EPOCH))) + .map(_.toInt) + .getOrElse(0) + + // State is nested by shuffleId FIRST, then keyed by (epoch, reducePartitionId) within it. The + // outer level exists so the two per-shuffle operations the ContextCleaner drives -- holdsShuffle + // and removeShuffle -- are O(1) map lookups instead of a scan of every live entry: holdsShuffle + // now runs for EVERY shuffle cleaned in a feature-on session, including the regular prefix + // shuffles this feature itself produces, so a flat key made that O(live entries) per cleanup. + // + // Queue values are AnyRef because a queue carries both record batches (Array[AnyRef]) and the + // EndOfStream marker. + private val queues = + new ConcurrentHashMap[Int, ConcurrentHashMap[(Int, Int), LinkedBlockingQueue[AnyRef]]]() + + // (epoch, reducePartitionId) keys, per shuffleId, whose reader has departed (its reduce task + // finished) and will drain no more. A writer stops feeding an abandoned partition and drops the + // rest. This covers the LIVE-partition early-stop case (e.g. a LIMIT reader that pulled enough + // and quit): without it the writer fills the partition's bounded queue and blocks forever. + private val abandoned = + new ConcurrentHashMap[Int, java.util.Set[(Int, Int)]]() + + /** + * Per-queue capacity in BATCHES (not rows), the backpressure bound and the heap-residency + * knob (see spark.shuffle.channel.queueCapacity). Set once by the channel manager at + * construction from that conf; defaults to 64 (with the default 1024-row batch, ~64K rows per + * reduce partition in flight) until a manager sets it. `@volatile` because the manager sets it + * on the driver while writer/reader threads read it. + */ + @volatile private var capacity = 64 + + /** Set the per-queue capacity in batches. Called by the channel manager from its conf. */ + private[pipelined] def setCapacity(batches: Int): Unit = { capacity = batches } + + /** The queue for one `(shuffleId, epoch, reducePartitionId)`, created on first access. */ + def queue(shuffleId: Int, epoch: Int, reducePartitionId: Int): LinkedBlockingQueue[AnyRef] = { Review Comment: Added manager `startRun`/`endRun` hooks tied to the job lifecycle. Job cleanup releases that epoch’s queues and abandon marks, including on failure or cancellation. Late task accesses cannot recreate state for an ended run. Unregister also preserves active epochs until their owning jobs finish. This matters when a concurrent action fails and its SQL cleanup runs while the first action is still using the shuffle. Added tests for repeated actions with a reachable dependency, epoch isolation during cleanup, late task access, failure cleanup, and unregister during an active run. ########## core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala: ########## @@ -2129,8 +2428,24 @@ private[spark] class DAGScheduler( // Job submitted, clear internal data. barrierJobIdToNumTasksCheckFailures.remove(jobId) + // For a pipelined job, stamp the per-run epoch (the jobId) into the job's properties BEFORE + // creating the ActiveJob, so submitMissingTasks -- which clones jobIdToActiveJob(jobId) + // .properties per stage -- carries the SAME epoch to every stage's tasks. Both the producer + // (writer) and the consumer (reader) of the one gang belong to this job, so both read one + // value; a different run is a different job, hence a different epoch, which keys the + // in-process channel rendezvous per run (see SPARK_PIPELINED_RUN_EPOCH). Copy the caller's + // Properties rather than mutating it. Inert for a non-pipelined job (property never set, + // never read). + val jobProperties = + if (hasPipelined && pipelinedManagerWantsLiveReduceHints) { + val p = Utils.cloneProperties(if (properties == null) new Properties() else properties) + p.setProperty(SparkContext.SPARK_PIPELINED_RUN_EPOCH, jobId.toString) Review Comment: The cross-job reuse restriction remains, and it is now documented. Added a non-AQE SQL test that verifies the second action is rejected, its cleanup does not disrupt the owning action, and a subsequent sequential action succeeds. The test runs with shuffle-file cleanup enabled. One qualification to the AQE example: execution of the same adaptive plan is synchronized through `withFinalPlanUpdate`, so that case can serialize rather than reach the cross-job rejection. I have not added multi-job sharing of the channel. ########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/PipelinedChannelShuffleManager.scala: ########## @@ -0,0 +1,154 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.local.pipelined + +import org.apache.spark.{ShuffleDependency, SparkConf, TaskContext} +import org.apache.spark.internal.config +import org.apache.spark.shuffle.{BaseShuffleHandle, PipelinedShuffleManager, ShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} + +/** + * A pipelined shuffle manager whose writer -> reader transport is an in-process bounded + * channel (see [[ChannelShuffleRendezvous]]) rather than the RPC streaming shuffle. It + * serves a [[org.apache.spark.PipelinedShuffleDependency]] on a single executor, letting + * the concurrent-stage scheduler run a shuffle's map and reduce stages at the same time + * while records flow between them in memory -- the in-process pipelined shuffle execution model. + * + * Selected via `spark.shuffle.manager.incremental`. + * + * Unlike the RPC streaming manager, this one needs no `StreamingShuffleOutputTracker`: it + * finds each reader/writer pair through the JVM-local [[ChannelShuffleRendezvous]] rather + * than a directory of writer host/port locations. It therefore declares + * `usesStreamingShuffleOutputTracker = false`, so `SparkEnv` creates no tracker and the + * scheduler registers the shuffle with none (a pipelined stage's availability is tracked on + * the stage itself, not in any output tracker). This is why it implements the + * `PipelinedShuffleManager` trait directly instead of subclassing the concrete streaming + * manager. + * + * This manager deliberately keeps NO per-shuffle registry. An early version recorded each + * shuffle's map-task count at registration and looked it up in getReader -- and lost it when + * an unregisterShuffle arrived BETWEEN registration and the job running, which happens + * legitimately: Dataset.rdd builds the RDD inside a SQL execution scope that ends (and, with + * spark.sql.classic.shuffleDependency.fileCleanup.enabled, removes the shuffle from every + * manager) before any job has run. The reader then saw a missing entry as numMaps = 0 and + * silently under-read the channel. The count is instead stamped into the shuffle handle at + * registration ([[ChannelShuffleHandle.numMaps]]): the handle travels with the dependency + * into every task, a plain Int field survives task serialization (the dependency's own `rdd` + * reference is @transient and is null inside a deserialized task), and no later unregister + * can take it away. + */ +private[spark] class PipelinedChannelShuffleManager(conf: SparkConf) + extends PipelinedShuffleManager { + + // The in-process rendezvous is JVM-local: on a multi-executor deployment each executor would + // get its own empty queue map, and every reader would block forever on data written in some + // other JVM -- a silent hang. Refuse to construct anywhere but local mode, so a + // misconfiguration fails loudly at startup instead. + require(org.apache.spark.util.Utils.isLocalMaster(conf), + "PipelinedChannelShuffleManager is an in-process (single-JVM) transport and requires " + + s"local mode; got master '${conf.get("spark.master", "")}'") + + // Rows accumulated per output partition before a batch is handed across the channel in one + // queue operation. Batching amortizes the queue's per-operation lock cost; per-row hand-off + // measured ~19x slower than a regular shuffle on a 20M-row repartition. + private val batchSize = conf.get(config.SHUFFLE_PIPELINED_CHANNEL_BATCH_SIZE) + + // Per-queue depth in batches (backpressure bound + heap-residency knob). Set the process-wide + // rendezvous from the conf at construction, before any writer/reader creates a queue. + ChannelShuffleRendezvous.setCapacity(conf.get(config.SHUFFLE_PIPELINED_CHANNEL_QUEUE_CAPACITY)) + + override def usesStreamingShuffleOutputTracker: Boolean = false Review Comment: Agreed. The live-reduce hint limits emitted data, not producer computation. The config documentation now states the per-job recomputation cost, including once per output partition for `toLocalIterator`, and recommends caching first. Added AQE-on/off tests using an accumulator to verify producer evaluation counts across repeated iterator actions, alongside result and epoch-cleanup checks. The recomputation behavior remains; this change documents and tests it. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEEnablePipelinedShuffle.scala: ########## @@ -0,0 +1,229 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.adaptive + +import scala.collection.mutable + +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.{BinaryExecNode, CoalesceExec, CollectLimitExec, CollectTailExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.exchange.{PipelinedShuffleEligibility, ReusedExchangeExec, ShuffleExchangeExec} +import org.apache.spark.sql.execution.joins.ShuffledJoin + +/** + * Opt-in (SPARK-57399). Flips eligible [[ShuffleExchangeExec]] + * nodes to `pipelined = true` under AQE, the adaptive counterpart of the non-AQE + * `EnablePipelinedShuffle` preparation rule (which is a no-op once the plan is wrapped in + * `AdaptiveSparkPlanExec`). Runs in `AdaptiveSparkPlanExec.queryStagePreparationRules`, so + * it is re-applied on every replanning round; the decision is deterministic on plan shape, + * and already-flipped exchanges are left alone. + * + * Placement policy. A flipped exchange has no + * map output statistics (it never materializes as a query stage: the DAGScheduler + * gang-runs it inline with its consumer in the final job -- see the pipelined case in + * `AdaptiveSparkPlanExec.createNonResultQueryStages`), so only exchanges whose statistics + * no AQE decision consumes are flipped: + * + * - "free" candidate: the path from the candidate to the plan root crosses no + * stats-sensitive node ([[BinaryExecNode]], another [[ShuffleExchangeExec]], or a + * query stage). Its own coalescing/skew handling is given up; nothing above needed its + * stats. + * - "join-paired" candidate: the immediate shuffle inputs of a [[ShuffledJoin]] whose + * path to the root is otherwise free, flipped only as a symmetric pair (an asymmetric + * flip would leave one side participating in AQE coalesce/skew and the other fixed). + * - everything else stays regular and materializes as usual -- those stages form the + * fully-materialized prefix the scheduler's mixed-job shape requires. + * + * A pipelined exchange supports every + * partitioning, so a SinglePartition exchange in a free position is simply a candidate + * itself. The walk stops below a flipped candidate: exchanges underneath stay regular and + * keep full AQE treatment. Candidates whose canonicalized form occurs more than once in + * the plan (including inside materialized stages and subqueries) are skipped: flipping + * them would trade AQE's stage reuse for duplicate recomputation, and a pipelined producer + * cannot be consumed twice. + */ +object AQEEnablePipelinedShuffle extends Rule[SparkPlan] { + + override def apply(plan: SparkPlan): SparkPlan = { + // Shared environment gate (opt-in flag, single-executor local mode, channel manager active), + // identical to the non-AQE rule's -- see PipelinedShuffleEligibility for why it is a + // correctness gate that must not drift between the two rules. + if (!PipelinedShuffleEligibility.enabled(plan, conf)) return plan + + flipEligibleExchanges(plan) + } + + /** + * The plan-shape core of the rule, factored out of [[apply]]'s environment guards (opt-in flag, + * local mode, channel manager) so it can be unit-tested on a hand-built plan directly. Collects + * the eligible exchanges and returns the plan with each flipped to `pipelined = true`. + */ + private[adaptive] def flipEligibleExchanges(plan: SparkPlan): SparkPlan = { + val shared = if (conf.exchangeReuseEnabled) duplicatedShuffleForms(plan) else Set.empty[Any] + // Collect the exchanges to flip BY IDENTITY (SparkPlan.id, unique per instance), not by the + // node itself: TreeNode overrides hashCode but not equals, so a HashSet[ShuffleExchangeExec] + // matches structurally, and the transformDown below would then flip EVERY exchange + // structurally equal to a collected one -- including a twin the collector deliberately left + // regular on a blocked path. That twin, if it sits below a regular boundary, makes + // classifyJobShuffleShape reject the whole job. Keying on the instance id flips exactly the + // nodes the collector chose, regardless of spark.sql.exchange.reuse (duplicatedShuffleForms, + // the only other guard, is empty when reuse is off). transformDown matches each ORIGINAL node + // before rebuilding it, so its id is the same instance id the collector recorded. + val toFlip = mutable.HashSet.empty[Int] + collectCandidates(plan, blocked = false, shared, toFlip) + if (toFlip.isEmpty) return plan + + // transformDown, NOT transformUp: candidates can be nested (a SinglePartition candidate + // above a hash candidate). transformUp rebuilds children first, so by the time it + // reaches the upper candidate that node is a NEW instance whose (already flipped) child + // no longer matches the collected original structurally, and the upper flip is silently + // dropped -- leaving a regular exchange above a pipelined one, which the scheduler then + // rejects. transformDown hands each candidate to the pattern before its subtree is + // rebuilt, so both nested flips apply. + plan.transformDown { + case s: ShuffleExchangeExec if toFlip.contains(s.id) => s.copy(pipelined = true) + } + } + + private def isCandidate(s: ShuffleExchangeExec, shared: Set[Any]): Boolean = Review Comment: Retained the explicit admission failure, but expanded the config documentation to explain that the default 200 partitions is usually too wide for local execution, that AQE does not coalesce pipelined exchanges, and that all producer and consumer tasks must fit concurrently. The setting guidance is therefore based on total group demand, not just requiring `spark.sql.shuffle.partitions` to be below the core count. The existing admission-failure tests remain. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala: ########## @@ -570,7 +569,15 @@ object ShuffleExchangeExec { rddWithPartitionIds, new PartitionIdPassthrough(part.numPartitions), serializer, - shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics), + // Copy rows only for a transport that hands object references to a concurrent + // consumer (the in-process channel). The RPC streaming transport detaches rows by + // serializing them promptly and must not pay an extra per-row copy on its path. And + // skip it when rddWithPartitionIds already copied: needToCopyObjectsBeforeShuffle makes + // that RDD emit (pid, row.copy()), so a second copy here would be redundant. + shuffleWriterProcessor = createShuffleWriteProcessor( + writeMetrics, + copyRows = SparkEnv.get.pipelinedShuffleManager.requiresDetachedRecords && Review Comment: Applied this suggestion. `needToCopyObjectsBeforeShuffle` now takes the pipelined flag and consults `pipelinedShuffleManager.requiresDetachedRecords` for that path. Rows are copied at the existing shuffle-input construction site. Removed `copyRows`, the write-processor override, and the incorrect comment, so the channel’s copy decision no longer depends on the blocking manager’s thresholds. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [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]
