基于其他列将值映射到特定列

Question

我有以下问题：

• 包含带有字符串 A、B 或 C 的 col1 的数据帧。
• 带有整数的第二个 col2。
• 以及其他三列 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。

如何实现此目的？

< code>df.withColumn()我不能使用(或者至少我不知道为什么),对于< code>val df2 = df.map(x =

提前期待您的帮助和感谢！

最好的，肯

Answer 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|
+----+----+----+----+----+