spark foreachPartition，如何获取每个分区的索引？

Question

spark foreachPartition , how to get an index of the partition (or sequence number, or something to identify the partition)?

val docs: RDD[String] = ...

println("num partitions: " + docs.getNumPartitions)

docs.foreachPartition((it: Iterator[String]) => {
  println("partition index: " + ???)
  it.foreach(...)
})

Answer 1

You can use TaskContext ( How to get ID of a map task in Spark? ):

import org.apache.spark.TaskContext

rdd.foreachPartition((it: Iterator[String]) => {
  println(TaskContext.getPartitionId)
})

Answer 2

Not exactly identical, but you can use RDD.mapPartitionsWithIndex and return an Iterator[Unit] as a result:

val rdd: RDD[Unit] = docs.mapPartitionsWithIndex { case (idx, it) => 
  println("partition index: " + ???)
  it.foreach(...)
}

But then you have to remember to materialize the RDD

An alternative would be to use mapPartitionsWithIndex to do logic related to transforming the data, and then using foreachRDD just to send the data externally.