viirya commented on code in PR #58097: URL: https://github.com/apache/spark/pull/58097#discussion_r3985505299
########## python/pyspark/sql/tests/test_pipelined_shuffle.py: ########## @@ -0,0 +1,105 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import unittest + +from pyspark.sql.functions import udf +from pyspark.sql.types import LongType +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class PipelinedShuffleTests(ReusedSQLTestCase): Review Comment: Fixed in `bbb1cefaa94`. Added `pyspark.sql.tests.test_pipelined_shuffle` to `pyspark_sql.python_test_goals` and verified that `determine_dangling_python_tests` no longer reports it. This was an integration step I missed. The earlier local suite results did not establish that CI could start; thanks for catching it. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala: ########## @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.spark.SparkEnv +import org.apache.spark.internal.Logging +import org.apache.spark.internal.config +import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager +import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, DeserializeToObjectExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.joins.CartesianProductExec +import org.apache.spark.sql.internal.SQLConf + +/** + * Shared environment gate for the two pipelined-shuffle enabling rules + * ([[EnablePipelinedShuffle]] non-AQE and `AQEEnablePipelinedShuffle` under AQE). This is a + * CORRECTNESS gate, not cosmetics: flipping an exchange to pipelined while the incremental manager + * is still the RPC streaming one would route it to an untested transport that reports + * `requiresDetachedRecords = false`, so the SQL layer would skip the row copy and silently corrupt + * rows shared across the writer/reader threads. Both rules must apply the identical gate, so it + * lives here rather than being copy-pasted into each `apply` (where the two could drift and split + * AQE vs non-AQE behavior). Each rule keeps only its own plan-shape logic. + */ +private[sql] object PipelinedShuffleEligibility extends Logging { + + // The flag/manager mismatch is a start-up misconfiguration, so warn once per JVM rather than on + // every query planned in the session. + private val mismatchWarned = new AtomicBoolean(false) + + /** Operators whose consumers cannot safely drain a bounded, single-reader channel. */ + def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match { + case _: CoalesceExec | _: CartesianProductExec | _: DeserializeToObjectExec | + _: CollectLimitExec | _: CollectTailExec | _: TakeOrderedAndProjectExec => true + case _ => false + } + + /** + * Estimate the complete group's task demand without executing or materializing the plan. + * Each exchange contributes its producer width; the root contributes the consumer width. + * Unknown widths conservatively retain regular execution. The scheduler still checks actual + * RDD widths and free slots at submission, including contention from other jobs. + */ + def fitsLocalCapacity(plan: SparkPlan, candidates: Set[Int]): Boolean = { + val exchanges = plan.collect { + case s: ShuffleExchangeExec if s.pipelined || candidates.contains(s.id) => s + }.groupBy(_.id).values.map(_.head).toSeq + val widths = plan.outputPartitioning.numPartitions +: + exchanges.map(_.child.outputPartitioning.numPartitions) + val sc = plan.session.sparkContext + val slots = sc.maxNumConcurrentTasks(sc.resourceProfileManager.defaultResourceProfile) + val demand = widths.map(_.toLong).sum + val fits = widths.forall(_ > 0) && demand <= slots + if (!fits) { + logDebug(s"Pipelined shuffle: estimated stage widths ${widths.mkString(",")} cannot " + + s"fit $slots local task slots; retaining regular shuffles.") + } + fits + } + + /** + * Whether the pipelined channel transport may be used for `plan` at all, independent of plan + * shape. Requires: the opt-in flag on; single-executor local mode (the in-process channel + * transport needs producer and consumer in one JVM); and the configured incremental manager + * actually being the in-process channel manager. Returns false (with a DEBUG diagnostic) when + * any gate fails, so the caller leaves the plan regular. + */ + def enabled(plan: SparkPlan, conf: SQLConf): Boolean = { + if (!conf.localPipelinedShuffleEnabled) { + return false + } + if (plan.session == null || !plan.session.sparkContext.isLocal) { + return false + } + // Batch only. `IncrementalExecution.preparations` inherits QueryExecution's list, so without + // this gate a streaming plan would be rewritten here: every micro-batch exchange (the + // state-store shuffles, the static side of a stream-static join) would be flipped to pipelined + // BEFORE `MarkPipelinedShuffleForRealTimeMode` runs. That contradicts what the Real-Time Mode + // rule deliberately does -- it leaves the static side regular, because pulling it into the gang + // would demand slots for stages that must instead finish first, failing admission. Streaming + // marks its own pipelined boundaries; this opt-in batch path must not pre-empt that decision. + // (`logicalLink.exists(_.isStreaming)` is the same signal InsertAdaptiveSparkPlan uses to keep + // AQE off streaming plans.) + // Dataset.rdd exposes arbitrary consumers beyond the SQL plan, including RDD shuffles, + // multi-partition reads and repeated reads of the same partition. Cached inputs also hide + // shuffle lineage: cache hits skip those readers, while misses may require regular stages. + def hasUnsupportedBoundary(p: SparkPlan): Boolean = p match { + case _: DeserializeToObjectExec | _: InMemoryTableScanExec => true Review Comment: Added `RDDScanExec` and `ExternalRDDScanExec` to the shared boundary exclusions, and applied `withRegularShuffle` at all four `UnionLoopExec` RDD access sites. The new RDD-input test runs `createDataFrame` over an RDD-level `reduceByKey`, followed by a SQL aggregation and `collect`, without passing through the Dataset RDD exit helper. The recursive-CTE test includes repartitioned iterations, disables conversion to LocalRelation, and uses a small channel capacity. It checks regular RDD lineage and the complete result. Both tests pass with AQE on and off. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleEligibility.scala: ########## @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.spark.SparkEnv +import org.apache.spark.internal.Logging +import org.apache.spark.internal.config +import org.apache.spark.shuffle.local.pipelined.PipelinedChannelShuffleManager +import org.apache.spark.sql.execution.{CoalesceExec, CollectLimitExec, CollectTailExec, DeserializeToObjectExec, SparkPlan, TakeOrderedAndProjectExec} +import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.columnar.InMemoryTableScanExec +import org.apache.spark.sql.execution.joins.CartesianProductExec +import org.apache.spark.sql.internal.SQLConf + +/** + * Shared environment gate for the two pipelined-shuffle enabling rules + * ([[EnablePipelinedShuffle]] non-AQE and `AQEEnablePipelinedShuffle` under AQE). This is a + * CORRECTNESS gate, not cosmetics: flipping an exchange to pipelined while the incremental manager + * is still the RPC streaming one would route it to an untested transport that reports + * `requiresDetachedRecords = false`, so the SQL layer would skip the row copy and silently corrupt + * rows shared across the writer/reader threads. Both rules must apply the identical gate, so it + * lives here rather than being copy-pasted into each `apply` (where the two could drift and split + * AQE vs non-AQE behavior). Each rule keeps only its own plan-shape logic. + */ +private[sql] object PipelinedShuffleEligibility extends Logging { + + // The flag/manager mismatch is a start-up misconfiguration, so warn once per JVM rather than on + // every query planned in the session. + private val mismatchWarned = new AtomicBoolean(false) + + /** Operators whose consumers cannot safely drain a bounded, single-reader channel. */ + def isUnsupportedConsumer(plan: SparkPlan): Boolean = plan match { + case _: CoalesceExec | _: CartesianProductExec | _: DeserializeToObjectExec | + _: CollectLimitExec | _: CollectTailExec | _: TakeOrderedAndProjectExec => true + case _ => false + } + + /** + * Estimate the complete group's task demand without executing or materializing the plan. + * Each exchange contributes its producer width; the root contributes the consumer width. + * Unknown widths conservatively retain regular execution. The scheduler still checks actual + * RDD widths and free slots at submission, including contention from other jobs. + */ + def fitsLocalCapacity(plan: SparkPlan, candidates: Set[Int]): Boolean = { + val exchanges = plan.collect { + case s: ShuffleExchangeExec if s.pipelined || candidates.contains(s.id) => s + }.groupBy(_.id).values.map(_.head).toSeq + val widths = plan.outputPartitioning.numPartitions +: Review Comment: Agreed. I retained the conservative unknown-width fallback for this PR rather than adding scan-specific width estimation. SQLConf and the PR description now explicitly state that ordinary non-bucketed file scans remain regular. Added a Parquet-read/repartition test in both AQE modes that verifies correct results and no pipelined exchange. This makes the current limitation explicit; it does not add pipelining support for those scans. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/QueryExecution.scala: ########## @@ -409,6 +409,23 @@ class QueryExecution( def assertExecutedPlanPrepared(): Unit = executedPlan + /** + * RDD consumers and partition-at-a-time iterators can outlive one SQL job. Give them a + * separate plan whose shuffles retain output, without changing this execution's plan. + * Keep the cloned session for lazy planning and AQE, not just for this method's call. + */ + private[sql] def withRegularShuffle: QueryExecution = { + val session = SparkSession.getOrCloneSessionWithConfigsOff( Review Comment: The fallback is now memoized per QueryExecution. It returns the original execution when the feature is disabled, the shared eligibility gate rejects the plan, or there is no shuffle. The shuffle check looks through adaptive plans and query stages. Added tests verifying that shuffle-free queries retain the original QE and repeated fallback requests return the same QE/session. Updated the iterator tests to verify reuse of retained regular-shuffle output across repeated actions. I did not call `SparkSession.close()` after an action: in classic Spark that stops the shared SparkContext. The fallback session remains alive with its QE, and artifact cleanup uses the existing ArtifactManager Cleaner. This reduces unnecessary clones and repeated planning; it does not introduce a new explicit session-disposal mechanism. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala: ########## @@ -0,0 +1,728 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.exchange + +import java.util.concurrent.{Callable, CountDownLatch, Executors, TimeUnit} + +import org.apache.spark.{PipelinedShuffleDependency, SparkEnv, SparkFunSuite} +import org.apache.spark.rdd.RDD +import org.apache.spark.shuffle.local.pipelined.ChannelShuffleRendezvous +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.storage.{RDDBlockId, StorageLevel} + +/** + * End-to-end SQL coverage of the pipelined channel path: a batch query whose hash + * exchange is rewritten to a pipelined shuffle (EnablePipelinedShuffle) and served by the + * in-process channel manager (PipelinedChannelShuffleManager), run through the + * concurrent-stage scheduler on a single executor. Self-manages its SparkSession because the + * shuffle manager and AQE-off gate are start-up configs. + */ +class PipelinedShuffleSqlSuite extends SparkFunSuite + with AdaptiveSparkPlanHelper with PipelinedShuffleTestSession { + + private def withPipelinedSession(body: SparkSession => Unit): Unit = + withPipelinedSession("pipelined-shuffle-sql", aqe = false)(body) + + test("batch repartition($k) runs end-to-end through the pipelined channel shuffle") { + withPipelinedSession { spark => + import spark.implicits._ + val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 10)).repartition($"k") + + // Single action only: a pipelined shuffle is single-shot, so collect exactly once and + // derive everything from that one result. + val rows = df.select($"id").as[Long].collect() + val ids = rows.toSet + + // The rule fired and the exchange is pipelined. + val pipelinedExchanges = collect(df.queryExecution.executedPlan) { + case s: ShuffleExchangeExec if s.pipelined => s + } + assert(pipelinedExchanges.nonEmpty, + s"expected a pipelined ShuffleExchangeExec; plan was:\n${df.queryExecution.executedPlan}") + + // Correctness: the same 1000 ids, repartitioned, all present exactly once. + assert(rows.length === 1000, s"expected 1000 rows, got ${rows.length}") + assert(ids === (0L until 1000L).toSet) + } + } + + test("a single keyed groupBy runs end-to-end through the pipelined channel shuffle") { + withPipelinedSession { spark => + import spark.implicits._ + val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7)) + .groupBy($"k").count() + + // Single action only (pipelined shuffle is single-shot). + val counts = df.as[(Long, Long)].collect().toMap + val pipelined = collect(df.queryExecution.executedPlan) { + case s: ShuffleExchangeExec if s.pipelined => s + } + assert(pipelined.nonEmpty, + s"expected pipelined exchange; plan:\n${df.queryExecution.executedPlan}") + // Each residue class 0..6 of 0..999. + val expected = (0L until 1000L).groupBy(_ % 7).map { case (k, vs) => (k, vs.size.toLong) } + assert(counts === expected) + } + } + + test("groupBy with ORDER BY (hash + range exchanges) is all-pipelined") { + // A trailing ORDER BY adds a RANGE exchange (global sort with 4 shuffle partitions -> + // RangePartitioning) on top of the groupBy's hash exchange. The relaxed rule pipelines + // BOTH (a mixed pipelined/regular job would be rejected). Range is the interesting case: + // RangePartitioner construction runs a SAMPLE job over the exchange's child -- which here + // reads the pipelined hash shuffle -- before the main job runs, so this also exercises + // two successive jobs over the same single-shot pipelined producer. + withPipelinedSession { spark => + import spark.implicits._ + val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 7)) + .groupBy($"k").count().orderBy($"k") + + val rows = df.as[(Long, Long)].collect() + val exchanges = collect(df.queryExecution.executedPlan) { + case s: ShuffleExchangeExec => s + } + assert(exchanges.nonEmpty && exchanges.forall(_.pipelined), + s"every exchange should be pipelined; plan:\n${df.queryExecution.executedPlan}") + // Pin the partitioning shapes so this test can't silently stop covering range. + val partitionings = exchanges.map(_.outputPartitioning.getClass.getSimpleName).sorted + assert(exchanges.exists(_.outputPartitioning.isInstanceOf[ + org.apache.spark.sql.catalyst.plans.physical.RangePartitioning]), + s"expected a RangePartitioning exchange, got: $partitionings; " + + s"plan:\n${df.queryExecution.executedPlan}") + // Result is correct AND globally ordered by k. + val expected = (0L until 1000L).groupBy(_ % 7).map { case (k, vs) => (k, vs.size.toLong) } + .toSeq.sortBy(_._1) + assert(rows.toSeq === expected) + } + } + + test("repartitionByRange (pure range exchange) runs through the pipelined channel shuffle") { + // A range exchange directly over the scan: RangePartitioner samples the scan (a job with + // no shuffle at all), then the main job runs the pipelined range shuffle. Verifies the + // channel transport is agnostic to the partitioner kind, and rows land range-partitioned. + withPipelinedSession { spark => + import spark.implicits._ + val df = spark.range(0, 1000, 1, 2).withColumn("k", ($"id" % 100)) + .repartitionByRange($"k") + // spark_partition_id() records which output partition each row landed in without + // leaving the DataFrame API (Dataset.rdd would execute a separate QueryExecution). + .select($"k", org.apache.spark.sql.functions.spark_partition_id().as("p")) + + val partitioned = df.as[(Long, Int)].collect().map { case (k, p) => (p, k) } + val exchanges = collect(df.queryExecution.executedPlan) { + case s: ShuffleExchangeExec => s + } + assert(exchanges.nonEmpty && exchanges.forall(_.pipelined), + s"expected a pipelined exchange; plan:\n${df.queryExecution.executedPlan}") + assert(exchanges.exists(_.outputPartitioning.isInstanceOf[ + org.apache.spark.sql.catalyst.plans.physical.RangePartitioning]), + s"expected RangePartitioning; plan:\n${df.queryExecution.executedPlan}") + + // No rows lost, and the partitioning is a genuine range split: key ranges of distinct + // partitions must not overlap. + assert(partitioned.length === 1000) + val ranges = partitioned.groupBy(_._1).map { case (p, rows) => + (p, rows.map(_._2).min, rows.map(_._2).max) + }.toSeq.sortBy(_._2) + ranges.sliding(2).foreach { + case Seq((p1, _, max1), (p2, min2, _)) => + assert(max1 <= min2, s"partitions $p1 and $p2 overlap: max($p1)=$max1 > min($p2)=$min2") + case _ => + } + } + } + + test("sort-merge join (both sides hash-exchanged) is all-pipelined and correct") { + // A shuffled join is the last TPC-DS transport shape not yet covered: both join inputs + // get a hash ShuffleExchangeExec. Disable broadcast so the join is a SortMergeJoin with + // two real shuffles; the relaxed rule pipelines both, and the concurrent-stage group + // (two producers + the join stage) runs together. + withPipelinedSession { spark => + import spark.implicits._ + spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1") + // Two structurally DIFFERENT inputs so exchange reuse does not collapse them into one + // ReusedExchange (which the rule would skip). Different ranges + key expressions. + val left = spark.range(0, 200, 1, 2).withColumn("k", ($"id" % 10)) + .select($"k", $"id".as("lv")) + val right = spark.range(0, 120, 1, 2).withColumn("k", ($"id" % 6)) + .select($"k", $"id".as("rv")) + val joined = left.join(right, "k") + + val rows = joined.select($"k", $"lv", $"rv").as[(Long, Long, Long)].collect() + val exchanges = collect(joined.queryExecution.executedPlan) { + case s: ShuffleExchangeExec => s + } + assert(exchanges.length >= 2 && exchanges.forall(_.pipelined), + s"both join inputs should be pipelined; plan:\n${joined.queryExecution.executedPlan}") + + // Ground truth: an equi-join on k over the two relations. + val l = (0L until 200L).map(i => (i % 10, i)) + val r = (0L until 120L).map(i => (i % 6, i)) + val expected = (for ((lk, lv) <- l; (rk, rv) <- r if lk == rk) yield (lk, lv, rv)).toSet + assert(rows.toSet === expected) + } + } + + test("global aggregate (single-partition exchange) runs through the pipelined channel") { + // An ungrouped aggregate requires AllTuples, planned as a SinglePartition exchange: the + // channel's numPartitions == 1 degenerate case (everything routes to queue 0). + withPipelinedSession { spark => + import spark.implicits._ + val df = spark.range(0, 1000, 1, 2).agg(org.apache.spark.sql.functions.sum($"id")) + + val result = df.as[Long].collect() + val exchanges = collect(df.queryExecution.executedPlan) { + case s: ShuffleExchangeExec => s + } + assert(exchanges.nonEmpty && exchanges.forall(_.pipelined), + s"expected a pipelined exchange; plan:\n${df.queryExecution.executedPlan}") + assert(exchanges.exists(_.outputPartitioning == + org.apache.spark.sql.catalyst.plans.physical.SinglePartition), + s"expected a SinglePartition exchange; plan:\n${df.queryExecution.executedPlan}") + assert(result.toSeq === Seq((0L until 1000L).sum)) + } + } + + private def assertRegularRDD(root: RDD[_]): Unit = { + val visited = scala.collection.mutable.Set.empty[Int] + val pending = scala.collection.mutable.Stack[RDD[_]](root) + while (pending.nonEmpty) { + val rdd = pending.pop() + if (visited.add(rdd.id)) { + rdd.dependencies.foreach { dep => + assert(!dep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) + pending.push(dep.rdd) + } + } + } + } + + for (aqe <- Seq(false, true)) { + test(s"typed and Python RDD exports retain regular shuffle lineage with AQE=$aqe") { + withPipelinedSession("pipelined-rdd-exports", aqe) { spark => + import spark.implicits._ + val df = spark.range(0, 4000, 1, 2).repartition(4) + assert(df.collect().length === 4000) + val typed = df.groupByKey(_ % 4).mapGroups { (key, rows) => + (key, rows.size.toLong) + } + val rdd = typed.rdd + assertRegularRDD(rdd) + assert(rdd.coalesce(1).collect().map(_._2).sum === 4000L) + val pythonRDD = df.asInstanceOf[org.apache.spark.sql.classic.Dataset[Long]] + .javaToPython.rdd + assertRegularRDD(pythonRDD) + assert(pythonRDD.coalesce(1).count() > 0) + val fromRDD = spark.createDataFrame(rdd).repartition(4).rdd + assertRegularRDD(fromRDD) + assert(fromRDD.coalesce(1).count() === 4L) + } + } + + test(s"lazy checkpoints retain regular shuffle lineage with AQE=$aqe") { + withPipelinedSession("pipelined-checkpoint", aqe) { spark => + withTempDir { dir => + spark.sparkContext.setCheckpointDir(dir.getCanonicalPath) + for (reliable <- Seq(false, true)) { + val df = spark.range(0, 4000, 1, 2).repartition(4) + assert(df.collect().length === 4000) + val checkpointed = if (reliable) { + df.checkpoint(eager = false) + } else { + df.localCheckpoint(eager = false) + } + val rdd = checkpointed.rdd + assertRegularRDD(rdd) + assert(rdd.coalesce(1).count() === 4000L) + assert(checkpointed.collect().sorted === (0L until 4000L).toArray) + } + } + } + } + + test(s"SQL cursor computes its producer once with AQE=$aqe") { + withPipelinedSession("pipelined-cursor", aqe) { spark => + spark.conf.set("spark.sql.scripting.enabled", "true") + spark.conf.set("spark.sql.scripting.cursorEnabled", "true") + spark.conf.set("spark.sql.classic.shuffleDependency.fileCleanup.enabled", "false") + val evaluated = spark.sparkContext.longAccumulator("cursor producer rows") + spark.udf.register("record_cursor_row", (id: Long) => { + evaluated.add(1) + id + }) + val result = spark.sql( + """BEGIN + | DECLARE v BIGINT; + | DECLARE total BIGINT DEFAULT 0; + | DECLARE i INT DEFAULT 0; + | DECLARE c CURSOR FOR + | SELECT /*+ REPARTITION(4) */ record_cursor_row(id) FROM range(100); + | OPEN c; + | WHILE i < 100 DO + | FETCH c INTO v; + | SET total = total + v; + | SET i = i + 1; + | END WHILE; + | CLOSE c; + | VALUES (total); + |END""".stripMargin) + assert(result.collect().head.getLong(0) === (0L until 100L).sum) + assert(evaluated.value === 100L) + } + } + + test(s"Dataset.rdd supports narrow and shuffle consumers with AQE=$aqe") { + withPipelinedSession("pipelined-rdd-boundary", aqe) { spark => + val rdd = spark.range(0, 4000, 1, 2).repartition(4).toDF().rdd + // Assert eligibility before executing a shape that would hang with the channel. + assertRegularRDD(rdd) + assert(rdd.coalesce(2).count() === 4000L) + assert(rdd.union(rdd).count() === 8000L) + assert(rdd.zip(rdd).count() === 4000L) + val grouped = rdd.map(r => (r.getLong(0) % 2, 1L)).reduceByKey(_ + _) + assert(grouped.collect().toMap === Map(0L -> 2000L, 1L -> 2000L)) + assert(rdd.repartition(2).count() === 4000L) + } + } + + test(s"cache construction and partial eviction use regular shuffles with AQE=$aqe") { + withPipelinedSession("pipelined-cache", aqe) { spark => + val df = spark.range(0, 4000, 1, 2).repartition(4).persist(StorageLevel.MEMORY_ONLY) + try { + val classicDf = df.asInstanceOf[org.apache.spark.sql.classic.Dataset[_]] + val cached = spark.sharedState.cacheManager.lookupCachedData(classicDf).get + .cachedRepresentation.cacheBuilder + assert(collect(cached.cachedPlan) { + case s: ShuffleExchangeExec if s.pipelined => s + }.isEmpty) + val expected = (0L until 4000L).toArray + assert(df.collect().sorted === expected) + val buffers = cached.cachedColumnBuffers + assert(buffers.getNumPartitions > 1) + SparkEnv.get.blockManager.removeBlock(RDDBlockId(buffers.id, 0)) Review Comment: Strengthened the cache test with `batchSize = 8` and `queueCapacity = 1`, so the data volume exceeds what an unread channel partition could buffer. An accumulator in the cached projection now records 4,000 evaluations on initial population, 5,000 after evicting partition 0, and still 5,000 after another read. This verifies that only the evicted partition is recomputed. The test retains the plan-policy assertion before execution. Also corrected the AQE repeated-action comment: local cleanup makes the stage unavailable and the next action recomputes it. That test does not exercise a genuine `FetchFailed` retry. ########## core/src/main/scala/org/apache/spark/shuffle/local/pipelined/ChannelShuffleWriterReader.scala: ########## @@ -0,0 +1,359 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.local.pipelined + +import java.util.Arrays + +import org.apache.spark.{SparkContext, SparkEnv, TaskContext} +import org.apache.spark.scheduler.MapStatus +import org.apache.spark.shuffle.{BaseShuffleHandle, ShuffleReader, ShuffleReadMetricsReporter, ShuffleWriteMetricsReporter, ShuffleWriter} + +/** + * Map-side of the in-process pipelined shuffle. Each input record is routed to the reduce + * partition its key hashes to and accumulated in a per-partition batch; a FULL batch (an + * `Array[AnyRef]` of `batchSize` pairs) is pushed onto that partition's shared queue in one + * queue operation (see [[ChannelShuffleRendezvous]]) -- the consumer stage, running + * concurrently, drains it batch by batch. No serialization, no disk, no network. + * + * Batching is what makes the transport viable for large unaggregated shuffles: the queue + * costs a lock acquisition per operation (~hundreds of ns under producer/consumer + * contention), so handing rows across one at a time costs that PER ROW -- measured at ~19x + * slower than a regular shuffle on a 20M-row repartition. Batching divides the lock traffic + * by `batchSize`, the same lesson as any object-batch transport. A batch + * array is handed off to the consumer and never touched again by the writer (a fresh array + * is allocated after each put), so ownership transfer is clean across threads. + */ +private[spark] class ChannelShuffleWriter[K, V]( + handle: BaseShuffleHandle[K, V, _], + mapId: Long, + batchSize: Int, + writeMetrics: ShuffleWriteMetricsReporter) + extends ShuffleWriter[K, V] with org.apache.spark.internal.Logging { + + require(batchSize > 0, s"batchSize must be positive, got $batchSize") + + private val dep = handle.dependency + private val partitioner = dep.partitioner + private val numPartitions = partitioner.numPartitions + private val shuffleId = handle.shuffleId + + // Per-run epoch (the jobId), read from the job-level local property the DAGScheduler set for a + // pipelined job. The reader of this gang reads the SAME value, so both address the same + // per-run queues in the process-wide rendezvous; a re-run of this shuffleId is a different job + // and gets a different epoch, keeping its queues physically separate. Absent (a core-RDD test + // path that never sets it, and never re-runs a shuffleId concurrently) means epoch 0. + private val runEpoch = ChannelShuffleRendezvous.epochOf(TaskContext.get()) + + // The reduce partitions this job actually reads, from the producer stage's task property + // (set by the DAGScheduler from the result stage's partitions). A record routed to a + // partition NOT in this set has no consumer -- putting it would fill that partition's + // bounded queue and, because the writer interleaves all partitions on one thread, block + // the writer before it can feed even the read partitions or emit their end-of-stream, + // deadlocking the job. So such records are dropped. Absent property (None) means every + // partition is live (the normal full-read case: collect, count, a full-partition job) and + // nothing is dropped. + private val liveReducePartitions: Option[Set[Int]] = + Option(TaskContext.get()) + .flatMap(tc => + Option(tc.getLocalProperty(SparkContext.SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS))) + .map(_.split(",").filter(_.nonEmpty).map(_.toInt).toSet) + + // Per-partition liveness, precomputed ONCE from the (static) live set: true iff a consumer + // reads this reduce partition at all. This is the hot-path gate -- checked per input record -- + // so it is a plain Array[Boolean] load, not a boxed Set lookup: on a large repartition the + // per-record path must not allocate (the transport's whole point is amortizing per-row cost). + // Absent property means every partition is live. The OTHER half of "worth writing" -- + // abandonment, which happens at runtime when a reader departs early (e.g. LIMIT) -- is dynamic + // and is checked where it matters (at hand-off, in putUnlessAbandoned), NOT per record: + // accumulating a few more rows into an in-memory batch for a since-abandoned partition is + // harmless because that batch is never put (putUnlessAbandoned drops it). + private val liveMask: Array[Boolean] = { + val mask = Array.fill(numPartitions)(true) + liveReducePartitions.foreach { live => + var p = 0 + while (p < numPartitions) { mask(p) = live.contains(p); p += 1 } + } + mask + } + + // Hand a batch to a partition's queue, but do NOT block forever if its reader departs: + // poll with a short timeout and bail out the moment the partition becomes abandoned. This + // is the cooperative unblock for the early-stop case -- abandon() also drains the queue to + // release a parked put, and this re-check ensures the writer then stops rather than + // re-filling. Returns false if the partition was abandoned before the batch was accepted. + // On a successful hand-off, records the batch's records and the time spent (including any + // backpressure wait) against the write metrics; a dropped/abandoned batch counts nothing, + // since those records are never shuffled out. `records` is the number of pairs in `batch` + // (a full batch is `batchSize`, a trimmed tail is shorter; the end-of-stream marker is 0). + private def putUnlessAbandoned(pid: Int, batch: AnyRef, records: Int): Boolean = { + val q = ChannelShuffleRendezvous.queue(shuffleId, runEpoch, pid) + val start = System.nanoTime() + while (!ChannelShuffleRendezvous.isAbandoned(shuffleId, runEpoch, pid)) { + // Wake on a task kill even without thread interruption. `isAbandoned` is set only by a + // reduce task that actually STARTED (its completion listener); if the reader for pid never + // started -- the group aborts while this producer is already filling queues (an + // unserializable consumer task, an early failure of another member, a job cancel) -- the + // mark never appears and this offer loop would park forever, pinning the executor slot + // (spark.job.interruptOnCancel defaults to false, so the kill does not interrupt the + // thread). Checking the TaskContext interrupt flag each cycle is the symmetric escape to + // the reader's takeItem. + Option(TaskContext.get()).foreach(_.killTaskIfInterrupted()) + if (q.offer(batch, 100, java.util.concurrent.TimeUnit.MILLISECONDS)) { + // A successful offer can race abandon(): abandon does `add(mark)` then `q.clear()`, so if + // it ran between the isAbandoned check above and this offer, our batch lands AFTER the + // clear and would be stranded in the queue (no reader will ever drain it). Re-check and + // clear it ourselves so nothing is left behind. The reader has departed, so discarding is + // correct; and it keeps the queue empty for removeShuffle rather than pinning a batch. + if (ChannelShuffleRendezvous.isAbandoned(shuffleId, runEpoch, pid)) { + q.clear() + return false + } + if (records > 0) { + writeMetrics.incRecordsWritten(records.toLong) + writeMetrics.incWriteTime(System.nanoTime() - start) + } + return true + } + } + false + } + + override def write(records: Iterator[Product2[K, V]]): Unit = { + // No stale-state reset is needed here: this run's queues and abandoned marks are keyed by + // runEpoch (the jobId), so an EARLIER run of this shuffleId (a RangePartitioner sampling job + // then the main job; executeTake batches; a re-executed classic plan) used a different epoch + // and its leftovers are physically separate -- this run starts against empty per-epoch state. + + // One in-progress batch per reduce partition, plus its fill count. A partition's batch array + // is allocated LAZILY, on its first record (batches(pid) starts null), so a map task pays for + // only the partitions it actually writes -- a wide shuffle (thousands of partitions) or a + // partial read (liveMask leaves most partitions dead) does not eagerly allocate + // numPartitions * batchSize empty slots up front. + val batches = new Array[Array[AnyRef]](numPartitions) + val sizes = new Array[Int](numPartitions) + // Records skipped because their reduce partition has no consumer (see liveMask). + // Reported at the end of write(). + var droppedRecords = 0L + + while (records.hasNext) { + val rec = records.next() + val pid = partitioner.getPartition(rec._1) + // Only accumulate for partitions a consumer reads (liveMask). Abandonment is not checked + // here -- it is handled at hand-off in putUnlessAbandoned (see liveMask's comment). + if (!liveMask(pid)) { + // This record is routed to a reduce partition the driver said no consumer reads, so it is + // dropped -- see liveMask. Count it: dropping is CORRECT only if the live set was computed + // correctly, and everything else here fails loudly (a wrong width fails a require, a + // reader-less live partition hangs the writer), while an under-approximated live set would + // instead lose rows quietly. A non-zero count at the end of a job whose result looks wrong + // is the thread to pull. + droppedRecords += 1 + } else { + if (batches(pid) == null) batches(pid) = new Array[AnyRef](batchSize) + // Records must already be detached from the producer's reused row buffers by the time + // they reach here (the producer reuses its output UnsafeRow across iterations, and the + // consumer reads on another thread). The copy is done in the SQL layer's + // ShuffleWriteProcessor for the pipelined path -- where InternalRow.copy() is available Review Comment: Corrected the comment to point to `ShuffleExchangeExec.prepareShuffleDependency` and `needToCopyObjectsBeforeShuffle` for the pipelined copy decision. The reference to the removed `ShuffleWriteProcessor` override is gone. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [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]
