Spark DataFrame Aggregation based on two or more Columns - apache-spark

Question

I want to write a UDAF for some customized aggregation based on more than one column. A simple example would be a dataframe with two columns, c1 and c2. For each row, I take the max of c1 and c2 (let's call it cmax), then I take the sum of cmax.
When I call df.agg(), it does not look like I can pass two or more columns to any aggregation method including UDAF. 1st question, is it true?
For this simple example, I could create another column called cmax, and do the aggregation on cmax. But in reality, I would need to do aggregation based on N combinations of columns and the results would be a collection of size N. I would want to loop the combinations within the update method in my UDAF. Therefore it would require N intermediate columns, which does not seem to be a clean solution to me. 2nd question, I wonder if creating intermediate columns is the way to do it, or if there is a better solution.
I noticed in RDD, the problem is much easier. I can pass the entire record to my aggregation function and I have access to all the data fields.

Answer 1

You can use as many columns in a UDAF as the signature of it's apply function accepts multiple Columns (from it's source code).
def apply(exprs: Column*): Column
You just have to make sure that the inputSchema returns a StructType reflecting the columns that you want to consume as your UDAF input.
For the case of columns c1 and c2 your UDAF has to implement a inputSchema with the following schema:
def inputSchema: StructType = StructType(Array(StructField("c1", DoubleType), StructField("c2", DoubleType)))
However if you want a more general solution, you can always initialize the custom UDAF with arguments that allows returning the right inputSchema. See the example below that allows defining an arbitrary StructType at construction time (Note that we don't verify that StructType is of DoubleType).
class MyMaxUDAF(schema: StructType) extends UserDefinedAggregateFunction {
def inputSchema: StructType = this.schema
def bufferSchema: StructType = StructType(Array(StructField("maxSum", DoubleType)))
def dataType: DataType = DoubleType
def deterministic: Boolean = true
def initialize(buffer: MutableAggregationBuffer): Unit = buffer(0) = 0.0
def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
buffer(0) = buffer.getDouble(0) + Array.range(0, input.length).map(input.getDouble).max
}
def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = buffer2 match {
case Row(buffer2Sum: Double) => buffer1(0) = buffer1.getDouble(0) + buffer2Sum
}
def evaluate(buffer: Row): Double = buffer match {
case Row(totalSum: Double) => totalSum
}
}
Your DataFrame containing values and a key for aggregation.
val df = spark.createDataFrame(Seq(
Entry(0, 1.0, 2.0, 3.0), Entry(0, 3.0, 1.0, 2.0), Entry(1, 6.0, 2.0, 2)
))
df.show
+-------+---+---+---+
|groupMe| c1| c2| c3|
+-------+---+---+---+
| 0|1.0|2.0|3.0|
| 0|3.0|1.0|2.0|
| 1|6.0|2.0|2.0|
+-------+---+---+---+
And using the UDAF we expect the sum of max being 6.0 and 6.0
val fields = Array("c1", "c2", "c3")
val struct = StructType(fields.map(StructField(_, DoubleType)))
val myMaxUDAF: MyMaxUDAF = new MyMaxUDAF(struct)
df.groupBy("groupMe").agg(myMaxUDAF(fields.map(df(_)):_*)).show
+-------+---------------------+
|groupMe|mymaxudaf(c1, c2, c3)|
+-------+---------------------+
| 0| 6.0|
| 1| 6.0|
+-------+---------------------+
There is a nice tutorial on UDAF. Unfortunately they don't cover multiple arguments.
https://ragrawal.wordpress.com/2015/11/03/spark-custom-udaf-example/