scala - 如何在每个条目的所有行条目上获得 scala-spark 中数组类型列的平均值？

Question

我得到了一个包含 512 个双元素的数组列，并且想要获得平均值。以一个length=3的数组列为例:

val x = Seq("2 4 6", "0 0 0").toDF("value").withColumn("value", split($"value", " "))
x.printSchema()
x.show()


root
 |-- value: array (nullable = true)
 |    |-- element: string (containsNull = true)

+---------+
|    value|
+---------+
|[2, 4, 6]|
|[0, 0, 0]|
+---------+

需要以下结果:

x.select(..... as "avg_value").show()

------------
|avg_value |
------------
|[1,2,3]   |
------------

Answer 1

将每个数组元素视为列并计算平均值，然后使用这些列构造数组:

val array_size = 3
val avgAgg = for (i <- 0 to array_size -1) yield avg($"value".getItem(i))
df.select(array(avgAgg: _*).alias("avg_value")).show(false)

给出:

+---------------+
|avg_value      |
+---------------+
|[1.0, 2.0, 3.0]|
+---------------+