web-dev-qa-db-fra.com

Comment définir une fonction d'agrégation personnalisée pour additionner une colonne de vecteurs?

J'ai un DataFrame de deux colonnes, ID de type Int et Vec de type Vector (org.Apache.spark.mllib.linalg.Vector).

Le DataFrame ressemble à ceci:

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]
....

Je voudrais faire une groupBy($"ID") puis appliquer une agrégation sur les lignes à l'intérieur de chaque groupe en additionnant les vecteurs.

La sortie souhaitée de l'exemple ci-dessus serait:

ID,SumOfVectors
1,[5,2,7]
2,[10,8,4]
3,[7,15,9]
...

Les fonctions d'agrégation disponibles ne fonctionneront pas, par exemple df.groupBy($"ID").agg(sum($"Vec") conduira à une ClassCastException.

Comment implémenter une fonction d'agrégation personnalisée qui me permet de faire la somme des vecteurs ou des tableaux ou toute autre opération personnalisée?

scalaapache-sparkapache-spark-sqlaggregate-functionsapache-spark-ml
22
Rami

Spark> = 3.0

Vous pouvez utiliser Summarizer avec sum

import org.Apache.spark.ml.stat.Summarizer

df
  .groupBy($"id")
  .agg(Summarizer.sum($"vec").alias("vec"))

Spark <= 3.0

Personnellement, je ne m'embêterais pas avec les UDAF. Il y a plus que verbeux et pas exactement rapide ( Spark UDAF avec ArrayType comme problèmes de performances de bufferSchema ) À la place, j'utiliserais simplement 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]|
// +---+--------------+

Et juste pour comparaison un UDAF "simple". Importations requises:

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

Définition de classe:

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)
}

Et un exemple d'utilisation:

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]|
// +---+--------------+

Voir aussi: Comment trouver la moyenne des colonnes vectorielles groupées dans Spark SQL? .

27
zero323

Je suggère ce qui suit (fonctionne sur Spark 2.0.2 en avant), il pourrait être optimisé mais c'est très bien, une chose que vous devez savoir à l'avance est la taille du vecteur lorsque vous créez l'instance 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()
}
0
Aviad Klein