Scala / Spark数据帧：找到与max相对应的列名

import org.apache.spark.sql.functions.{col,greatest,lit,struct}

首先让我们创建一个结构列表,其中第一个元素是值,第二个是列名：

val structs = dfIn.columns.tail.map(
  c => struct(col(c).as("v"),lit(c).as("k"))
)

像这样的结构可以传递到最大,如下所示：

dfIn.withColumn("maxCol",greatest(structs: _*).getItem("k"))

+---+---+---+---+------+
| id| c0| c1| c2|maxCol|
+---+---+---+---+------+
| r0|  0|  2|  3|    c2|
| r1|  1|  0|  0|    c0|
| r2|  0|  2|  2|    c2|
+---+---+---+---+------+

请注意,在连接的情况下,它将采用序列中稍后出现的元素(按字典顺序排列(x,“c2”)>(x,“c1”)).如果出于某种原因这是不可接受的,您可以在以下时间明确减少：

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

val max_col = structs.reduce(
  (c1,c2) => when(c1.getItem("v") >= c2.getItem("v"),c1).otherwise(c2)
).getItem("k")

dfIn.withColumn("maxCol",max_col)

+---+---+---+---+------+
| id| c0| c1| c2|maxCol|
+---+---+---+---+------+
| r0|  0|  2|  3|    c2|
| r1|  1|  0|  0|    c0|
| r2|  0|  2|  2|    c1|
+---+---+---+---+------+

如果是可空列,则必须对其进行调整,例如通过合并到-Inf的值.