如何在spark Dataframe中的nest数组上设计UDAF

Question

输入数据：

+---+----+----+
|idx|  v1|  v2|
+---+----+----+
|  a|   1|   3|
|  a|null|   2|
|  a|   4|   5|
|  b|   6|   1|
|  b|   7|null|
+---+----+----+

我想要的是什么：

+---+-------------------------------------------+                               
|idx|total                                      |
+---+-------------------------------------------+
|b  |[WrappedArray(6, 7), WrappedArray(1)]      |
|a  |[WrappedArray(1, 4), WrappedArray(3, 2, 5)]|
+---+-------------------------------------------+

我知道我可以得到这个

df.groupBy("idx").agg(array(collect_list(col("v1")), collect_list(col("v2"))));

但我希望通过JAVA中的 UDAF 来实现结果。