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

Question

我在 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' 的情况不同。

标签： scaladataframeapache-spark

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