基于另一列将值映射到特定列

Posted 2023-04-13

【中文标题】基于另一列将值映射到特定列

【英文标题】：Mapping a value into a specific column based on annother column

【发布时间】：2017-03-01 14:48:39

【问题描述】：

我有以下问题：

一个DataFrame 包含col1 和字符串A、B 或C。 第二个col2 和Integer。 还有另外三列col3、col4 和col5（这些列也被命名为A、B 和C）。

因此，

 col1 - col2 - A (col3) - B (col4) - C (col5)
|--------------------------------------------
   A      6
   B      5
   C      6

应该得到

 col1 - col2 - A (col3) - B (col4) - C (col5)
|--------------------------------------------
   A      6       6
   B      5                  5
   C      6                              6

现在我想遍历每一行并根据col1 中的条目将col2 中的整数分配给A、B 或C 列。

我如何做到这一点？

df.withColumn() 我不能使用（或者至少我不知道为什么），val df2 = df.map(x => x ) 也是如此。

期待您的帮助，提前致谢！

最好的，肯

【参考方案1】：

在键和目标列之间创建映射：

val mapping = Seq(("A", "col3"), ("B", "col4"), ("C", "col5"))

用它来生成列序列：

import org.apache.spark.sql.functions.when

val exprs = mapping.map  case (key, target) => 
  when($"col1" === key, $"col2").alias(target) 

添加星号并选择：

val df = Seq(("A", 6), ("B", 5), ("C", 6)).toDF("col1", "col2")
df.select($"*" +: exprs: _*)

结果是：

+----+----+----+----+----+
|col1|col2|col3|col4|col5|
+----+----+----+----+----+
|   A|   6|   6|null|null|
|   B|   5|null|   5|null|
|   C|   6|null|null|   6|
+----+----+----+----+----+