Apply UDF to multiple columns in Spark Dataframe

user_3598286 0
Viewed 18002

I have a dataframe which looks like below 

| id| age|   rbc|  bgr| dm|cad|appet| pe|ane|classification|
+---+----+------+-----+---+---+-----+---+---+--------------+
|  3|48.0|normal|117.0| no| no| poor|yes|yes|           ckd|
....
....
....

I have written an UDF to convert categorical yes, no, poor, normal into binary 0s and 1s

def stringToBinary(stringValue: String): Int = {
    stringValue match {
        case "yes" => return 1
        case "no" => return 0
        case "present" => return 1
        case "notpresent" => return 0
        case "normal" => return 1
        case "abnormal" => return 0
    }
}

val stringToBinaryUDF = udf(stringToBinary _)

I am applying this to the dataframe as follows

val newCol = stringToBinaryUDF.apply(col("pc")) //creates the new column with formatted value
val refined1 = noZeroDF.withColumn("dm", newCol) //adds the new column to original

How can I pass multiple columns into the UDF so that I don't have to repeat myself for other categorical columns?

apache-spark user-defined-functions scala
3 Answers

You could also use foldLeft function. Having your UDF called stringToBinaryUDF:

import org.apache.spark.sql.functions._

val categoricalColumns = Seq("rbc", "cad", "rbc", "pe", "ane")
val refinedDF = categoricalColumns
    .foldLeft(noZeroDF) { (accumulatorDF: DataFrame, columnName: String) =>
         accumulatorDF
            .withColumn(columnName, stringToBinaryUDF(col(columnName)))
     }

That will respect the immutability and functional programming.

user_29649110
Related