How to write multiple If statements in Spark UDF

Question

My requirement is to create categories for age.I am trying to write multiple if conditions in UDF but its taking else condition. my code is below.

My data

1,Ashok,23,asd
2,Joi,27,dfs
3,Sam,30,dft
4,Bob,37,dat

my code

val sqlContext = new org.apache.spark.sql.SQLContext(sc)
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.sql._
import org.apache.spark.sql.SaveMode
import sqlContext.implicits._
val a = sc.textFile("file2.txt")
a.foreach(println)

val coder: (Int=>String)=(arg:Int)=>{if(arg>20&&arg<27) "20-27";if(arg>30&&arg<37) "30-37"; else "38+"}

val co = udf(coder)

val a2 = a1.select(col("Id"),col("Name"),col("Age"),col("Dpt"))

a2.withColumn("range",co(col("Age"))).show()

Output i am getting

1,Ashok,23,asd,38+
2,Joi,27,dfs,38+
3,Sam,30,dft,38+
4,Bob,37,dat,38+

For each row it is showing 38+, please suggest the syntax.

user_7224597 · Answer

You should be using if - else if - else .

Also, you're skipping boundary values when you check for arg>27 && arg<30

Your UDF should look something like this :

val co = udf { (x: Int) => 
    if (x >= 20 && x <=27) "20-27"
    else if (x > 27 && x<=37 ) "28-37"
    else "38+" 
}

// co: org.apache.spark.sql.UserDefinedFunction = UserDefinedFunction(,StringType,List(IntegerType))

df.withColumn("range" , co($"age" ) ).show

// +---+-----+---+---+-----+
// | id| name|age|dpt|range|
// +---+-----+---+---+-----+
// |  1|Ashok| 23|asd|20-27|
// |  2|  Joi| 27|dfs|28-37|
// |  3|  Sam| 30|dft|28-37|
// |  4|  Bob| 37|dat|  38+|
// +---+-----+---+---+-----+

How to write multiple If statements in Spark UDF

1 Answers