如何使用scala查找sparkDataframe中由另一列值分组的列中的数组之和

jaql4c8m  于 2021-07-14  发布在  Spark
关注(0)|答案(1)|浏览(356)

我有一个 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]

它只不过是将列c1分组,然后按顺序查找列value中的值之和。请帮忙,我在谷歌找不到任何方法。

ejk8hzay

ejk8hzay1#

这不是很复杂。正如您所提到的,您可以简单地按“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]|
+---+------------------+

相关问题