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