Spark driver获取每个task返回值方法

在Spark计算中，如果想要在driver端获取executor执行task返回的结果，在实际的操作中可能会想到2种方法：一是将每条数据打上标签输出到结果表中，再通过sql执行查询操作；二是通过Spark自定义累加器，就可以在executor端将结果累加到driver端使用，不过这种方式的效率也不高，具体的实现也很麻烦。

下面介绍一种常用的简便操作方法：

其实，这种操作我们最先想到的应该是count函数，因为他就是将task的返回值返回到driver端，然后进行聚合的。我们可以从idea count函数点击进去，可以看到

 def count(): Long = sc.runJob(this, Utils.getIteratorSize _).sum

也即是sparkcontext的runJob方法。

Utils.getIteratorSize _这个方法主要是计算每个iterator的元素个数，也即是每个分区的元素个数，返回值就是元素个数：

/**
   * Counts the number of elements of an iterator using a while loop rather than calling
   * [[scala.collection.Iterator#size]] because it uses a for loop, which is slightly slower
   * in the current version of Scala.
   */
  def getIteratorSize[T](iterator: Iterator[T]): Long = {
    var count = 0L
    while (iterator.hasNext) {
      count += 1L
      iterator.next()
    }
    count
  }

然后就是runJob返回的是一个数组，每个数组的元素就是我们task执行函数的返回值，然后调用sum就得到我们的统计值了。

那么我们完全可以借助这个思路实现我们开头的目标，下面直接上代码了：

def main(args: Array[String]): Unit = {
  val sparkConf: SparkConf = new SparkConf().setAppName("Master").setMaster("local[2]")
  val sc = new SparkContext(sparkConf)

  val rdd: RDD[(String, String)] = sc
    .textFile("D:\\demo.txt")
    .map(line => {
      val arr: Array[String] = line.split(",")
      (arr(0), arr(1))
    })

  rdd.count()
  val func = (iter: Iterator[(String, String)]) => {
    var count = 0
    iter.foreach(each => {
      count += 1
    })
    (TaskContext.getPartitionId(), count)
  }

  val tuples: Array[(Int, Int)] = sc.runJob(rdd, func)
  tuples.foreach(println)
  sc.stop()
}