基于列对 Spark DataFrame 进行分区

Question

我正在尝试使用 groupByKey() 根据列“b”对 Spark DataFrame 进行分区，但我最终在同一分区中有不同的组。

这是我正在使用的数据框和代码：

df:
+---+---+
|  a|  b|
+---+---+
|  4|  2|
|  5|  1|
|  1|  4|
|  2|  2|
+---+---+


 val partitions = df.map(x => x.getLong(1)).distinct().count().toInt
 val df2 = df.map(r => (r.getLong(1), r)).groupByKey(partitions)
 val gb = df2.mapPartitions(iterator => {
            val rows = iterator.toList
            println(rows)
            iterator
            })

The printed rows are:
Partition 1: List((2,CompactBuffer([4,2], [2,2])))
Partition 2: List((4,CompactBuffer([1,4])), (1,CompactBuffer([5,1])))

第 4 组和第 1 组在同一个分区 (2) 中，我想将它们放在不同的分区中，你知道怎么做吗？

Desired output:
Partition 1: List((2,CompactBuffer([4,2], [2,2])))
Partition 2: List((4,CompactBuffer([1,4])))
Partition 3: List((1,CompactBuffer([5,1])))

附：为了给你一些背景信息，我这样做是因为我需要使用来自所有其他行的数据更新 DataFrame 中的行，这些行对特定列共享相同的值。因此 map() 是不够的，我目前正在尝试使用 mapPartitions() ，其中每个分区将包含对特定列具有相同值的所有行。如果您知道更好的方法，请随时告诉我:)

非常感谢！

克莱德X

Answer 1

这听起来像您正在尝试做的事情，可以通过使用窗口函数来完成：https://databricks.com/blog/2015/07/15/introducing-window-functions-in-spark-sql.html