Select a column based on another column's value in Spark Dataframe using Scala

user_17432580 0
Viewed 846

I have a dataframe with 5 columns - sourceId, score_1, score_3, score_4 and score_7. The values of sourceId column can be [1, 3, 4, 7]. I want to convert this into another dataframe that has the columns sourceId and score, where score depends on the value of the sourceId column.

sourceId score_1 score_3 score_4 score_7
1 0.3 0.7 0.45 0.21
4 0.15 0.66 0.73 0.47
7 0.34 0.41 0.78 0.16
3 0.77 0.1 0.93 0.67

So if sourceId = 1, we select value of score_1 for that record, if sourceId = 3, we select value of score_3, and so on...

Result would be

sourceId score
1 0.3
4 0.73
7 0.16
3 0.1

What would be the best way to do this in Spark?

apache-spark-sql apache-spark dataframe scala
1 Answers

Chaining multiple when expressions on id column values:

val ids = Seq(1, 3, 4, 7)

val scoreCol = ids.foldLeft(lit(null)) { case (acc, id) =>
  when(col("sourceId")===id, col(s"score_$id")).otherwise(acc)
}

val df2 = df.withColumn("score", scoreCol)

Or building a map expression from score_* columns and use it to get score values:

val scoreMap = map(
  df.columns
    .filter(_.startsWith("score_"))
    .flatMap(c => Seq(lit(c.split("_")(1)), col(c))): _*
)

val df2 = df.withColumn("score", scoreMap(col("sourceId")))
user_13865510
Related