scala：从变量列列表中获取与最大列值对应的列名

Posted 2023-04-15

【中文标题】scala：从变量列列表中获取与最大列值对应的列名

【英文标题】：scala: get column name corresponding to max column value from variable columns list

【发布时间】：2020-06-08 21:17:24

【问题描述】：

我在 databricks 笔记本中有以下工作解决方案作为测试。

var maxcol =  udf((col1: Long, col2: Long, col3: Long) => 
var res = ""
  if (col1 > col2 && col1 > col3) res = "col1"
  else if (col2 > col1 && col2 > col3) res = "col2"
  else res = "col3"
  res
)

val someDF = Seq(
  (8, 10, 12, "bat"),
  (64, 61, 59, "mouse"),
  (-27, -30, -15, "horse")
).toDF("number1", "number2", "number3", "word")
.withColumn("maxColVal", greatest("number1", "number2", "number3"))
.withColumn("maxColVal_Name", maxcol(col("number1"), col("number2"), col("number3")))

display(someDF)

有什么方法可以使这个通用吗？我有一个用例来使变量列传递给这个 UDF，并且仍然获得最大列名作为与具有最大值的列相对应的输出。 与上面我在 UDF 中硬编码列名 'col1'、'col2' 和 'col3' 的情况不同。

【问题讨论】：

你用的是什么版本？ 2.4？

如果 col1 == col2 和 col1 > col3，你不想要 col3 怎么办？还有一些其他类似的案例

是的，2.4；是的，没错。

那么当两个或更多列具有相同和最大值时，您想要什么？ col1 == col2 = 5 and col3 = 4，你的逻辑返回 4。

【参考方案1】：

在下面使用：

    val df = List((1,2,3,5,"a"),(4,2,3,1,"a"),(1,20,3,1,"a"),(1,22,22,2,"a")).toDF("mycol1","mycol2","mycol3","mycol4","mycol5")

//list all your columns among which you want to find the max value
    val colGroup = List(df("mycol1"),df("mycol2"),df("mycol3"),df("mycol4"))

//list column value -> column name of the columns among which you want to find max value column NAME
    val colGroupMap = List(df("mycol1"),lit("mycol1"),
df("mycol2"),lit("mycol2"),
df("mycol3"),lit("mycol3"),
df("mycol4"),lit("mycol4"))

    var maxcol =  udf((colVal: Map[Int,String]) => 
      colVal.max._2  //you can easily find the column name of the max column value
    )

    df.withColumn("maxColValue",greatest(colGroup:_*)).withColumn("maxColVal_Name",maxcol(map(colGroupMap:_*))).show(false)

    +------+------+------+------+------+-----------+--------------+
    |mycol1|mycol2|mycol3|mycol4|mycol5|maxColValue|maxColVal_Name|
    +------+------+------+------+------+-----------+--------------+
    |1     |2     |3     |5     |a     |5          |mycol4        |
    |4     |2     |3     |1     |a     |4          |mycol1        |
    |1     |20    |3     |1     |a     |20         |mycol2        |
    |1     |22    |22    |2     |a     |22         |mycol3        |
    +------+------+------+------+------+-----------+--------------+