Explode multiple nested columns, perform agg and join all the tables - scala

Question

I was wondering if there is another option much more efficient to do this job, for example:
val df0 = df.select($"id", explode($"event.x0") as "n_0" ).groupBy("id").agg(sum("n_0") as "0")
val df1 = df.select($"id", explode($"event.x1") as "n_1").groupBy("id").agg(sum("n_1") as "1")
val df2 = df.select($"id", explode($"event.x2") as "n_2").groupBy("id").agg(sum("n_2") as "2")
val df3 = df.select($"id", explode($"event.x3") as "n_3").groupBy("id").agg(sum("n_3") as "3)
val final_df = df.join(df0, "id").join(df1, "id").join(df2, "id").join(df3, "id")
I was trying something like this:
val df_x = df.select($"id", $"event", explode($"event.x0") as "0" )
.select($"id", $"event", $"0", explode($"event.x1") as "1")
.select($"id", $"event", $"0", $"1", explode($"event.x2") as "2")
.groupBy("id")
.agg(sum("0") as "0", sum("1") as "1", sum("2") as "2")
val final_df = df.join(df_x, "id")
Despite it runs much more faster!!!! The aggregations values are wrong, so it does not work actually :( !
Any ideas to decrease the amount of joins ?

Answer 1

Assuming each id doesn't have too many matching records, you can use the collect_list aggregation function to collect all matching arrays into an array-of-arrays, and then a User Defined Function to sum over these nested arrays:
val flattenAndSum = udf[Int, mutable.Seq[mutable.Seq[Int]]] { seqOfArrays => seqOfArrays.flatten.sum }
val sums = df.groupBy($"id").agg(
collect_list($"event.x0") as "arr0",
collect_list($"event.x1") as "arr1",
collect_list($"event.x2") as "arr2",
collect_list($"event.x3") as "arr3"
).select($"id",
flattenAndSum($"arr0") as "0",
flattenAndSum($"arr1") as "1",
flattenAndSum($"arr2") as "2",
flattenAndSum($"arr3") as "3"
)
df.join(sums, "id")
Alternatively, if that assumption cannot be made, you can create a User Defined Aggregation Function to perform the flatten-and-sum on the fly. This is safer and potentially faster but requires a bit more work:
// implement a UDAF:
class FlattenAndSum extends UserDefinedAggregateFunction {
override def inputSchema: StructType = new StructType().add("arr", ArrayType(IntegerType))
override def bufferSchema: StructType = new StructType().add("sum", IntegerType)
override def dataType: DataType = IntegerType
override def deterministic: Boolean = true
override def initialize(buffer: MutableAggregationBuffer): Unit = buffer.update(0, 0)
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
val current = buffer.getAs[Int](0)
val toAdd = input.getAs[Seq[Int]](0).sum
buffer.update(0, current + toAdd)
}
override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
buffer1.update(0, buffer1.getAs[Int](0) + buffer2.getAs[Int](0))
}
override def evaluate(buffer: Row): Any = buffer.getAs[Int](0)
}
// use it in aggregation:
val flattenAndSum = new FlattenAndSum()
val sums = df.groupBy($"id").agg(
flattenAndSum($"event.x0") as "0",
flattenAndSum($"event.x1") as "1",
flattenAndSum($"event.x2") as "2",
flattenAndSum($"event.x3") as "3"
)
df.join(sums, "id")