CODEWITHSUNDEEP
scaladataframeapache-spark
user3868051
user3868051

Reputation: 1249

scala: get column name corresponding to max column value from variable columns list

I have the following working solution in a databricks notebook as test.

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)

enter image description here

Is there any way to make this generic? I have a usecase to make variable columns pass to this UDF and still get the max column name as output corresponding to the column having max value. Unlike above where I have hard coded the column names 'col1', 'col2' and 'col3' in the UDF.

Upvotes: 0

Views: 184

Answers (1)

Parvez Patel
Parvez Patel

Reputation: 211

Use below:

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

Upvotes: 2

Related Questions