如何在 Spark sql withColumn 中動態構建列名

Question

我有如下示例數據。

我想根據country /地區字段值填充另一列exp 。

像下面的東西

df.withColumn("exp",col(s"exp_$country"))

這樣相應的country號碼就可以放在那里。

但上面的代碼錯誤說：

無法解析國家

Output 我需要的是

任何幫助表示贊賞。

Answer 1

您可以從國家列表中鏈接多個when表達式：

val countries = Seq("us", "uk", "ind")

val expCol = countries.foldLeft(lit(null)) { case (acc, country) =>
  when(col("country")===country, col(s"exp_$country")).otherwise(acc)
}

val df1 = df.withColumn("exp", expCol)

---

或者，如果您更喜歡從exp_*列創建 map 表達式country -> exp ，而不是使用 map 創建exp列：

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

val df1 = df.withColumn("exp", mapCountries(col("country")))