Spark Tutorial (Part 4) - The Swiss Army Knife: UDF/UDTF

Introduction

UDF/UDTF

By this point in your exploration, you’ve probably noticed that the built-in tools already cover most day-to-day data transformation needs, whether you’re working with a Spark DataFrame or a pandas-API-on-Spark DataFrame. But what if you want to write your own function — one that takes some arguments, runs your custom logic, and hands back exactly the result you want, just like writing a regular program? Is that possible in Spark?

It is, and Spark provides a set of tools built exactly for this: UDF / UDTF. What they do, in short, is take your data row by row, run it through your function, and output the result as a DataFrame. The tradeoff is that they’re typically slower than Spark’s built-in native functions — but in exchange, you get the flexibility to shape the output exactly to your needs, which makes them genuinely useful.

So what’s the actual difference between UDF and UDTF? Let’s look at the comparison below.

UDF vs. UDTF: What’s the Difference

The core distinction between a UDF (User-Defined Function) and a UDTF (User-Defined Table Function) comes down to two things: the mapping relationship between input and output, and the structure of what’s returned:

• UDF is one-to-one: it takes in one row of data and returns a single scalar value.
• UDTF is one-to-many: it takes in one row of data and, via yield, expands it into a full table with zero or more rows and multiple columns.

📊 UDF vs. UDTF Comparison Table

Let’s walk through a few real examples to see how UDF and UDTF are actually written.

UDF

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, FloatType, DoubleType

# Define the function
@udf(returnType=DoubleType(), useArrow=True)
def channel_bonus(revenue: int, channel: str):
    if channel == "POS":
        return revenue * 1.1
    elif channel == "Online":
        return revenue * 1.5
    else:
        return revenue

# Define the schema
schema = StructType([
    StructField("Months", IntegerType(), nullable=False),
    StructField("Revenue", IntegerType(), nullable=False),
    StructField("Channel", StringType(), nullable=False)
])

# Create the Spark session
spark = SparkSession.builder.appName("Spark_udf").getOrCreate()

# The dataset
data = [
    (1, 100, "POS"),
    (1, 150, "Online"),
    (2, 200, "POS"),
    (2, 250, "App"),
    (3, 300, "Online"),
    (3, 350, "Wholesale"),
    (1, 400, "App"),
    (2, 450, "POS")
]

# Convert to a Spark DataFrame
ps_df = spark.createDataFrame(data, schema)
ps_df.show()

# Apply the function
ps_df.withColumn('bonus', channel_bonus(col('Revenue'), col('Channel'))).show()

As you can see from the result, when defining the function you need to specify the return type up front. Once that’s done, you can call it directly inside a DataFrame operation using function_name(p1, p2) syntax.

Using a UDF in Spark SQL

There are two things to keep in mind when calling a custom function from Spark SQL:

1. Register the DataFrame as a temp view
2. Register the function so SQL can call it

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# To use it in a Spark SQL query

# Create the temp table
ps_df.createOrReplaceTempView("q1_revenue")

# Register the function
spark.udf.register("channel_bonus", channel_bonus)

# Run the SQL query
spark.sql("select Months, Revenue, Channel, channel_bonus(Revenue, Channel) as bonus from q1_revenue").show()

UDTF

UDTF really shines in situations where you need to expand a nested structure into multiple rows. That’s a bit abstract in the abstract, so let’s look at a concrete example.

The code below still needs a defined return format up front, but here you’re defining a class with a method named specifically eval, and instead of ending with return, it ends with yield. Since we need to preserve the original columns at query time, we call this through Spark SQL — and since we need to iterate and expand each row, we pair it with LATERAL so it loops through the results correctly.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col, lit
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, FloatType, DoubleType
from pyspark.sql.functions import udtf
import re

# Define the function
@udtf(returnType="keyword: string, length: int")
class ExtractHashtags:
    def eval(self, sentence: str):
        keywords = re.findall(r"#([\w-]+)", sentence)
        for keyword in keywords:
            yield (keyword, len(keyword))


# Create the Spark session
spark = SparkSession.builder.appName("Spark_udtf").getOrCreate()

# Define the schema
schema = StructType([
    StructField("sentence", StringType(), nullable=False)
])

data = [
    ("Hello world this is #big-data channel",),
    ("How can we use #Spark to process data",),
    ("Maybe we need to use some tool such as #airflow #redshift",)
]

ps_df = spark.createDataFrame(data, schema)

spark.udtf.register("extract_hashtags", ExtractHashtags)
ps_df.createOrReplaceTempView("ps_df")

spark.sql("""
    SELECT *
    FROM ps_df, LATERAL extract_hashtags(sentence)
""").show()

Closing Thoughts

These tools are all ways to build your own custom processing logic, and some of the concepts here might take a couple of tries to really click. I’d strongly recommend running through the examples yourself to get comfortable with them. You can also check out the official docs for a more complete reference: UDF and UDTF official guide