如何在 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")))