Get column from another map type column spark scala dataframe

Question

I have a dataframe having 1 column as a case class format like this

case class FeaturizedDataset(
    indices: Array[String],
    values: Array[Float]
)

The table is something like this

|sourceId|scoreMapping           |
|--------|-----------------------|
|3       |{[1,3,4],[0.1,0.2,0.3]}|
|4       |{[1,3,4],[0.1,0.2,0.3]}|
|1       |{[1,3,4],[0.1,0.2,0.3]}|
|4       |{[1,3,4],[0.1,0.2,0.3]}|

I would like to get the score corresponding the score ID in scala. How can I do this?

For e.g., if sourceId = 3 => score = 0.2, if sourceId = 1 => score =0.1 ...

Answer 1

Using arrays_zip + filter functions:

val df1 = df.withColumn(
  "tmp",
  arrays_zip(col("scoreMapping.indices"), col("scoreMapping.values"))
).withColumn(
  "score",
  filter(col("tmp"), x => x("0") === col("sourceId"))(0)("1")
).drop("tmp")

You zip scoreMapping indices and values then filtering where indices is equal to sourceId.