[SPARK-58078][CORE] Add PipelinedShuffleDependency

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.

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[_, _, _]])
+  }
+
 }