使用groupBy在DataFrame中激发聚合SparseVectors

Question

我有以下包含GET和id: Int的DataFrame

stuff: breeze.linalg.SparseVector

我想以预期的输出运行以下代码

val test = spark.createDataFrame(Seq(
  (1, Vectors.sparse(5, Array(0, 1), Array(1.0, 2.0))), 
  (1, Vectors.sparse(5, Array(1, 2), Array(1.0, 2.0))), 
  (2, Vectors.sparse(5, Array(2, 3), Array(2.0, 3.0)))
)).toDF("id", "stuff")

test.show
+---+-------------------+
| id|              stuff|
+---+-------------------+
|  1|(5,[0,1],[1.0,2.0])|
|  1|(5,[1,2],[1.0,2.0])|
|  2|(5,[2,3],[2.0,3.0])|
+---+-------------------+

我创建了一个简单的UDF（非常丑陋，但是确实可以完成工作），但是它只能在行级别上使用（不适用于groupBy / agg）：

test.groupBy("id").agg(mergeVectors($"stuff"))
+---+-------------------------+
| id|                    stuff|
+---+-------------------------+
|  1|(5,[0,1,2],[1.0,3.0,2.0])|
|  2|(5,[2,3],      [2.0,3.0])|
+---+-------------------------+