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

xfb7svmp  于 2021-05-29  发布在  Spark
关注(0)|答案(1)|浏览(303)

我在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,并且仍然得到max列名作为输出,对应于具有max值的列。与上面不同的是,我硬编码了udf中的列名“col1”、“col2”和“col3”。

ezykj2lf

ezykj2lf1#

使用方法如下:

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

相关问题