Spark Scala实现DataFrame列内多值的平均值与元素计数

针对数百万行规模的Spark DataFrame，要计算逗号分隔数值字符串列的平均值和元素个数，优先采用Spark内置高阶函数实现，避免数据膨胀与shuffle操作，确保处理性能最优。

核心思路 #

直接对单条记录的_2列做列内转换计算，不拆分行：

1. 将逗号分隔的字符串拆分为数组
2. 把字符串数组转换为数值类型数组
3. 用数组长度直接获取元素计数
4. 通过数组求和除以长度得到平均值（附带空值/空数组的边界处理）

代码实现 #

Spark 3.0+ 版本（推荐，性能最优）

利用transform、array_sum、size等内置高阶函数，无需自定义UDF：

import org.apache.spark.sql.functions._

val resultDF = originalDF
  .withColumn("_2_array", split(col("_2"), ","))
  .withColumn("_2_num_array", transform(col("_2_array"), x => x.cast(DoubleType)))
  .withColumn("Count_2", size(col("_2_num_array")))
  .withColumn("Average_2", 
    when(col("Count_2") === 0, lit(null))
    .otherwise(array_sum(col("_2_num_array")) / col("Count_2"))
  )
  .select("_1", "_2", "Average_2", "Count_2")
  // 可选：将平均值保留两位小数
  .withColumn("Average_2", round(col("Average_2"), 2))

Spark 2.x 版本

由于2.x不支持transform，可直接对数组进行类型转换（Spark 2.3+支持数组类型转换）：

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.DoubleType

val resultDF = originalDF
  .withColumn("_2_array", split(col("_2"), ",").cast(ArrayType(DoubleType)))
  .withColumn("Count_2", size(col("_2_array")))
  .withColumn("Average_2", 
    when(col("Count_2") === 0, lit(null))
    .otherwise(array_sum(col("_2_array")) / col("Count_2"))
  )
  .select("_1", "_2", "Average_2", "Count_2")
  .withColumn("Average_2", round(col("Average_2"), 2))

性能优势 #

• 无数据膨胀：避免使用explode拆分行，不会导致数据量倍数增长，大幅降低内存与IO开销
• 无shuffle操作：所有计算均为单条记录的列内转换，无需跨节点数据传输，适配百万级及更大规模数据集
• 内置函数优化：Spark内置函数经过底层深度优化，性能远高于自定义UDF

边界情况处理 #

• 若_2为空字符串，split会生成包含空字符串的数组，cast为DoubleType后会得到null，array_sum会自动忽略null；可根据业务需求调整为视为空数组（计数为0，平均值为null）
• 若_2包含非数值字符串，可先过滤无效元素：

// 过滤数组中的非数值元素
.withColumn("_2_num_array", transform(col("_2_array"), x => when(x.cast(DoubleType).isNotNull, x.cast(DoubleType))))
.withColumn("_2_num_array", filter(col("_2_num_array"), x => x.isNotNull))

内容的提问来源于stack exchange，提问作者ADK