Spark SQL UDF function

UDF function

def main(args: Array[String]): Unit = {
    
    val session: SparkSession = SparkSession
      .builder()
      .config("dfs.ha.namenodes.nameservices","xxxxx")
      .appName("sql demo")
      .master("local")
      .enableHiveSupport()
      .getOrCreate()

    // Custom function 
    val code: UserDefinedFunction = session.udf.register("code", (arg: Int) => {
    
      if (arg > 5) "F" else "M"
    })


    val sql_1:String =
      s""" |select * from |sparktuning.course_pay |""".stripMargin

    val df: DataFrame = session.sql(sql_1)
	
	// usage 1.
    df.select("chapterid")
      .withColumn("flag",code(col("chapterid")))
      .withColumnRenamed("chapterid","new_chapterid")
      .show(false)
      
		+-------------+----+
		|new_chapterid|flag|
		+-------------+----+
		|4            |M   |
		|7            |F   |
		|8            |F   |
		|5            |M   |
		|4            |M   |
		|9            |F   |
		|7            |F   |
		|0            |M   |
		|5            |M   |
		|4            |M   |
		|4            |M   |
		|0            |M   |
		|1            |M   |
		|4            |M   |
		|9            |F   |
		+-------------+----+
	
	
	// usage 2.
    val sql_2:String =
      s""" |select | code(chapterid) Sex |from |sparktuning.course_pay |""".stripMargin

    session.sql(sql_2).show(false)

		+---+
		|Sex|
		+---+
		|M  |
		|F  |
		|F  |
		|M  |
		|M  |
		|F  |
		|F  |
		|M  |
		|M  |
		|M  |
		|M  |
		|M  |
		|M  |
		|M  |
		|F  |
		+---+


  }