Memory 通过键重新分区大型数据集并为每个组逐批分别应用 function 的有效方法

Question

我有一个带有"groupName"列的大火花 scala 数据集。 数据记录分布在不同的分区上。 我想按"groupName"将记录分组在一起，逐批收集并在整个批次上应用 function。 “批次”是指同一组的预定义数量的记录（我们称之为maxBatchCount ）。 通过“逐批”，我的意思是我想有效地使用 memory，而不是将所有分区收集到 memory。

更具体地说，批次 function 包括整个批次的序列化、压缩和加密。 稍后将其转换为另一个数据集，使用partitionBy("groupName")将其写入 hdfs。 因此，我无法避免完全洗牌。

有没有一种简单的方法可以做到这一点？ 我做了一些下面描述的尝试，但TL/DR似乎有点过于复杂，最终在 Java memory 问题上失败了。

---

细节

我尝试使用 repartition repartition("groupName") 、 mapPartitions和Iterator的grouped(maxBatchCount)方法的组合，这似乎非常适合该任务。 但是，重新分区仅确保相同groupName的记录将在同一分区中，但单个分区可能具有来自多个不同groupName的记录（如果#groups > #partitions）并且它们可以分散在分区内。 所以现在我仍然需要先在每个分区内进行一些分组。 问题是从 mapPartition 我得到了一个似乎没有这样的 API 的Iterator ，我不想将所有数据收集到 memory。

然后我尝试使用Iterator的partition方法来增强上述解决方案。 这个想法是首先迭代完整的分区以构建所有当前组的Set ，然后使用Iterator.partition为每个当前组构建一个单独的迭代器。 然后像以前一样使用grouped 。

它是这样的——为了说明，我使用了两个 Ints 的简单案例 class， groupName实际上是mod3列，通过对 Range 中的每个number应用modulo 3 function 创建：

  case class Mod3(number: Int, mod3: Int)
  val maxBatchCount = 5
  val df = spark.sparkContext.parallelize(Range(1,21))
     .toDF("number").withColumn("mod3", col("number") % 3)

  // here I choose #partitions < #groups for illustration
  val dff = df.repartition(1, col("mod3"))

  val dsArr = dff.as[Mod3].mapPartitions(partitionIt => {
    // we'll need 2 iterations
    val (it1, it2) = partitionIt.duplicate
    
    // first iterate to create a Set of all present groups
    val mod3set = it1.map(_.mod3).toSet
    
    // build partitioned iterators map (one for each group present)
    var it: Iterator[Mod3] = it2 // init var
    val itMap = mod3set.map(mod3val => {
      val (filteredIt, residueIt) = it.partition(_.mod3 == mod3val)
      val pair = (mod3val -> filteredIt)
      it = residueIt
      pair
    }).toMap

    mod3set.flatMap(mod3val => {
      itMap(mod3val).grouped(maxBatchCount).map(grp => {
        val batch = grp.toList
        batch.map(_.number).toArray[Int] // imagine some other batch function
      })
    }).toIterator
  }).as[Array[Int]]

  val dsArrCollect = dsArr.collect
  dsArrCollect.map(_.toList).foreach(println)

在使用小数据进行测试时，这似乎工作得很好，但是在使用实际数据运行时（在具有 20 个执行程序的实际 spark 集群上，每个 2 个内核）我收到java.lang.OutOfMemoryError: GC overhead limit exceeded

• 请注意，在我的实际数据组中，大小是高度倾斜的，其中一个组的大小大约是所有组的 rest 组合的大小（我猜 GC memory 问题与该组有关）。 因此，我还尝试在repartition中合并一个辅助中性列，但它没有帮助。

将不胜感激这里的任何指示，谢谢！

Answer 1

我认为您对重新分区 + map 分区有正确的方法。 问题是您的 map 分区 function 最终将整个分区加载到 memory 中。

第一个解决方案可能是增加分区的数量，从而减少分区中的组/数据的数量。

另一种解决方案是使用 partitionIt.flatMap 并一次处理 1 条记录，最多仅累积 1 组数据

• 使用 sortWithinPartitions 使同一组中的记录是连续的
• 在 flatMap function 中，累积您的数据并跟踪组更改。