Scala/Spark 数据帧：找到对应于最大值的列名答案

【问题标题】：Scala/Spark dataframes: find the column name corresponding to the maxScala/Spark 数据帧：找到对应于最大值的列名
【发布时间】：2017-07-18 00:40:59
【问题描述】：

在 Scala/Spark 中，有一个数据框：

val dfIn = sqlContext.createDataFrame(Seq(
  ("r0", 0, 2, 3),
  ("r1", 1, 0, 0),
  ("r2", 0, 2, 2))).toDF("id", "c0", "c1", "c2")

我想计算一个新列maxCol，其中包含对应于最大值（每行）的列的名称。在这个例子中，输出应该是：

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

实际上数据框有 60 多列。因此需要一个通用的解决方案。

Python Pandas 中的等价物（是的，我知道，我应该与 pyspark 进行比较...）可能是：

dfOut = pd.concat([dfIn, dfIn.idxmax(axis=1).rename('maxCol')], axis=1)

【问题讨论】：

你一般有多少列？
我有大约 60 列
与最大列相比有多少？
不知道是不是这个question的副本。显然，拥有 60 列会改变一组可行的解决方案。反正有下面answer
谢谢@Wilmerton！这是一个很好的解决方案，展示了 Scala/Spark 数据帧与 Python/Pandas 数据帧的优雅；）（同样，pandas 数据帧没有分布，因此比较并不真正相关）

标签： scala apache-spark dataframe apache-spark-sql argmax

【解决方案1】：

通过一个小技巧，您可以使用greatest 函数。所需的进口：

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

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

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

这样的结构可以传递给greatest，如下所示：

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")）。如果由于某种原因这是不可接受的，您可以使用when 显式减少：

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

如果是 nullable 列，您必须对此进行调整，例如将 coalescing 的值调整为 -Inf。

【讨论】：