将聚合列添加到 Spark DataFrame

Question

我有一个 Spark DataFrame，如下所示：

| id | value | bin |
|----+-------+-----|
|  1 |   3.4 |   2 |
|  2 |   2.6 |   1 |
|  3 |   1.8 |   1 |
|  4 |   9.6 |   2 |

我有一个函数f，它接受一个值数组并返回一个数字。我想在上面的数据框中添加一列，其中每一行中新列的值是 f 对于所有具有相同 bin 条目的 value 条目的值，即：

| id | value | bin | f_value       |
|----+-------+-----+---------------|
|  1 |   3.4 |   2 | f([3.4, 9.6]) |
|  2 |   2.6 |   1 | f([2.6, 1.8]) |
|  3 |   1.8 |   1 | f([2.6, 1.8]) |
|  4 |   9.6 |   2 | f([3.4, 9.6]) |

由于我需要聚合每个bin 的所有values，我不能使用withColumn 函数来添加这个新列。在用户定义的聚合函数进入 Spark 之前，最好的方法是什么？

Answer 1

以下代码未经测试，只是一个想法。

在 Hive 中，可以使用 collect_list 函数这样完成。

val newDF = sqlContext.sql(
    "select bin, collect_list() from aboveDF group by bin")

下一个 join aboveDF 和 newDF 在 bin 上。

这是你要找的吗？