dataframe - 使用 Scala 根据 Spark Dataframe 中另一列的值选择一列

Question

我有一个包含 5 列的数据框 - sourceId , score_1 , score_3 , score_4和score_7 。 sourceId 的值列可以是[1, 3, 4, 7] 。我想将其转换为另一个具有列 sourceId 的数据框和score ，其中分数取决于 sourceId 的值专栏。

<表类=“s-表”> <标题> 源ID 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

所以如果 sourceId = 1 ，我们选择值score_1对于该记录，如果 sourceId = 3 ，我们选择值score_3 ，等等...

结果是

<表类=“s-表”> <标题> 源ID 得分 <正文> 1 0.3 4 0.73 7 0.16 3 0.1

在 Spark 中执行此操作的最佳方法是什么？

Answer 1

在 id 列值上链接多个 when 表达式:

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)

---

或者从 score_* 列构建 map 表达式并使用它来获取 score 值:

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")))