scala 如何定义自定义聚合函数来对一列向量求和？
声明：本页面是StackOverFlow热门问题的中英对照翻译，遵循CC BY-SA 4.0协议，如果您需要使用它，必须同样遵循CC BY-SA许可，注明原文地址和作者信息，同时你必须将它归于原作者（不是我）：StackOverFlow
原文地址: http://stackoverflow.com/questions/33899977/
Warning: these are provided under cc-by-sa 4.0 license. You are free to use/share it, But you must attribute it to the original authors (not me):
StackOverFlow
How to define a custom aggregation function to sum a column of Vectors?
提问by Rami
I have a DataFrame of two columns, IDof type Intand Vecof type Vector(org.apache.spark.mllib.linalg.Vector).
我有一个两列的 DataFrame，ID类型Int和Vec类型Vector( org.apache.spark.mllib.linalg.Vector)。
The DataFrame looks like follow:
DataFrame 如下所示：
ID,Vec
1,[0,0,5]
1,[4,0,1]
1,[1,2,1]
2,[7,5,0]
2,[3,3,4]
3,[0,8,1]
3,[0,0,1]
3,[7,7,7]
....
I would like to do a groupBy($"ID")then apply an aggregation on the rows inside each group by summing the vectors.
我想做一个groupBy($"ID")然后通过对向量求和对每个组内的行应用聚合。
The desired output of the above example would be:
上面例子的期望输出是：
ID,SumOfVectors
1,[5,2,7]
2,[10,8,4]
3,[7,15,9]
...
The available aggregation functions will not work, e.g. df.groupBy($"ID").agg(sum($"Vec")will lead to an ClassCastException.
可用的聚合函数将不起作用，例如df.groupBy($"ID").agg(sum($"Vec")会导致 ClassCastException。
How to implement a custom aggregation function that allows me to do the sum of vectors or arrays or any other custom operation?
如何实现允许我对向量或数组求和或任何其他自定义操作的自定义聚合函数？
回答by zero323
Spark >= 3.0
火花 >= 3.0
You can use Summarizerwith sum
你可以用Summarizer与sum
import org.apache.spark.ml.stat.Summarizer
df
.groupBy($"id")
.agg(Summarizer.sum($"vec").alias("vec"))
Spark <= 3.0
火花 <= 3.0
Personally I wouldn't bother with UDAFs. There are more than verbose and not exactly fast (Spark UDAF with ArrayType as bufferSchema performance issues) Instead I would simply use reduceByKey/ foldByKey:
就我个人而言，我不会打扰 UDAF。不仅冗长而且不完全（Spark UDAF 与 ArrayType 作为 bufferSchema 性能问题）相反，我会简单地使用reduceByKey/ foldByKey：
import org.apache.spark.sql.Row
import breeze.linalg.{DenseVector => BDV}
import org.apache.spark.ml.linalg.{Vector, Vectors}
def dv(values: Double*): Vector = Vectors.dense(values.toArray)
val df = spark.createDataFrame(Seq(
(1, dv(0,0,5)), (1, dv(4,0,1)), (1, dv(1,2,1)),
(2, dv(7,5,0)), (2, dv(3,3,4)),
(3, dv(0,8,1)), (3, dv(0,0,1)), (3, dv(7,7,7)))
).toDF("id", "vec")
val aggregated = df
.rdd
.map{ case Row(k: Int, v: Vector) => (k, BDV(v.toDense.values)) }
.foldByKey(BDV.zeros[Double](3))(_ += _)
.mapValues(v => Vectors.dense(v.toArray))
.toDF("id", "vec")
aggregated.show
// +---+--------------+
// | id| vec|
// +---+--------------+
// | 1| [5.0,2.0,7.0]|
// | 2|[10.0,8.0,4.0]|
// | 3|[7.0,15.0,9.0]|
// +---+--------------+
And just for comparison a "simple" UDAF. Required imports:
只是为了比较一个“简单的”UDAF。所需的进口：
import org.apache.spark.sql.expressions.{MutableAggregationBuffer,
UserDefinedAggregateFunction}
import org.apache.spark.ml.linalg.{Vector, Vectors, SQLDataTypes}
import org.apache.spark.sql.types.{StructType, ArrayType, DoubleType}
import org.apache.spark.sql.Row
import scala.collection.mutable.WrappedArray
Class definition:
类定义：
class VectorSum (n: Int) extends UserDefinedAggregateFunction {
def inputSchema = new StructType().add("v", SQLDataTypes.VectorType)
def bufferSchema = new StructType().add("buff", ArrayType(DoubleType))
def dataType = SQLDataTypes.VectorType
def deterministic = true
def initialize(buffer: MutableAggregationBuffer) = {
buffer.update(0, Array.fill(n)(0.0))
}
def update(buffer: MutableAggregationBuffer, input: Row) = {
if (!input.isNullAt(0)) {
val buff = buffer.getAs[WrappedArray[Double]](0)
val v = input.getAs[Vector](0).toSparse
for (i <- v.indices) {
buff(i) += v(i)
}
buffer.update(0, buff)
}
}
def merge(buffer1: MutableAggregationBuffer, buffer2: Row) = {
val buff1 = buffer1.getAs[WrappedArray[Double]](0)
val buff2 = buffer2.getAs[WrappedArray[Double]](0)
for ((x, i) <- buff2.zipWithIndex) {
buff1(i) += x
}
buffer1.update(0, buff1)
}
def evaluate(buffer: Row) = Vectors.dense(
buffer.getAs[Seq[Double]](0).toArray)
}
And an example usage:
以及一个示例用法：
df.groupBy($"id").agg(new VectorSum(3)($"vec") alias "vec").show
// +---+--------------+
// | id| vec|
// +---+--------------+
// | 1| [5.0,2.0,7.0]|
// | 2|[10.0,8.0,4.0]|
// | 3|[7.0,15.0,9.0]|
// +---+--------------+
See also: How to find mean of grouped Vector columns in Spark SQL?.
回答by Aviad Klein
I suggest the following (works on Spark 2.0.2 onward), it might be optimized but it's very nice, one thing you have to know in advance is the vector size when you create the UDAF instance
我建议以下（适用于 Spark 2.0.2 以上），它可能会被优化，但它非常好，你必须提前知道的一件事是创建 UDAF 实例时的向量大小
import org.apache.spark.ml.linalg._
import org.apache.spark.mllib.linalg.WeightedSparseVector
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._
class VectorAggregate(val numFeatures: Int)
extends UserDefinedAggregateFunction {
private type B = Map[Int, Double]
def inputSchema: StructType = StructType(StructField("vec", new VectorUDT()) :: Nil)
def bufferSchema: StructType =
StructType(StructField("agg", MapType(IntegerType, DoubleType)) :: Nil)
def initialize(buffer: MutableAggregationBuffer): Unit =
buffer.update(0, Map.empty[Int, Double])
def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
val zero = buffer.getAs[B](0)
input match {
case Row(DenseVector(values)) => buffer.update(0, values.zipWithIndex.foldLeft(zero){case (acc,(v,i)) => acc.updated(i, v + acc.getOrElse(i,0d))})
case Row(SparseVector(_, indices, values)) => buffer.update(0, values.zip(indices).foldLeft(zero){case (acc,(v,i)) => acc.updated(i, v + acc.getOrElse(i,0d))}) }}
def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
val zero = buffer1.getAs[B](0)
buffer1.update(0, buffer2.getAs[B](0).foldLeft(zero){case (acc,(i,v)) => acc.updated(i, v + acc.getOrElse(i,0d))})}
def deterministic: Boolean = true
def evaluate(buffer: Row): Any = {
val Row(agg: B) = buffer
val indices = agg.keys.toArray.sorted
Vectors.sparse(numFeatures,indices,indices.map(agg)).compressed
}
def dataType: DataType = new VectorUDT()
}
