Spark User Defined Functions

A user defined function, or UDF, lets you run your own Python logic on DataFrame columns. Built-in functions cover most needs. A UDF fills the gaps when a business rule has no ready-made function, such as a custom grading formula or a special text clean-up routine.

The Layman View

Spark is a large kitchen with a standard set of tools. A UDF is a hand-made tool that you bring from home. The kitchen accepts it, but every ingredient must travel to your tool and back. That trip costs time.

Creating a Simple UDF

Write a normal Python function, wrap it with udf, and declare the return type.

from pyspark.sql import functions as F
from pyspark.sql.types import StringType

def grade(score):
    if score is None:
        return "Unknown"
    if score >= 80:
        return "A"
    if score >= 60:
        return "B"
    return "C"

grade_udf = F.udf(grade, StringType())

df = spark.createDataFrame([("Asha", 91), ("Ben", 67), ("Chitra", 45)], ["name", "score"])
df.withColumn("grade", grade_udf("score")).show()

name    score  grade
Asha     91     A
Ben      67     B
Chitra   45     C

Decorator Style

@F.udf(returnType=StringType())
def initials(name):
    return "".join(part[0].upper() for part in name.split())

df.withColumn("initials", initials("name")).show()

Using a UDF in SQL

Register the function by name and call it inside any SQL query.

spark.udf.register("grade_sql", grade, StringType())
df.createOrReplaceTempView("students")
spark.sql("SELECT name, grade_sql(score) AS grade FROM students").show()

Why Plain UDFs Run Slowly

Spark stores data in the JVM. A plain Python UDF forces Spark to serialize each row, send it to a Python worker, run the function, and send the answer back.

JVM executor                    Python worker
-------------                   -------------
row 1 --serialize-->             run grade(row 1)
row 2 --serialize-->             run grade(row 2)
row 3 --serialize-->             run grade(row 3)
<--serialize-- results ---------

The optimizer also treats a UDF as a black box. It cannot push filters through it or reorder the work around it.

Pandas UDFs: The Faster Option

A pandas UDF processes a whole batch of values at once through Apache Arrow. Batches remove most of the per-row cost. Use them for numeric work and for logic that already exists in pandas.

import pandas as pd
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def to_fahrenheit(celsius: pd.Series) -> pd.Series:
    return celsius * 9 / 5 + 32

temps = spark.createDataFrame([(0.0,), (25.0,), (100.0,)], ["celsius"])
temps.withColumn("fahrenheit", to_fahrenheit("celsius")).show()

Plain UDF   : 1 row  per call  --> many small calls
Pandas UDF  : 10,000 rows per call --> few large calls

Handling Null Values

Spark passes null as None to your function. A function that calls a method on None raises an error and stops the job. Test for None at the top of every UDF.

Replace a UDF with Built-ins When Possible

# UDF version
upper_udf = F.udf(lambda s: s.upper() if s else None, StringType())
df.withColumn("n", upper_udf("name"))

# Built-in version: faster and optimizable
df.withColumn("n", F.upper("name"))

The when and otherwise functions replace the grade function above. Try them before you write a UDF.

df.withColumn("grade",
    F.when(F.col("score") >= 80, "A")
     .when(F.col("score") >= 60, "B")
     .otherwise("C")).show()

Decision Guide

  • Start with a built-in function. It runs fastest.
  • Use a pandas UDF when the logic needs Python libraries or vector math.
  • Use a plain UDF for small data or rare, complex rules.
  • Declare an exact return type. A wrong type produces null values without a warning.
  • Keep UDFs free of side effects such as writing files. Spark may run a function more than once.
Previous lesson
Back to course
Next lesson