PySpark SQL UDF (User Defined Functions)

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

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

What is PySpark UDF?

UDFs, or User-Defined Functions, will already be familiar to you if you come from an SQL background. Most traditional RDBMS systems support user-defined functions, where you register custom logic in the database and then use it in SQL queries just like any built-in function.

PySpark UDFs work much like UDFs in traditional databases. In PySpark, you write a function using Python and then wrap it with the udf() API or register it as a SQL UDF so it can be used in DataFrame operations or SQL queries.

Why do we need a PySpark UDF?

UDFs are used to extend the functions of the framework and reuse this function on several DataFrame. For example, if you wanted to convert the first letter of a word in a sentence to capital case, PySpark built-in features don’t have this function; hence, you can create it as a UDF and reuse this as needed on many DataFrames. Once UDFs are created, they can be reused on several DataFrames and SQL expressions.

Before you create any UDF, do your research to check if a similar function you want is already available in Spark SQL Functions. PySpark 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 you create UDF’s you need to design them very carefully, otherwise you will come across performance issues.

Create a DataFrame

Before we jump in creating a UDF, first let’s create a PySpark DataFrame.

# Create DataFrame
from pyspark.sql import SparkSession

spark = SparkSession.builder.appName('AIStudyLabs.com').getOrCreate()

columns = ["Seqno", "Name"]
data = [("1", "john jones"), ("2", "tracey smith"), ("3", "amy sanders")]

df = spark.createDataFrame(data=data, schema=columns)
df.show(truncate=False)

Yields below output.

+-----+------------+
|Seqno|Name        |
+-----+------------+
|1    |john jones  |
|2    |tracey smith|
|3    |amy sanders |
+-----+------------+

Create a Function

The first step in creating a UDF is creating a Python function. Snippet below creates a function convertCase() that takes a string parameter and converts the first letter of every word to a capital letter. UDF takes parameters of your choice and returns a value.

def convertCase(str):
    resStr = ""
    arr = str.split(" ")
    for x in arr:
        resStr = resStr + x[0:1].upper() + x[1:len(x)] + " "
    return resStr

Note that there might be a better way to write this function. But for the sake of this article, I am not worried much about the performance and better ways.

Now convert the convertCase() function into a UDF by passing it to the PySpark SQL udf() method. Make sure to import the package before using it.

from pyspark.sql.functions import col, udf
from pyspark.sql.types import StringType

# Converting function to UDF
convertUDF = udf(lambda z: convertCase(z), StringType())

Note: The default return type of udf() is StringType, so you can also define the above UDF without explicitly specifying the return type.

# Converting function to UDF
# StringType() is by default hence not required
convertUDF = udf(lambda z: convertCase(z))

Using UDF with DataFrame

Now you can apply convertUDF() on a DataFrame column, just like any other built-in function.

df.select(col("Seqno"), convertUDF(col("Name")).alias("Name")) \
  .show(truncate=False)

This result below the output.

+-----+-------------+
|Seqno|Name         |
+-----+-------------+
|1    |John Jones   |
|2    |Tracey Smith |
|3    |Amy Sanders  |
+-----+-------------+

Registering PySpark UDF to use it on SQL

To use the convertCase() function on PySpark SQL, you need to register the function with PySpark using spark.udf.register().

# Using UDF on SQL 
spark.udf.register("convertUDF", convertCase, StringType())
df.createOrReplaceTempView("NAME_TABLE")
spark.sql("select Seqno, convertUDF(Name) as Name from NAME_TABLE") \
     .show(truncate=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 below error.

columns = ["Seqno", "Name"]
data = [("1", "john jones"), ("2", "tracey smith"), ("3", "amy sanders"), ('4', None)]

df2 = spark.createDataFrame(data=data, schema=columns)
df2.show(truncate=False)

df2.createOrReplaceTempView("NAME_TABLE2")
spark.sql("select convertUDF(Name) from NAME_TABLE2") \
     .show(truncate=False)

Note that in the above example, the record with Seqno 4 has a value of None in the name column. Since the UDF does not handle null values, applying it on the DataFrame results in the error shown below. Remember that in Python, None is treated as a null value.

Things to keep in mind:

• It is always recommended to handle null values inside a UDF rather than checking for null externally.
• If you are unable to perform a null check within the UDF, then at least use IF or CASE WHEN in your DataFrame or SQL logic to check for nulls and call the UDF only when necessary.

Performance concern using UDF

UDFs are a black box to PySpark, 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.