[SPARK-24615-NG1]Add API to support configuring preferred resource per rdd

core/src/main/scala/org/apache/spark/rdd/CartesianRDD.scala

@@ -22,6 +22,7 @@ import java.io.{IOException, ObjectOutputStream}
 import scala.reflect.ClassTag
 
 import org.apache.spark._
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 
 private[spark]
@@ -70,6 +71,13 @@ class CartesianRDD[T: ClassTag, U: ClassTag](
     (rdd1.preferredLocations(currSplit.s1) ++ rdd2.preferredLocations(currSplit.s2)).distinct
   }
 
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    val currSplit = split.asInstanceOf[CartesianPartition]
+    rdd1.getPreferredResources(currSplit.s1)
+      .mergeOther(rdd2.getPreferredResources(currSplit.s2))
+      .mergeOther(super.getPreferredResources(split))
+  }
+
   override def compute(split: Partition, context: TaskContext): Iterator[(T, U)] = {
     val currSplit = split.asInstanceOf[CartesianPartition]
     for (x <- rdd1.iterator(currSplit.s1, context);

core/src/main/scala/org/apache/spark/rdd/CoGroupedRDD.scala

@@ -25,6 +25,7 @@ import scala.reflect.ClassTag
 
 import org.apache.spark._
 import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.util.Utils
 import org.apache.spark.util.collection.{CompactBuffer, ExternalAppendOnlyMap}
@@ -126,6 +127,15 @@ class CoGroupedRDD[K: ClassTag](
     array
   }
 
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    split.asInstanceOf[CoGroupPartition].narrowDeps.map {
+      case Some(s) => s.rdd.getPreferredResources(s.split)
+      case None => PreferredResources.EMPTY
+    }
+      .reduce { (s1, s2) => s1.mergeOther(s2) }
+      .mergeOther(super.getPreferredResources(split))
+  }
+
   override val partitioner: Some[Partitioner] = Some(part)
 
   override def compute(s: Partition, context: TaskContext): Iterator[(K, Array[Iterable[_]])] = {

core/src/main/scala/org/apache/spark/rdd/CoalescedRDD.scala

@@ -25,6 +25,7 @@ import scala.language.existentials
 import scala.reflect.ClassTag
 
 import org.apache.spark._
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 
 /**
@@ -122,6 +123,16 @@ private[spark] class CoalescedRDD[T: ClassTag](
   override def getPreferredLocations(partition: Partition): Seq[String] = {
     partition.asInstanceOf[CoalescedRDDPartition].preferredLocation.toSeq
   }
+
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    val parentResources = split.asInstanceOf[CoalescedRDDPartition].parents.map { p =>
+      prev.getPreferredResources(p)
+    }
+
+    parentResources
+      .reduce { (r1, r2) => r1.mergeOther(r2) }
+      .mergeOther(super.getPreferredResources(split))
+  }
 }
 
 /**

core/src/main/scala/org/apache/spark/rdd/PartitionPruningRDD.scala

@@ -21,6 +21,7 @@ import scala.reflect.ClassTag
 
 import org.apache.spark.{NarrowDependency, Partition, TaskContext}
 import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.rdd.resource.PreferredResources
 
 private[spark] class PartitionPruningRDDPartition(idx: Int, val parentSplit: Partition)
   extends Partition {
@@ -66,8 +67,12 @@ class PartitionPruningRDD[T: ClassTag](
 
   override protected def getPartitions: Array[Partition] =
     dependencies.head.asInstanceOf[PruneDependency[T]].partitions
-}
 
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    prev.getPreferredResources(split.asInstanceOf[PartitionPruningRDDPartition].parentSplit)
+      .mergeOther(super.getPreferredResources(split))
+  }
+}
 
 @DeveloperApi
 object PartitionPruningRDD {

core/src/main/scala/org/apache/spark/rdd/PartitionerAwareUnionRDD.scala

@@ -22,6 +22,7 @@ import java.io.{IOException, ObjectOutputStream}
 import scala.reflect.ClassTag
 
 import org.apache.spark.{OneToOneDependency, Partition, SparkContext, TaskContext}
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 
 /**
@@ -94,6 +95,15 @@ class PartitionerAwareUnionRDD[T: ClassTag](
     location.toSeq
   }
 
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    val partition = split.asInstanceOf[PartitionerAwareUnionRDDPartition]
+    partition.rdds.zip(partition.parents).map { case (rdd, p) =>
+      rdd.getPreferredResources(p)
+    }
+      .reduce { (s1, s2) => s1.mergeOther(s2) }
+      .mergeOther(super.getPreferredResources(split))
+  }
+
   override def compute(s: Partition, context: TaskContext): Iterator[T] = {
     val parentPartitions = s.asInstanceOf[PartitionerAwareUnionRDDPartition].parents
     rdds.zip(parentPartitions).iterator.flatMap {

core/src/main/scala/org/apache/spark/rdd/PartitionwiseSampledRDD.scala

@@ -22,6 +22,7 @@ import java.util.Random
 import scala.reflect.ClassTag
 
 import org.apache.spark.{Partition, TaskContext}
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 import org.apache.spark.util.random.RandomSampler
 
@@ -67,4 +68,9 @@ private[spark] class PartitionwiseSampledRDD[T: ClassTag, U: ClassTag](
     thisSampler.setSeed(split.seed)
     thisSampler.sample(firstParent[T].iterator(split.prev, context))
   }
+
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    prev.getPreferredResources(split.asInstanceOf[PartitionwiseSampledRDDPartition].prev)
+      .mergeOther(super.getPreferredResources(split))
+  }
 }

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -40,6 +40,7 @@ import org.apache.spark.partial.BoundedDouble
 import org.apache.spark.partial.CountEvaluator
 import org.apache.spark.partial.GroupedCountEvaluator
 import org.apache.spark.partial.PartialResult
+import org.apache.spark.rdd.resource.{PreferredResources, ResourcesPreferenceBuilder}
 import org.apache.spark.storage.{RDDBlockId, StorageLevel}
 import org.apache.spark.util.{BoundedPriorityQueue, Utils}
 import org.apache.spark.util.collection.{OpenHashMap, Utils => collectionUtils}
@@ -136,6 +137,14 @@ abstract class RDD[T: ClassTag](
    */
   protected def getPreferredLocations(split: Partition): Seq[String] = Nil
 
+  private[spark] def getPreferredResources(split: Partition): PreferredResources = {
+    dependencies.headOption.flatMap {
+      case d: OneToOneDependency[_] =>
+        Some(d.rdd.getPreferredResources(split).mergeOther(this.preferredResources))
+      case _ => None
+    }.getOrElse(preferredResources)
+  }
+
   /** Optionally overridden by subclasses to specify how they are partitioned. */
   @transient val partitioner: Option[Partitioner] = None
 
@@ -158,6 +167,21 @@ abstract class RDD[T: ClassTag](
     this
   }
 
+  @transient protected var preferredResources: PreferredResources = PreferredResources.EMPTY
+
+  private[spark] def setPreferredResources(info: PreferredResources): Unit = {
+    this.preferredResources = info
+  }
+
+  def withResources(): ResourcesPreferenceBuilder[T] = {
+    new ResourcesPreferenceBuilder(this)
+  }
+
+  def clearResources(): this.type = {
+    preferredResources = PreferredResources.EMPTY
+    this
+  }
+
   /**
    * Mark this RDD for persisting using the specified level.
    *

core/src/main/scala/org/apache/spark/rdd/SubtractedRDD.scala

@@ -30,6 +30,7 @@ import org.apache.spark.Partitioner
 import org.apache.spark.ShuffleDependency
 import org.apache.spark.SparkEnv
 import org.apache.spark.TaskContext
+import org.apache.spark.rdd.resource.PreferredResources
 
 /**
  * An optimized version of cogroup for set difference/subtraction.
@@ -84,6 +85,15 @@ private[spark] class SubtractedRDD[K: ClassTag, V: ClassTag, W: ClassTag](
     array
   }
 
+  override private[spark] def getPreferredResources(split: Partition): PreferredResources = {
+    split.asInstanceOf[CoGroupPartition].narrowDeps.map {
+      case Some(s) => s.rdd.getPreferredResources(s.split)
+      case None => PreferredResources.EMPTY
+    }
+      .reduce { (s1, s2) => s1.mergeOther(s2) }
+      .mergeOther(super.getPreferredResources(split))
+  }
+
   override val partitioner = Some(part)
 
   override def compute(p: Partition, context: TaskContext): Iterator[(K, V)] = {

core/src/main/scala/org/apache/spark/rdd/UnionRDD.scala

@@ -26,6 +26,7 @@ import scala.reflect.ClassTag
 
 import org.apache.spark.{Dependency, Partition, RangeDependency, SparkContext, TaskContext}
 import org.apache.spark.annotation.DeveloperApi
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 
 /**
@@ -39,7 +40,7 @@ import org.apache.spark.util.Utils
  */
 private[spark] class UnionPartition[T: ClassTag](
     idx: Int,
-    @transient private val rdd: RDD[T],
+    @transient val rdd: RDD[T],
     val parentRddIndex: Int,
     @transient private val parentRddPartitionIndex: Int)
   extends Partition {
@@ -112,4 +113,10 @@ class UnionRDD[T: ClassTag](
     super.clearDependencies()
     rdds = null
   }
+
+  override private[spark] def getPreferredResources(split: Partition): PreferredResources = {
+    val partition = split.asInstanceOf[UnionPartition[_]]
+    partition.rdd.getPreferredResources(partition.parentPartition)
+      .mergeOther(super.getPreferredResources(split))
+  }
 }

core/src/main/scala/org/apache/spark/rdd/ZippedPartitionsRDD.scala

@@ -22,11 +22,12 @@ import java.io.{IOException, ObjectOutputStream}
 import scala.reflect.ClassTag
 
 import org.apache.spark.{OneToOneDependency, Partition, SparkContext, TaskContext}
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 
 private[spark] class ZippedPartitionsPartition(
     idx: Int,
-    @transient private val rdds: Seq[RDD[_]],
+    @transient val rdds: Seq[RDD[_]],
     @transient val preferredLocations: Seq[String])
   extends Partition {
 
@@ -74,6 +75,15 @@ private[spark] abstract class ZippedPartitionsBaseRDD[V: ClassTag](
     super.clearDependencies()
     rdds = null
   }
+
+  override private[spark] def getPreferredResources(split: Partition): PreferredResources = {
+    val partition = split.asInstanceOf[ZippedPartitionsPartition]
+    partition.rdds.zip(partition.partitions).map { case (rdd, p) =>
+      rdd.getPreferredResources(p)
+    }
+      .reduce { (s1, s2) => s1.mergeOther(s2) }
+      .mergeOther(super.getPreferredResources(split))
+  }
 }
 
 private[spark] class ZippedPartitionsRDD2[A: ClassTag, B: ClassTag, V: ClassTag](

core/src/main/scala/org/apache/spark/rdd/ZippedWithIndexRDD.scala

@@ -20,6 +20,7 @@ package org.apache.spark.rdd
 import scala.reflect.ClassTag
 
 import org.apache.spark.{Partition, TaskContext}
+import org.apache.spark.rdd.resource.PreferredResources
 import org.apache.spark.util.Utils
 
 private[spark]
@@ -67,4 +68,9 @@ class ZippedWithIndexRDD[T: ClassTag](prev: RDD[T]) extends RDD[(T, Long)](prev)
     val parentIter = firstParent[T].iterator(split.prev, context)
     Utils.getIteratorZipWithIndex(parentIter, split.startIndex)
   }
+
+  private[spark] override def getPreferredResources(split: Partition): PreferredResources = {
+    prev.getPreferredResources(split.asInstanceOf[ZippedWithIndexRDDPartition].prev)
+      .mergeOther(super.getPreferredResources(split))
+  }
 }