在Spark/Scala中计算DataFrame内SparkVector元素的求和/均值及稀疏向量求和方法

我来帮你解决这两个Spark/Scala里处理Vector的问题，都是实际项目里常遇到的场景：

Spark的DataFrame本身没有内置的Vector聚合函数，我们需要结合MLlib的Vector工具类和自定义UDF来实现。这里分两种常见场景说明：

场景1：计算每个Vector自身的元素总和与平均值 #

针对每行的Vector单独计算，我们可以基于org.apache.spark.ml.linalg.Vector的API写自定义UDF：

import org.apache.spark.sql.functions.{udf, col}
import org.apache.spark.ml.linalg.Vector

// 计算单个Vector的元素总和
val vectorSumUdf = udf((vec: Vector) => vec.toArray.sum)

// 计算单个Vector的元素平均值
val vectorAvgUdf = udf((vec: Vector) => vec.toArray.sum / vec.size)

// 使用示例：假设你的DataFrame目标列名为`features`
val dfWithSumAvg = originalDF
  .withColumn("vector_sum", vectorSumUdf(col("features")))
  .withColumn("vector_avg", vectorAvgUdf(col("features")))

场景2：计算整个列所有Vector的元素全局总和与平均值 #

如果要统计所有行Vector的元素总和再算全局平均值，先对每行Vector求和，再做全局聚合：

// 先计算每行Vector的总和，再聚合得到全局总和
val globalSum = originalDF
  .select(vectorSumUdf(col("features")).alias("row_sum"))
  .agg(sum("row_sum"))
  .first()
  .getDouble(0)

// 计算全局平均值：总元素数=行数×单个Vector长度（假设所有Vector长度一致）
val sampleVectorSize = originalDF.select(col("features").size).first().getInt(0)
val totalElements = originalDF.count() * sampleVectorSize
val globalAvg = globalSum / totalElements

注意：如果你的Vector长度不一致，需要额外写UDF统计每个Vector的size，再聚合得到总元素数。

绝对不要把SparseVector转换成DenseVector——这会完全浪费稀疏向量的内存优势。最优方法是直接利用SparseVector仅存储非零元素的特性，操作它的values数组：

需求1：对每行的SparseVector自身元素求和 #

写一个针对SparseVector的轻量UDF，直接求和非零元素数组：

import org.apache.spark.sql.functions.udf
import org.apache.spark.ml.linalg.SparseVector

val sparseSumUdf = udf((sv: SparseVector) => sv.values.sum)

// 使用示例：假设你的PPR结果DataFrame名为pprDF
val pprWithRowSum = pprDF.withColumn("pagerank_row_sum", sparseSumUdf(col("pageranks")))

这个方法效率极高，只遍历实际存储的非零元素，完全跳过默认的零值，适合大规模稀疏数据场景。

需求2：对整个列所有SparseVector的元素全局求和 #

先计算每行的Vector和，再做全局聚合：

val globalPagerankSum = pprDF
  .select(sparseSumUdf(col("pageranks")).alias("row_sum"))
  .agg(sum("row_sum"))
  .first()
  .getDouble(0)

额外提示：如果你的SparseVector规模极大，担心UDF的序列化开销，可以用mapPartitions批量处理分区内的Vector，但大多数场景下，上面的UDF方法已经足够高效，代码也更简洁易维护。

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