使用ArrayType作为bufferSchema性能的Spark UDAF [英] Spark UDAF with ArrayType as bufferSchema performance issues

问题描述

我正在研究返回一个元素数组的UDAF.

I'm working on a UDAF that returns an array of elements.

每次更新的输入都是索引和值的元组.

The input for each update is a tuple of index and value.

UDAF的作用是对同一索引下的所有值求和.

What the UDAF does is to sum all the values under the same index.

示例:

对于输入(索引，值):(2,1)，(3,1)，(2,3)

For input(index,value) : (2,1), (3,1), (2,3)

应返回(0,0,4,1，...，0)

should return (0,0,4,1,...,0)

逻辑工作正常，但是我遇到了 update方法的问题，我的实现仅为每行更新1个单元格，但是该方法中的最后一个分配实际上复制整个阵列-这是多余的，而且非常耗时.

The logic works fine, but I have an issue with the update method, my implementation only updates 1 cell for each row, but the last assignment in that method actually copies the entire array - which is redundant and extremely time consuming.

仅此一项任务即可完成我的查询执行时间的 98％.

This assignment alone is responsible for 98% of my query execution time.

我的问题是，我该如何减少时间?是否可以在缓冲区数组中分配1个值而不必替换整个缓冲区?

My question is, how can I reduce that time? Is it possible to assign 1 value in the buffer array without having to replace the entire buffer?

P.S .:我正在使用Spark 1.6，并且无法在短期内对其进行升级，因此请坚持使用适用于此版本的解决方案.

P.S.: I'm working with Spark 1.6, and I cannot upgrade it anytime soon, so please stick to solution that would work with this version.

class SumArrayAtIndexUDAF() extends UserDefinedAggregateFunction{

  val bucketSize = 1000

  def inputSchema: StructType =  StructType(StructField("index",LongType) :: StructField("value",LongType) :: Nil)

  def dataType: DataType = ArrayType(LongType)

  def deterministic: Boolean = true

  def bufferSchema: StructType = {
    StructType(
      StructField("buckets", ArrayType(LongType)) :: Nil  
    )
  }

  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = new Array[Long](bucketSize)
  }

  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    val index = input.getLong(0)
    val value = input.getLong(1)

    val arr = buffer.getAs[mutable.WrappedArray[Long]](0)

    buffer(0) = arr   // TODO THIS TAKES WAYYYYY TOO LONG - it actually copies the entire array for every call to this method (which essentially updates only 1 cell)
  }

    override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    val arr1 = buffer1.getAs[mutable.WrappedArray[Long]](0)
    val arr2 = buffer2.getAs[mutable.WrappedArray[Long]](0)

    for(i <- arr1.indices){
      arr1.update(i, arr1(i) + arr2(i))
    }

    buffer1(0) = arr1
  }

  override def evaluate(buffer: Row): Any = {
    buffer.getAs[mutable.WrappedArray[Long]](0)
  }
}

推荐答案

TL; DR 要么不使用UDAF要么使用原始类型代替ArrayType.

TL;DR Either don't use UDAF or use primitive types in place of ArrayType.

没有UserDefinedFunction

Without UserDefinedFunction

这两种解决方案都应避免在内部和外部表示之间进行昂贵的处理.

Both solutions should skip expensive juggling between internal and external representation.

使用标准聚合和pivot

Using standard aggregates and pivot

这使用标准的SQL聚合.虽然内部进行了优化，但是随着键的数量和阵列大小的增加，它可能会变得昂贵.

This uses standard SQL aggregations. While optimized internally it might be expensive when number of keys and size of the array grow.

输入:

val df = Seq((1, 2, 1), (1, 3, 1), (1, 2, 3)).toDF("id", "index", "value")

您可以:

import org.apache.spark.sql.functions.{array, coalesce, col, lit}

val nBuckets = 10
@transient val values = array(
  0 until nBuckets map (c => coalesce(col(c.toString), lit(0))): _*
)

df
  .groupBy("id")
  .pivot("index", 0 until nBuckets)
  .sum("value")
  .select($"id", values.alias("values"))

+---+--------------------+                                                      
| id|              values|
+---+--------------------+
|  1|[0, 0, 4, 1, 0, 0...|
+---+--------------------+

将RDD API与combineByKey/aggregateByKey一起使用.

Using RDD API with combineByKey / aggregateByKey.

使用可变缓冲区将旧的byKey聚合简单化.没有花哨的声音，但在广泛的输入范围内应该表现良好.如果您怀疑输入稀疏，则可以考虑使用更有效的中间表示形式，例如可变的Map.

Plain old byKey aggregation with mutable buffer. No bells and whistles but should perform reasonably well with wide range of inputs. If you suspect input to be sparse, you may consider more efficient intermediate representation, like mutable Map.

rdd
  .aggregateByKey(Array.fill(nBuckets)(0L))(
    { case (acc, (index, value)) => { acc(index) += value; acc }},
    (acc1, acc2) => { for (i <- 0 until nBuckets) acc1(i) += acc2(i); acc1}
  ).toDF

+---+--------------------+
| _1|                  _2|
+---+--------------------+
|  1|[0, 0, 4, 1, 0, 0...|
+---+--------------------+

将UserDefinedFunction与原始类型一起使用

Using UserDefinedFunction with primitive types

据我了解的内部原理，性能瓶颈是

As far as I understand the internals, performance bottleneck is ArrayConverter.toCatalystImpl.

每次调用似乎都被调用

It look like it is called for each call MutableAggregationBuffer.update, and in turn allocates new GenericArrayData for each Row.

如果我们将bufferSchema重新定义为:

If we redefine bufferSchema as:

def bufferSchema: StructType = {
  StructType(
    0 to nBuckets map (i => StructField(s"x$i", LongType))
  )
}

update和merge都可以表示为缓冲区中原始值的普通替换.呼叫链将保持相当长的时间，但

both update and merge can be expressed as plain replacements of primitive values in the buffer. Call chain will remain pretty long, but it won't require copies / conversions and crazy allocations. Omitting null checks you'll need something similar to

val index = input.getLong(0)
buffer.update(index, buffer.getLong(index) + input.getLong(1))

和

for(i <- 0 to nBuckets){
  buffer1.update(i, buffer1.getLong(i) + buffer2.getLong(i))
}

分别.

最后evaluate应该使用Row并将其转换为输出Seq:

Finally evaluate should take Row and convert it to output Seq:

 for (i <- 0 to nBuckets)  yield buffer.getLong(i)

请注意，在此实现中，可能的瓶颈是merge.尽管它不会引入任何新的性能问题，但对于 M 存储桶，每个对merge的调用都是 O(M).

Please note that in this implementation a possible bottleneck is merge. While it shouldn't introduce any new performance problems, with M buckets, each call to merge is O(M).

具有 K 个唯一键和 P 分区，在最坏的情况下，每个键，在每个分区上至少发生一次.这有效地增加了merge组件与 O(M * N * K)的共谋性.

With K unique keys, and P partitions it will be called M * K times in the worst case scenario, where each key, occurs at least once on each partition. This effectively increases complicity of the merge component to O(M * N * K).

通常，您对此无能为力.但是，如果您对数据分布做出特定的假设(数据稀疏，密钥分布是统一的)，则可以稍微简化一下操作，然后先随机播放:

In general there is not much you can do about it. However if you make specific assumptions about the data distribution (data is sparse, key distribution is uniform), you can shortcut things a bit, and shuffle first:

df
  .repartition(n, $"key")
  .groupBy($"key")
  .agg(SumArrayAtIndexUDAF($"index", $"value"))

如果满足假设，则应该:

If the assumptions are satisfied it should:

• 通过改组稀疏对，而不是像密集数组一样的Rows，以反常的方式减小改组大小.
• 仅使用更新(每个 O(1))汇总数据，可能仅将其作为索引的子集.

• Counterintuitively reduce shuffle size by shuffling sparse pairs, instead of dense array-like Rows.
• Aggregate data using updates only (each O(1)) possibly touching only as subset of indices.

但是，如果不满足一个或两个假设，则可以预期随机播放的大小会增加，而更新数量将保持不变.同时，数据偏斜会使情况变得比update-shuffle-merge情形还要糟.

However if one or both assumptions are not satisfied, you can expect that shuffle size will increase while number of updates will stay the same. At the same time data skews can make things even worse than in update - shuffle - merge scenario.

使用Aggregator ，并在Dataset中输入"strong":

Using Aggregator with "strongly" typed Dataset:

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
import org.apache.spark.sql.{Encoder, Encoders}

class SumArrayAtIndex[I](f: I => (Int, Long))(bucketSize: Int)  extends Aggregator[I, Array[Long], Seq[Long]]
    with Serializable {
  def zero = Array.fill(bucketSize)(0L)
  def reduce(acc: Array[Long], x: I) = {
    val (i, v) = f(x)
    acc(i) += v
    acc
  }

  def merge(acc1: Array[Long], acc2: Array[Long]) = {
    for {
      i <- 0 until bucketSize
    } acc1(i) += acc2(i)
    acc1
  }

  def finish(acc: Array[Long]) = acc.toSeq

  def bufferEncoder: Encoder[Array[Long]] = Encoders.kryo[Array[Long]]
  def outputEncoder: Encoder[Seq[Long]] = ExpressionEncoder()
}

可以按如下所示使用

val ds = Seq((1, (1, 3L)), (1, (2, 5L)), (1, (0, 1L)), (1, (4, 6L))).toDS

ds
  .groupByKey(_._1)
  .agg(new SumArrayAtIndex[(Int, (Int, Long))](_._2)(10).toColumn)
  .show(false)

+-----+-------------------------------+
|value|SumArrayAtIndex(scala.Tuple2)  |
+-----+-------------------------------+
|1    |[1, 3, 5, 0, 6, 0, 0, 0, 0, 0] |
|2    |[0, 11, 0, 0, 0, 0, 0, 0, 0, 0]|
+-----+-------------------------------+

注意:

另请参见 SPARK-27296 -用户定义的聚合函数(UDAF)有一个主要的效率问题

这篇关于使用ArrayType作为bufferSchema性能的Spark UDAF的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！