У меня есть DataFrame из двух столбцов, ID
типа Int
и Vec
типа Vector
(org.apache.spark.mllib.linalg.Vector
).
DataFrame выглядит следующим образом:
ID,Vec
1,[0,0,5]
1,[4,0,1]
1,[1,2,1]
2,[7,5,0]
2,[3,3,4]
3,[0,8,1]
3,[0,0,1]
3,[7,7,7]
....
Я хотел бы сделать groupBy($"ID")
, затем применить агрегацию по строкам внутри каждой группы, суммируя векторы.
Желаемый результат приведенного выше примера будет:
ID,SumOfVectors
1,[5,2,7]
2,[10,8,4]
3,[7,15,9]
...
Доступные функции агрегации не будут работать, например. df.groupBy($"ID").agg(sum($"Vec")
приведет к исключению ClassCastException.
Как реализовать пользовательскую функцию агрегации, которая позволяет мне делать сумму векторов или массивов или любую другую пользовательскую операцию?