User-Defined Functions (UDFs)

Spark tutorial · PySpark.in

User-Defined Functions (UDFs) and Pandas UDFs in PySpark

Overview: Sometimes the built-in functions (like those we've covered above) might not fulfill a specific need – perhaps you have some complex business logic or need to use a Python library on your Spark data. User-Defined Functions (UDFs) allow you to write custom transformations in Python and apply them to DataFrame columns. There are two main types in PySpark:

Using UDFs should be a last resort only when Spark's built-in functions cannot achieve the result. This is because UDFs, especially standard ones, are not optimized by Spark's Catalyst optimizer and can be much slower. They break the "SQL" optimizations and operate on data one row at a time in Python, which is expensive. Pandas UDFs mitigate some of this cost by operating on chunks of data (thus reducing overhead of function call per row) and using vectorized pandas/Numpy operations internally, which can be 10x-100x faster than standard UDFs in many cases.

When to use UDFs:

When NOT to use UDFs:

Let's illustrate how to write and use both a standard UDF and a Pandas UDF. We will use the employees DataFrame from section 1 to demonstrate, for example:

Writing a Standard PySpark UDF

To create a UDF, you define a Python function and then use pyspark.sql.functions.udf to register it, specifying the return type. Common return types are found in pyspark.sql.types (e.g., StringType, IntegerType, etc.).

Example: A UDF to mask email addresses. Suppose we want to anonymize the email by hiding part of the username. We'll keep the first 2 characters of the email username and replace the rest with *, maintaining the domain.

```

from pyspark.sql.functions import udf,col

from pyspark.sql.types import StringType

from pyspark.sql import SparkSession

# Define the masking function in Python

def mask_email(email):

parts = email.split("@")

if len(parts) != 2:

return email # if it's not a standard email, return as-is

username, domain = parts

if len(username) <= 2:

masked_username = "*" * len(username)

else:

masked_username = username[:2] + "*" * (len(username) - 2)

return masked_username + "@" + domain

# Convert it to a UDF

mask_email_udf = udf(mask_email, StringType())

spark = SparkSession.builder.appName("BuiltInFuncDemo").getOrCreate()

# Sample employee data

data = [

("Alice", "Smith", 34, "alice.smith@xyz.com", "Sales", 80000),

("Bob", "Jones", 23, "bob.jones@xyz.com", "Marketing", 50000),

("Charlie", "Lee", 36, "charlie.lee@xyz.com", "Sales", 120000),

("Diana", "Diaz", 29, "diana.diaz@xyz.com", "Finance", 70000)

]

columns = ["first_name", "last_name", "age", "email", "department", "salary"]

df = spark.createDataFrame(data, columns)

# Use the UDF on the DataFrame

df_masked = df.withColumn("masked_email", mask_email_udf(col("email")))

df_masked.select("email", "masked_email").show(truncate=False)

```

We can verify the masking:

This UDF worked as expected. However, note the performance consideration:

Best Practice: If this masking needed to be done often or on large datasets, consider if it could be done with Spark native functions (for example, regexp_replace with a regex might do a similar mask, or using substring logic). Only use the UDF because writing this logic with built-ins would be very convoluted (this one could potentially be done with regex though). And if you use it, be aware of the performance cost.

One improvement is to use Pandas UDF for such tasks if possible, as shown next.

Writing a Pandas UDF (Vectorized UDF)

A Pandas UDF leverages Apache Arrow to send a batch of data to Python, operate on it with pandas (which uses C code under the hood for vectorized operations), and return a batch. This amortizes the Python overhead over many rows. Pandas UDFs are declared using pandas_udf and require specifying the return data type. You can write them as a function with the decorator @pandas_udf or use pandas_udf(function, returnType).

Example: Pandas UDF to add 10 to the age (as a trivial example). This operation could easily be done with Spark’s col("age") + 10, but we illustrate the syntax and concept.

```

import pandas as pd

from pyspark.sql.functions import pandas_udf,col

from pyspark.sql.types import IntegerType

from pyspark.sql import SparkSession

# Define a pandas UDF (Series to Series) that adds 10 to a number

add_ten_udf = pandas_udf(lambda s: s + 10, IntegerType())

spark = SparkSession.builder.appName("BuiltInFuncDemo").getOrCreate()

# Sample employee data

data = [

("Alice", "Smith", 34, "alice.smith@xyz.com", "Sales", 80000),

("Bob", "Jones", 23, "bob.jones@xyz.com", "Marketing", 50000),

("Charlie", "Lee", 36, "charlie.lee@xyz.com", "Sales", 120000),

("Diana", "Diaz", 29, "diana.diaz@xyz.com", "Finance", 70000)

]

columns = ["first_name", "last_name", "age", "email", "department", "salary"]

df = spark.createDataFrame(data, columns)

# Use the pandas UDF on the 'age' column

df_age_plus = df.select("first_name", add_ten_udf(col("age")).alias("age_plus_10"))

df_age_plus.show()

```

We see it correctly added 10 to each age. Under the hood, Spark will take batches of the "age" column (as a pandas Series), add 10 to the entire Series in one operation (which is fast), and return the new Series. This happens for each partition in parallel, using Arrow to avoid a lot of serialization overhead. The result is then combined back into a Spark DataFrame.

Performance Consideration: A Pandas UDF can be much faster than a standard UDF. In fact, vectorized UDFs can approach the performance of built-in Scala UDFs and are often 10x+ faster than equivalent row-by-row Python UDFs. However, they're still typically slower than pure built-in expressions in Spark, because of the serialization and Python execution (and you must ensure Arrow is installed and versions match). They also use more memory because they materialize data in memory as pandas objects. So, use them when necessary, but not frivolously.

Another example (for demonstration only): We could vectorize our email masking example. Pandas has string methods, so we could do something like:

```

import pandas as pd

from pyspark.sql.functions import pandas_udf, col

from pyspark.sql.types import StringType

from pyspark.sql import SparkSession

# Create Spark session

spark = SparkSession.builder.appName("EmailMask").getOrCreate()

# Pandas UDF to mask email

@pandas_udf(StringType())

def mask_email_pandas(series: pd.Series) -> pd.Series:

return series.str.replace(r'(?<=..).*?(?=@)', lambda m: "*" * len(m.group(0)), regex=True)

# Sample DataFrame

data = [

("alice.smith@xyz.com",),

("bob.jones@abc.com",),

("charlie.brown@test.com",)

]

df = spark.createDataFrame(data, ["email"])

# Apply masking UDF

df.withColumn("masked_email", mask_email_pandas(col("email"))).show(truncate=False)

```

This uses a regex with a lookbehind (?<=..) and lookahead (?=@) to replace the middle part of the username with asterisks. This would likely be faster than our pure Python UDF because it operates with pandas vectorized string functions (which are C optimized). But writing that regex is arguably more complex than the Python logic we wrote, so it's a trade-off between code simplicity and performance. (For actual use, Spark's own regexp_replace could do a similar thing without any UDF at all.)

When to use Pandas UDF vs standard UDF:

Important: Using UDFs (either type) can limit Spark's ability to optimize the query. For example, if you use a UDF in a filter, Spark can't push that filter down to the data source or do it as early as possible, because it doesn't understand the UDF's logic. So sometimes UDF-heavy pipelines might be slower not just due to UDF cost, but due to lost optimization opportunities.

Best Practices Recap for UDFs:

UDF vs Pandas UDF vs Built-in: A Quick Performance Thought

To summarize, imagine we needed to do a custom operation on a million records:

In testing, Pandas UDFs can be an order of magnitude faster than plain UDFs. So if you absolutely must use a UDF, and your Spark version and environment support Arrow, go for Pandas UDF.

Conclusion: UDFs give you flexibility to extend PySpark with custom logic. Use them sparingly and wisely. Always ask: Can I do this with existing Spark functions or SQL? If yes, do it that way. If no, then implement as efficiently as possible (vectorize if you can). Also, document the UDF well, because it will be a black box in your pipeline – other developers (or even yourself in the future) should easily understand what it's doing.

By mastering when and how to use UDFs, along with all the built-in functions covered in previous sections, you can handle a vast array of data transformation scenarios in PySpark, balancing practicality and performance.

More Spark tutorials

All tutorials · Try the free PySpark compiler · Practice challenges