在Spark/Scala中，如何对数据框的多列进行逐元素数组求和？

Question

我有一个数据框，其中可能有多个数组类型的列，例如“Array1”，“Array2”等。这些数组列应该具有相同数量的元素。我需要计算一个新的数组类型列，它将按元素逐个相加。我该怎么做？

Spark版本 = 2.3

例如：

输入：

|Column1| ... |ArrayColumn2|ArrayColumn3|
|-------| --- |------------|------------|
|T1     | ... |[1, 2 , 3]  | [2, 5, 7]

输出：

|Column1| ... |AggregatedColumn|
|-------| --- |------------|
|T1.    | ... |[3, 7 , 10]

数组列数不固定，因此我需要一个通用的解决方案。我将有一个列列表，需要对其进行聚合。

谢谢！

- Diamondhead

---

所有的数组列都有相同数量的元素吗？ - Srinivas

是的，数组列将具有相同数量的元素。我会更新问题以澄清这一点。 - Diamondhead

Answer 1

考虑使用 inline 和高阶函数 aggregate（在 Spark 2.4+ 中可用）来计算从数组类型列中逐元素求和，然后使用 groupBy/agg 将逐元素求和分组回数组：

val df = Seq(
  (101, Seq(1, 2), Seq(3, 4), Seq(5, 6)),
  (202, Seq(7, 8), Seq(9, 10), Seq(11, 12))
).toDF("id", "arr1", "arr2", "arr3")

val arrCols = df.columns.filter(_.startsWith("arr")).map(col)

对于Spark 3.0+

df.
  withColumn("arr_structs", arrays_zip(arrCols: _*)).
  select($"id", expr("inline(arr_structs)")).
  select($"id", aggregate(array(arrCols: _*), lit(0), (acc, x) => acc + x).as("pos_elem_sum")).
  groupBy("id").agg(collect_list($"pos_elem_sum").as("arr_elem_sum")).
  show
// +---+------------+
// | id|arr_elem_sum|
// +---+------------+
// |101|     [9, 12]|
// |202|    [27, 30]|
// +---+------------+

对于Spark 2.4及以上版本

df.
  withColumn("arr_structs", arrays_zip(arrCols: _*)).
  select($"id", expr("inline(arr_structs)")).
  select($"id", array(arrCols: _*).as("arr_pos_elems")).
  select($"id", expr("aggregate(arr_pos_elems, 0, (acc, x) -> acc + x)").as("pos_elem_sum")).
  groupBy("id").agg(collect_list($"pos_elem_sum").as("arr_elem_sum")).
  show

对于Spark 2.3或更低版本

val sumArrElems = udf{ (arr: Seq[Int]) => arr.sum }

df.
  withColumn("arr_structs", arrays_zip(arrCols: _*)).
  select($"id", expr("inline(arr_structs)")).
  select($"id", sumArrElems(array(arrCols: _*)).as("pos_elem_sum")).
  groupBy("id").agg(collect_list($"pos_elem_sum").as("arr_elem_sum")).
  show

- Leo C

---

@Diamondhead，刚注意到您已经指定使用Spark 2.3。回答已扩展以涵盖各种Spark版本。 - Leo C

Answer 2

像 array(ArrayColumn2[0]+ArrayColumn3[0], ArrayColumn2[1]+...) 这样的SQL表达式可用于计算预期结果。

val df = ...

//get all array columns
val arrayCols = df.schema.fields.filter(_.dataType.isInstanceOf[ArrayType]).map(_.name)

//get the size of the first array of the first row
val firstArray = arrayCols(0)
val arraySize = df.selectExpr(s"size($firstArray)").first().getAs[Int](0)

//generate the sql expression for the sums
val sums = (for( i <-0 to arraySize-1)
  yield arrayCols.map(c=>s"$c[$i]").mkString("+")).mkString(",")
//sums = ArrayColumn2[0]+ArrayColumn3[0],ArrayColumn2[1]+ArrayColumn3[1],ArrayColumn2[2]+ArrayColumn3[2]

//create a new column using sums
df.withColumn("AggregatedColumn", expr(s"array($sums)")).show()

输出：

+-------+------------+------------+----------------+
|Column1|ArrayColumn2|ArrayColumn3|AggregatedColumn|
+-------+------------+------------+----------------+
|     T1|   [1, 2, 3]|   [2, 5, 7]|      [3, 7, 10]|
+-------+------------+------------+----------------+

使用这个单一（长）的SQL表达式将避免在网络上传输数据，从而提高性能。

- werner