Is it possible to unit test Spark UDAFs?

Question

Spark UDAFs require that you implement several methods, specifically `def update(buffer: MutableAggregationBuffer, input: Row): Unit` and `def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit` Suppose I have a UDAF X, 4 rows `(r0, r1, r2, r3)` and two aggregation buffers `A, B` in my test. I want to see that this code produces the intended result: ``` X.update(A, r0) X.update(A, r1) X.update(B, r2) X.update(B, r3) X.merge(A, B) X.evaluate(A) ``` Same as calling X.update on each of the 4 rows with just one buffer: ``` X.update(A, r0) X.update(A, r1) X.update(A, r2) X.update(A, r3) X.evaluate(A) ``` This way the correctness of both methods is tested. However, I don't know how to write such a test: it does not seem that user code can instantiate any implementation of `MutableAggregationBuffer`. If I simply make a DF out of my 4 rows, and try to use `groupBy().agg(...)` to call my UDAF, Spark won't even try to merge them in this specific manner - since it's a small number of r…

Accepted Answer

`MutableAggregationBuffer` is just an abstract class. You can easily create your own implementation, for example one like this:

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

It won't replace a "real thing" but should be just enough for simple testing scenarios.

Is it possible to unit test Spark UDAFs?

Solution

Problem