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

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

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

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