斯卡拉/星火dataframes：找到对应的最大列名

Question

在斯卡拉/星火，有一个数据帧：

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熊猫（是的，我知道，我应该pyspark比较...）的等效可能是：

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

Answer 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")）。通过coalescing

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列的情况下，你必须调整此，例如值-Inf：如果由于某种原因，这是不能接受的，你可以用when明确减少。