Spark中的动态列选择（基于另一列的值）

Question

使用给定的Spark DataFrame：

> df.show()

+---+-----+---+---+---+---+
| id|delay| p1| p2| p3| p4|
+---+-----+---+---+---+---+
|  1|    3|  a|  b|  c|  d|
|  2|    1|  m|  n|  o|  p|
|  3|    2|  q|  r|  s|  t|
+---+-----+---+---+---+---+

如何动态选择一列，以便新的col列是p{delay}现有列的结果？

> df.withColumn("col", /* ??? */).show()

+---+-----+---+---+---+---+----+
| id|delay| p1| p2| p3| p4| col|
+---+-----+---+---+---+---+----+
|  1|    3|  a|  b|  c|  d|   c|   // col = p3
|  2|    1|  m|  n|  o|  p|   m|   // col = p1
|  3|    2|  q|  r|  s|  t|   r|   // col = p2
+---+-----+---+---+---+---+----+

Answer 1

我能想到的最简单的解决方案是使用delay array作为索引：

import org.apache.spark.sql.functions.array

df.withColumn("col", array($"p1", $"p2", $"p3", $"p4")($"delay" - 1))

Answer 2

一种选择是创建一个从数字到列名的映射，然后使用foldLeft用相应的值更新col列：

val cols = (1 to 4).map(i => i -> s"p$i")

(cols.foldLeft(df.withColumn("col", lit(null))){ 
   case (df, (k, v)) => df.withColumn("col", when(df("delay") === k, df(v)).otherwise(df("col"))) 
}).show
+---+-----+---+---+---+---+---+    
| id|delay| p1| p2| p3| p4|col|
+---+-----+---+---+---+---+---+
|  1|    3|  a|  b|  c|  d|  c|
|  2|    1|  m|  n|  o|  p|  m|
|  3|    2|  q|  r|  s|  t|  r|
+---+-----+---+---+---+---+---+