跨数据集行的数组的元素总和 - Spark Scala

Question

我正在尝试根据“id”列对以下数据集进行分组，并按元素对“values”列中的数组求和。 我如何使用 Scala 在 Spark 中做到这一点？

输入：（2 列的数据集，String 类型的 column1 和 Array[Int] 类型的 column2）

| id | values |
---------------
| A | [12,61,23,43]
| A | [43,11,24,45]
| B | [32,12,53,21]
| C | [11,12,13,14]
| C | [43,52,12,52]
| B | [33,21,15,24]

预期输出：（数据集或数据框）

| id | values |
---------------
| A | [55,72,47,88]
| B | [65,33,68,45]
| C | [54,64,25,66]

注意：结果必须是灵活和动态的。 也就是说，即使有 1000 列或即使文件有几个 TB 或 PB，该解决方案仍然应该有效。

Answer 1

当你说它必须是灵活的时，我有点不确定你的意思，但在我的头顶上，我可以想到几种方法。 第一个（在我看来是最漂亮的）一个使用udf ：

// Creating a small test example
val testDF = spark.sparkContext.parallelize(Seq(("a", Seq(1,2,3)), ("a", Seq(4,5,6)), ("b", Seq(1,3,4)))).toDF("id", "arr")
val sum_arr = udf((list: Seq[Seq[Int]]) => list.transpose.map(arr => arr.sum))

testDF
  .groupBy('id)
  .agg(sum_arr(collect_list('arr)) as "summed_values")

但是，如果您有数十亿个相同的 ID，那么collect_list当然会成为一个问题。 在这种情况下，您可以执行以下操作：

testDF
  .flatMap{case Row(id: String, list: Seq[Int]) => list.indices.map(index => (id, index, list(index)))}
  .toDF("id", "arr_index", "arr_element")
  .groupBy('id, 'arr_index)
  .agg(sum("arr_element") as "sum")
  .groupBy('id)
  .agg(collect_list('sum) as "summed_values")

Answer 2

以下单行解决方案对我有用

ds.groupBy("Country").agg(array((0 until n).map(i => sum(col("Values").getItem(i))) :_* ) as "Values")