How to apply udf on a dataframe and on a column in scala?

guest 0
Viewed 46

I am beginner to scala. I tried scala REPL window in intellij. I have a sample df and trying to test udf function not builtin for understanding.

df:

scala> import org.apache.spark.sql.SparkSession
 val spark: SparkSession = SparkSession.builder.appName("elephant").config("spark.master", "local[*]").getOrCreate()
 val df = spark.createDataFrame(Seq(("A",1),("B",2),("C",3))).toDF("Letter", "Number")
 df.show()

output:

|Letter|Number|
+------+------+
|     A|     1|
|     B|     2|
|     C|     3|
+------+------+

udf for dataframe filter:

scala> def kill_4(n: String) : Boolean = {
     | if (n =="A"){ true} else {false}} // please validate if its correct ???

I tried

df.withColumn("new_col", kill_4(col("Letter"))).show() // please tell correct way???

error error: type mismatch

Second: I tried direct filter:

df.filter(kill_4(col("Letter"))).show()

output desired

+------+------+
|Letter|Number|
+------+------+
|     B|     2| 
|     C|     3| 
+------+------+-
apache-spark scala
2 Answers

You can register udf and use it in code as follows:

import org.apache.spark.sql.functions.col

def kill_4(n: String) : Boolean = {
     if (n =="A"){ true } else {false}
}
 
val kill_udf = udf((x: String) => kill_4(x))

df.select(col("Letter"),col("Number")
    kill_udf(col("Letter")).as("Kill_4") ).show(false)
Aman Sehgal0

Please look at the databricks documentation on scala user defined funcitons.

https://docs.databricks.com/spark/latest/spark-sql/udf-scala.html

enter image description here

You do not need the spark session to create a dataframe. I removed that code.

enter image description here

Your function had a couple bugs. Since it is very small, I created a inline one. The udf() call allows the function to be used with dataframes. The call to register allows it to be used with Spark SQL.

A quick SQL statement shows the function works.

enter image description here

Last but not least, we need the udf() and col() functions for the last statement to work.

In short, these three snippets solve your problem.

CRAFTY DBA0
Related