如何使用scala在spark数据帧中找到由另一个列值分组的列中的数组总和

时间:2019-06-26 13:23:38

标签: scala apache-spark dataframe

我有一个dataframe,如下所示

c1             Value
A             Array[47,97,33,94,6]
A             Array[59,98,24,83,3]
A             Array[77,63,93,86,62]
B             Array[86,71,72,23,27]
B             Array[74,69,72,93,7]
B             Array[58,99,90,93,41]
C             Array[40,13,85,75,90]
C             Array[39,13,33,29,14]
C             Array[99,88,57,69,49]

我需要如下输出。

c1             Value
A             Array[183,258,150,263,71]
B             Array[218,239,234,209,75]
C             Array[178,114,175,173,153]

除了对 column 列进行分组并按顺序查找 column 列中的值之和外,什么都不是。 请帮助,我在google中找不到任何方法。

1 个答案:

答案 0 :(得分:1)

它不是很复杂。正如您提到的那样,您可以简单地按“ c1”分组并按索引汇总数组索引的值。

首先生成一些数据:

val df = spark.range(6)
    .select('id % 3 as "c1",
            array((1 to 5).map(_ => floor(rand * 10)) : _*) as "Value")
df.show()
+---+---------------+
| c1|          Value|
+---+---------------+
|  0|[7, 4, 7, 4, 0]|
|  1|[3, 3, 2, 8, 5]|
|  2|[2, 1, 0, 4, 4]|
|  0|[0, 4, 2, 1, 8]|
|  1|[1, 5, 7, 4, 3]|
|  2|[2, 5, 0, 2, 2]|
+---+---------------+

然后,我们需要遍历数组的值以聚合它们。这与我们创建它们的方式非常相似:

val n = 5 // if you know the size of the array
val n = df.select(size('Value)).first.getAs[Int](0) // If you do not
df
    .groupBy("c1")
    .agg(array((0 until n).map(i => sum(col("Value").getItem(i))) :_* ) as "Value")
    .show()
+---+------------------+
| c1|             Value|
+---+------------------+
|  0|[11, 18, 15, 8, 9]|
|  1|  [2, 10, 5, 7, 4]|
|  2|[7, 14, 15, 10, 4]|
+---+------------------+