Introduction

Spark SQL UDF (a.k.a. User-Defined Function) is a highly useful feature of Spark SQL & DataFrame that extends the Spark built-in capabilities.

Note: UDFs are the most expensive operations hence use them only you have no choice and when essential.

What is Spark UDF?

UDF a.k.a User Defined Function, If you are coming from an SQL background, UDFs are nothing new to you as most of the traditional RDBMS databases support User Defined Functions, and Spark UDFs are similar to these.

In Spark, you create a UDF by writing a function in a language you prefer to use for Spark. For example, if you are using Spark with Java, you create a UDF in Java language and wrap it with the udf() function or register it as a UDF to use it on DataFrame and SQL respectively.

Why do we need a Spark UDF?

UDFs are used to extend the functions of the framework and re-use this function on several DataFrames. For example, if you wanted to convert every first letter of a word in a sentence to a capital case, Spark's built-in features don't have this function, hence you can create it as a UDF and reuse this as needed on many DataFrames. UDFs, once created, can be reused on several DataFrames and SQL expressions.

Before you create any UDF, do your research to check if the similar function you wanted is already available in Spark SQL Functions. Spark SQL provides several predefined common functions and many more new functions are added with every release. Hence, it is best to check before you reinvent the wheel.

When creating UDFs, you need to design them very carefully; otherwise, you will come across performance issues.

Create a DataFrame

Before we jump into creating a UDF, first let’s create a Spark DataFrame.

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.functions;
import java.util.Arrays;
import java.util.List;

SparkSession spark = SparkSession.builder().appName("UDF Example").master("local").getOrCreate();

List<Row> data = Arrays.asList(
  RowFactory.create("1", "Be the change that you wish to see in the world"),
  RowFactory.create("2", "Everyone thinks of changing the world, but no one thinks of changing himself."),
  RowFactory.create("3", "The purpose of our lives is to be happy.")
);

StructType schema = new StructType(new StructField[]{
  new StructField("Seqno", DataTypes.StringType, false, Metadata.empty()),
  new StructField("Quote", DataTypes.StringType, false, Metadata.empty())
});

Dataset<Row> df = spark.createDataFrame(data, schema);
df.show(false);

Yields the below output.

+-----+-----------------------------------------------------------------------------+
|Seqno|Quote                                                                        |
+-----+-----------------------------------------------------------------------------+
|1    |Be the change that you wish to see in the world                              |
|2    |Everyone thinks of changing the world, but no one thinks of changing himself.|
|3    |The purpose of our lives is to be happy.                                     |
+-----+-----------------------------------------------------------------------------+

Create a Function

The first step in creating a UDF is creating a Java function. Below is a function convertCase() which takes a string parameter and converts the first letter of every word to a capital letter. UDFs take parameters of your choice and return a value.

import org.apache.spark.api.java.function.MapFunction;

public static String convertCase(String strQuote) {
    String[] arr = strQuote.split(" ");
    StringBuilder builder = new StringBuilder();
    for (String word : arr) {
        builder.append(Character.toUpperCase(word.charAt(0)))
               .append(word.substring(1)).append(" ");
    }
    return builder.toString().trim();
}

Create Spark UDF to use it on DataFrame

Now convert this function convertCase() to a UDF by passing the function to Spark SQL udf(), this function is available in the org.apache.spark.sql.functions package. Make sure you import this package before using it.

import org.apache.spark.sql.api.java.UDF1;
import org.apache.spark.sql.expressions.UserDefinedFunction;

UserDefinedFunction convertUDF = functions.udf((UDF1<String, String>) s -> convertCase(s), DataTypes.StringType);

Now you can use convertUDF() on a DataFrame column.

// Using with DataFrame
df.select(functions.col("Seqno"), convertUDF.apply(functions.col("Quote")).as("Quote")).show(false);

This results in the below output.

+-----+-----------------------------------------------------------------------------+
|Seqno|Quote                                                                        |
+-----+-----------------------------------------------------------------------------+
|1    |Be The Change That You Wish To See In The World                              |
|2    |Everyone Thinks Of Changing The World, But No One Thinks Of Changing Himself.|
|3    |The Purpose Of Our Lives Is To Be Happy.                                     |
+-----+-----------------------------------------------------------------------------+

Registering Spark UDF to use it on SQL

In order to use the convertCase() function on Spark SQL, you need to register the function with Spark using spark.udf().register().

// Using it on SQL
spark.udf().register("convertUDF", (UDF1<String, String>) s -> convertCase(s), DataTypes.StringType);
df.createOrReplaceTempView("QUOTE_TABLE");
spark.sql("SELECT Seqno, convertUDF(Quote) FROM QUOTE_TABLE").show(false);

This yields the same output as the previous example.

null check

UDFs are error-prone when not designed carefully. For example, when you have a column that contains the value null on some records, and not handling null inside a UDF function returns the below error.

Exception in thread "main" org.apache.spark.SparkException: Failed to execute user defined function(anonfun$1: (string) => string)
    at org.apache.spark.sql.catalyst.expressions.ScalaUDF.eval(ScalaUDF.scala:1066)
    at org.apache.spark.sql.catalyst.expressions.Alias.eval(namedExpressions.scala:152)
    at org.apache.spark.sql.catalyst.expressions.InterpretedMutableProjection.apply(Projection.scala:92)
    at org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation$anonfun$apply$24$anonfun$applyOrElse$23.apply(Optimizer.scala:1364)
    at org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation$anonfun$apply$24$anonfun$applyOrElse$23.apply(Optimizer.scala:1364)
    at scala.collection.TraversableLike$anonfun$map$1.apply(TraversableLike.scala:234)
    at scala.collection.TraversableLike$anonfun$map$1.apply(TraversableLike.scala:234)
    at scala.collection.immutable.List.foreach(List.scala:392)
    at scala.collection.TraversableLike$class.map(TraversableLike.scala:234)
    at scala.collection.immutable.List.map(List.scala:296)
    at org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation$anonfun$apply$24.applyOrElse(Optimizer.scala:1364)
    at org.apache.spark.sql.catalyst.optimizer.ConvertToLocalRelation$anonfun$apply$24.applyOrElse(Optimizer.scala:1359)

It is always best practice to check for null inside a UDF function rather than checking for null outside.

Performance concern using UDF

UDFs are a black box to Spark hence it can’t apply optimization and you will lose all the optimization Spark does on DataFrame/Dataset. When possible, you should use Spark SQL built-in functions as these functions provide optimization.