如何通过mapPartitions迭代器创建DataFrame？

Question

我有一个带有id列的DataFrame。我想对每个id的行进行一些计算（不仅是聚合），并输出每个id一行的新DataFrame，其中包含计算结果。

我正在尝试通过重新分配ID然后使用mapPartitions来做到这一点：

df.repartition(col("id")).mapPartitions(iter => {
    val dfSubset = // iter to DataFrame?
    // Computations on dfSubset
})

但是如何从iter创建一个DataFrame？然后，目标是在包含ID的所有行的DataFrame dfSubset上进行计算。

编辑：

repartition(col("id"))不会为每个id创建1个分区。我们应该改用groupBy("id")。

Answer 1

您要寻找的是对groupBy做一个id，然后定义自己的User Defined Aggregate Function。如果需要所有列，则可以构造所述列的结构并传递给聚合函数。

df
.groupBy("id")
.agg(myUdaf(struct(df.columns.filter(_ != "id").map(col(_)):_*)).as("result")).show()