Github user squito commented on a diff in the pull request:
    --- Diff: core/src/main/scala/org/apache/spark/rdd/MapPartitionsRDD.scala 
    @@ -34,8 +35,29 @@ private[spark] class MapPartitionsRDD[U: ClassTag, T: 
       override def getPartitions: Array[Partition] = firstParent[T].partitions
    -  override def compute(split: Partition, context: TaskContext): 
Iterator[U] =
    -    f(context, split.index, firstParent[T].iterator(split, context))
    +  override def compute(split: Partition, context: TaskContext): 
Iterator[U] = {
    +    val input = firstParent[T].iterator(split, context)
    +    // If our task has data property accumulators we need to keep track of 
which partitions
    +    // are fully processed and what we are currently processing.
    +    if (context.taskMetrics.hasDataPropertyAccumulators()) {
    +      // Use -1 for shuffle id since we are not in a shuffle.
    +      val shuffleWriteId = -1
    +      // Set the ID of the RDD and partition being processed. We need to 
do this per
    +      // element since we chain the iterator transformations together
    +      val data ={x =>
    +        context.setRDDPartitionInfo(id, shuffleWriteId, split.index)
    +        x
    +      }
    +      val wrappedData = Utils.signalWhenEmpty(data,
    +        () => context.taskMetrics.markFullyProcessed(id, shuffleWriteId, 
    +      // We also set it before the first call to the user function in case 
the user provides a
    +      // function which access a data property accumulator before 
accessing any elements.
    +      context.setRDDPartitionInfo(id, shuffleWriteId, split.index)
    +      f(context, split.index, wrappedData)
    +    } else {
    +      f(context, split.index, input)
    --- End diff --
    I'd add a comment here if there aren't any data property accumulators, we 
want to avoid the overhead that comes with the the extra iterator wrapping.

If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at or file a JIRA ticket
with INFRA.

To unsubscribe, e-mail:
For additional commands, e-mail:

Reply via email to