spark foreachPartition, how to get an index of each partition?

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)
})