diff --git a/core/src/main/scala/org/apache/spark/Dependency.scala b/core/src/main/scala/org/apache/spark/Dependency.scala index da97aff5b344d..05dfed3372af7 100644 --- a/core/src/main/scala/org/apache/spark/Dependency.scala +++ b/core/src/main/scala/org/apache/spark/Dependency.scala @@ -258,6 +258,59 @@ class ShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( } +/** + * :: DeveloperApi :: + * A [[ShuffleDependency]] whose output can be read incrementally: a consumer stage may begin + * reading the shuffle output while the producer stage is still running, rather than waiting for the + * producer's full, materialized output. + * + * This is a first-class dependency kind, peer to [[NarrowDependency]] and [[ShuffleDependency]]. It + * is intended to be the marker the `DAGScheduler` will use to decide that the producer and consumer + * stages connected by this edge may run concurrently (a "pipelined group"), and that the shuffle + * layer should serve this shuffle with an incremental shuffle implementation. A plain + * [[ShuffleDependency]] keeps the existing semantics: its output is fully materialized before any + * consumer reads it. + * + * This class only declares the capability. On its own it behaves exactly like its parent + * [[ShuffleDependency]] -- code that matches `ShuffleDependency` continues to treat it as an + * ordinary (materialized) shuffle -- so introducing it changes no existing behavior. The concurrent + * scheduling and incremental-shuffle behavior will be added separately by the components that match + * on this type. + * + * The name is *pipelined* rather than *streaming*: reading producer output as it is produced is a + * general execution capability (software-pipelining of dependent stages), not specific to + * streaming. Streaming / real-time mode is the first caller, but nothing here is + * streaming-specific. + * + * Note: the parent's `checksumMismatchFullRetryEnabled` / + * `checksumMismatchQueryLevelRollbackEnabled` parameters are intentionally not exposed here, so they + * stay at their `false` defaults for a pipelined shuffle. Their checksum retry / query-level + * rollback recomputes and re-runs succeeding stages after a mismatch; in a pipelined group any + * failure aborts the whole group and the caller reruns from scratch, so that stage-level recompute + * never fires -- the mechanism is moot by construction (it is also incompatible with a consumer that + * has already read the output incrementally). Leaving the params unset keeps the idiom unreachable. + */ +@DeveloperApi +class PipelinedShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag]( + _rdd: RDD[_ <: Product2[K, V]], + partitioner: Partitioner, + serializer: Serializer = SparkEnv.get.serializer, + keyOrdering: Option[Ordering[K]] = None, + aggregator: Option[Aggregator[K, V, C]] = None, + mapSideCombine: Boolean = false, + shuffleWriterProcessor: ShuffleWriteProcessor = new ShuffleWriteProcessor, + rowBasedChecksums: Array[RowBasedChecksum] = ShuffleDependency.EMPTY_ROW_BASED_CHECKSUMS) + extends ShuffleDependency[K, V, C]( + _rdd, + partitioner, + serializer, + keyOrdering, + aggregator, + mapSideCombine, + shuffleWriterProcessor, + rowBasedChecksums) + + /** * :: DeveloperApi :: * Represents a one-to-one dependency between partitions of the parent and child RDDs. diff --git a/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala b/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala index 4d5f599fb12ab..13a37708c029a 100644 --- a/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala +++ b/core/src/test/scala/org/apache/spark/shuffle/ShuffleDependencySuite.scala @@ -17,6 +17,7 @@ package org.apache.spark.shuffle import org.apache.spark._ +import org.apache.spark.rdd.RDD case class KeyClass() @@ -64,4 +65,55 @@ class ShuffleDependencySuite extends SparkFunSuite with LocalSparkContext { assert(dep.combinerClassName == None) } + test("PipelinedShuffleDependency is a ShuffleDependency and preserves its fields") { + sc = new SparkContext("local", "test", conf.clone()) + val rdd: RDD[(KeyClass, ValueClass)] = + sc.parallelize(1 to 5, 4).map(_ => (KeyClass(), ValueClass())) + val partitioner = new HashPartitioner(2) + val dep = new PipelinedShuffleDependency[KeyClass, ValueClass, ValueClass](rdd, partitioner) + + // It is a first-class dependency kind, but IS-A ShuffleDependency: code that matches + // ShuffleDependency continues to see it as an ordinary shuffle, so it changes no existing + // behavior. The concurrent-scheduling behavior is keyed on the subtype elsewhere. + assert(dep.isInstanceOf[ShuffleDependency[_, _, _]]) + assert(dep.partitioner === partitioner) + assert(dep.keyClassName == classOf[KeyClass].getName) + assert(dep.valueClassName == classOf[ValueClass].getName) + assert(dep.rdd === rdd) + // Construction goes through the normal ShuffleDependency path: the shuffle is registered with + // the ShuffleManager (a handle is produced) and a second instance gets a distinct shuffleId. + assert(dep.shuffleHandle != null) + val dep2 = new PipelinedShuffleDependency[KeyClass, ValueClass, ValueClass](rdd, partitioner) + assert(dep2.shuffleId != dep.shuffleId) + + // The checksum retry / query-level rollback params are intentionally not exposed by the + // subclass (see PipelinedShuffleDependency scaladoc): their stage-level recompute is moot for a + // pipelined group, so they must stay at their false defaults. + assert(!dep.checksumMismatchFullRetryEnabled) + assert(!dep.checksumMismatchQueryLevelRollbackEnabled) + } + + test("PipelinedShuffleDependency forwards non-default constructor args to ShuffleDependency") { + sc = new SparkContext("local", "test", conf.clone()) + val rdd: RDD[(KeyClass, ValueClass)] = + sc.parallelize(1 to 5, 4).map(_ => (KeyClass(), ValueClass())) + val partitioner = new HashPartitioner(2) + val aggregator = new Aggregator[KeyClass, ValueClass, ValueClass]( + v => v, (c, _) => c, (c1, _) => c1) + val dep = new PipelinedShuffleDependency[KeyClass, ValueClass, ValueClass]( + rdd, partitioner, aggregator = Some(aggregator), mapSideCombine = true) + + // The subclass forwards its constructor args positionally to ShuffleDependency; assert the + // non-default ones land where expected rather than on the wrong parameter. + assert(dep.aggregator.contains(aggregator)) + assert(dep.mapSideCombine) + } + + test("an ordinary ShuffleDependency is not a PipelinedShuffleDependency") { + sc = new SparkContext("local", "test", conf.clone()) + val rdd = sc.parallelize(1 to 5, 4).map(_ => (KeyClass(), ValueClass())).groupByKey() + val dep = rdd.dependencies.head.asInstanceOf[ShuffleDependency[_, _, _]] + assert(!dep.isInstanceOf[PipelinedShuffleDependency[_, _, _]]) + } + }