How can I aggregate a Spark dataframe to get a sparse vector using Scala?

I have a dataframe like the one below in Spark and I want to group it by a column id

and then for each row in the grouped data I need to create a sparse vector with elements from the weight

column with the indices specified by the column index

. The length of the sparse vector is known, say 1000 for this example.

Dataframe df

:

+-----+------+-----+
|   id|weight|index|
+-----+------+-----+
|11830|     1|    8|
|11113|     1|    3|
| 1081|     1|    3|
| 2654|     1|    3|
|10633|     1|    3|
|11830|     1|   28|
|11351|     1|   12|
| 2737|     1|   26|
|11113|     3|    2|
| 6590|     1|    2|
+-----+------+-----+

      

        
        
        
      

    

I read this which is similar to what I want to do but for rdd. Does anyone know a good way to do this for a dataframe in Spark using Scala?

My attempt so far is to first collect the weights and indices as lists like this:

val dfWithLists = df
    .groupBy("id")
    .agg(collect_list("weight") as "weights", collect_list("index") as "indices"))

      

        
        
        
      

    

which looks like this:

+-----+---------+----------+
|   id|  weights|   indices|
+-----+---------+----------+
|11830|   [1, 1]|   [8, 28]|
|11113|   [1, 3]|    [3, 2]|
| 1081|      [1]|       [3]|
| 2654|      [1]|       [3]|
|10633|      [1]|       [3]|
|11351|      [1]|      [12]|
| 2737|      [1]|      [26]|
| 6590|      [1]|       [2]|
+-----+---------+----------+

      

        
        
        
      

    

Then I define udf and do something like this:

import org.apache.spark.mllib.linalg.{Vector, Vectors}
import org.apache.spark.sql.functions.udf

def toSparseVector: ((Array[Int], Array[BigInt]) => Vector) = {(a1, a2) => Vectors.sparse(1000, a1, a2.map(x => x.toDouble))}
val udfToSparseVector = udf(toSparseVector)

val dfWithSparseVector = dfWithLists.withColumn("SparseVector", udfToSparseVector($"indices", $"weights"))

      

        
        
        
      

    

but that doesn't seem to work and it seems like there should be an easier way to do it without having to collect the weights and indices into lists first.

I am new to Spark, Dataframes and Scala so any help is greatly appreciated.