TL;DR 要么不使用UDAF,要么在ArrayType
的位置上使用原始类型。
不使用UserDefinedFunction
这两种解决方案都应该避免在内部和外部表示之间进行昂贵的切换。
使用标准聚合和pivot
这使用标准SQL聚合。尽管在内部进行了优化,但当键的数量和数组的大小增长时,它可能会很昂贵。
给定输入:
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...|
+
使用 combineByKey
/ aggregateByKey
的 RDD API。
使用可变缓冲区的普通的 byKey
聚合。没有特殊花哨的功能,但适用于各种输入并表现良好。如果您怀疑输入为稀疏数据,则可以考虑使用更高效的中间表示,如可变的 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
据我所知,性能瓶颈在于ArrayConverter.toCatalystImpl
。
看起来它是在每个调用MutableAggregationBuffer.update
时被调用的,而且每次为Row
分配新的GenericArrayData
。
如果我们将bufferSchema
重新定义为:
def bufferSchema: StructType = {
StructType(
0 to nBuckets map (i => StructField(s"x$i", LongType))
)
}
无论是update
还是merge
,都可以被表达为在缓冲区中对原始值进行简单的替换。调用链会保持相当长,但它不需要复制/转换和疯狂的分配。省略null
检查,您需要类似于:
val index = input.getLong(0)
buffer.update(index, buffer.getLong(index) + input.getLong(1))
and
for(i <- 0 to nBuckets){
buffer1.update(i, buffer1.getLong(i) + buffer2.getLong(i))
}
分别。最后,
evaluate
应该接受
Row
并将其转换为输出的
Seq
:
for (i <- 0 to nBuckets) yield buffer.getLong(i)
请注意,在这个实现中,一个可能成为瓶颈的地方是
merge
。虽然它不应该引入任何新的性能问题,但是对于
M 个桶,每次调用
merge
的时间复杂度为
O(M)。
在
K个唯一键和
P个分区的情况下,最坏情况下将被调用
M * K次,其中每个键在每个分区中至少出现一次。这将有效地将
merge
组件的复杂度增加到
O(M * N * K)。
一般来说,你无法做太多关于它。但是,如果你对数据分布做出具体的假设(数据是稀疏的,键分布是均匀的),你可以简化一些东西,并首先进行洗牌:
df
.repartition(n, $"key")
.groupBy($"key")
.agg(SumArrayAtIndexUDAF($"index", $"value"))
如果假设得到满足,它应该会:
- 通过洗牌稀疏对而不是密集的类似数组
Rows
来出人意料地减少洗牌大小。
- 仅使用更新的数据聚合(每个O(1)),可能只涉及子集索引。
但是,如果一个或两个假设未得到满足,你可以期望洗牌大小会增加,而更新数量会保持不变。同时,数据偏斜可能会比update
- shuffle
- merge
方案更糟糕。
在带有“强类型”Dataset
的情况下使用Aggregator
:
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)存在严重的效率问题