Spark: grouping rows in array by key

Question

I have a spark dataset like this one:

key id val1 val2 val3
1   a  a1   a2   a3
2   a  a4   a5   a6
3   b  b1   b2   b3
4   b  b4   b5   b6
5   b  b7   b8   b9
6   c  c1   c2   c3

I would like to group all rows by id in a list or array like this:

(a, ([1   a  a1   a2   a3], [2   a  a4   a5   a6]) ),
(b, ([3   b  b1   b2   b3], [4   b  b4   b5   b6], [5   b  b7   b8   b9]) ),
(c, ([6   c  c1   c2   c3]) )

I have used map to output key/value pairs with the right key but I have troubles in building the final key/array.

Can anybody help with that?

Answer 1

how about this:

import org.apache.spark.sql.functions._
df.withColumn("combined",array("key","id","val1","val2","val3")).groupby("id").agg(collect_list($"combined"))

The Array function converts the columns into an array of column and then its a simple groupby with collect_list