scala - Spark foreachPartition，如何获取每个分区的索引？

Question

spark foreachPartition，如何获取分区的索引(或者序列号，或者识别分区的东西)？

val docs: RDD[String] = ...

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

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

Answer 1

您可以使用TaskContext (How to get ID of a map task in Spark?):

import org.apache.spark.TaskContext

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