apache-spark - 如何检查给定键的所有记录是否已经在同一个分区中？

Question

我想尽可能避免按键重新分区数据集，并知道给定键的所有记录是否已经在同一个分区中。

Spark 中是否有内置函数可以给我答案？

Answer 1

不是内置的，但如果您假设特定的分区器，则很容易实现您自己的功能:

import org.apache.spark.rdd.RDD
import org.apache.spark.Partitioner
import scala.reflect.ClassTag

def checkDistribution[K : ClassTag, V : ClassTag](
   rdd: RDD[(K, V)], partitioner: Partitioner) = 
  // If partitioner is set we compare partitioners 
  rdd.partitioner.map(_ == partitioner).getOrElse {
    // Otherwise check if correct number of partitions 
    rdd.partitions.size ==  partitioner.numPartitions &&
    //  And check if distribution matches partitioner
    rdd.keys.mapPartitionsWithIndex((i, iter) => 
      Iterator(iter.forall(x => partitioner.getPartition(x) == i))
    ).fold(true)(_ && _)
  }

几个测试:

import org.apache.spark.HashPartitioner

val rdd = sc.range(0, 20, 5).map((_, None))

未分区，无效分布:

checkDistribution(rdd, new HashPartitioner(10))

Boolean = false

已分区，无效的分区程序:

checkDistribution(
  rdd.partitionBy(new HashPartitioner(5)),
  new HashPartitioner(10)
)

Boolean = false

已分区的有效分区程序:

checkDistribution(
  rdd.partitionBy(new HashPartitioner(10)),
  new HashPartitioner(10)
)

Boolean = true

未分区，有效分布:

checkDistribution(
  rdd.partitionBy(new HashPartitioner(10)).map(identity),
  new HashPartitioner(10)
)

Boolean = true

在不假设特定分区器的情况下，想到的唯一选项需要 shuffle，因此它不太可能是一种改进。

def checkDistribution[K : ClassTag, V : ClassTag](rdd: RDD[(K, V)]) =
   rdd.keys.mapPartitionsWithIndex((i, iter) => iter.map((_, i)))
     .combineByKey(
       x => Seq(x), 
       (x: Seq[Int], y: Int) => x, 
       (x: Seq[Int], y: Seq[Int]) => x ++ y)  // Should be more or less OK
     .values
     .mapPartitions(iter => Iterator(iter.forall(_.size == 1)))
     .fold(true)(_ && _)

一种可能的改进是您可以使用相同的逻辑来自动定义 Partitioner对于数据。如果您 collectAsMap之前 values并检查所有 Seqs大小为 1，您有一个有效的分区器，可以保证没有网络流量。