Поэлементная сумма массива по строкам набора данных - Spark Scala

Question

Я пытаюсь сгруппировать указанный ниже набор данных на основе столбца «id» и поэлементно суммировать массивы в столбце «values». Как мне это сделать в Spark, используя Scala?

Ввод: (набор данных из 2 столбцов, столбец 1 типа String и столбец 2 типа Array [Int])

| 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 столбцов или даже если файл имеет несколько ТБ или ПБ, решение все равно остается в силе.

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")

Однако, если у вас есть миллиарды идентичных идентификаторов, 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")