是否有可能对Spark UDAF进行单元测试?

时间:2017-04-30 05:24:58

标签: scala unit-testing apache-spark apache-spark-sql user-defined-functions

Spark UDAF要求您实现多种方法,特别是 def update(buffer: MutableAggregationBuffer, input: Row): Unitdef merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit

假设我的测试中有一个UDAF X,4行(r0, r1, r2, r3)和两个聚合缓冲区A, B。 我想看到这段代码产生了预期的结果:

X.update(A, r0)
X.update(A, r1)
X.update(B, r2)
X.update(B, r3)
X.merge(A, B)
X.evaluate(A)

与仅使用一个缓冲区在4行中的每一行上调用X.update相同:

X.update(A, r0)
X.update(A, r1)
X.update(A, r2)
X.update(A, r3)
X.evaluate(A)

这样就测试了两种方法的正确性。 但是,我不知道如何编写这样的测试:用户代码似乎不能实例化MutableAggregationBuffer的任何实现。

如果我只是从我的4行中创建一个DF,并尝试使用groupBy().agg(...)来调用我的UDAF,Spark甚至不会尝试以这种特定的方式合并它们 - 因为它是少量的行,它不需要。

1 个答案:

答案 0 :(得分:4)

MutableAggregationBuffer只是一个抽象类。您可以轻松创建自己的实现,例如:

import org.apache.spark.sql.expressions._

class DummyBuffer(init: Array[Any]) extends MutableAggregationBuffer {
  val values: Array[Any] = init
  def update(i: Int, value: Any) = values(i) = value
  def get(i: Int): Any = values(i)
  def length: Int = init.size
  def copy() = new DummyBuffer(values)
}

它不会取代真实的东西"但对于简单的测试场景应该足够了。