UDFs (User Defined Functions) in Apache Spark - Step-by-Step Execution

Concept Flow - UDFs (User Defined Functions)

Define UDF function

↓

Register UDF with Spark

↓

Apply UDF to DataFrame column

↓

Spark runs UDF on each row

↓

New column with UDF results

↓

Show or use transformed DataFrame

You first write a function, then tell Spark about it as a UDF. Spark applies it to each row, creating a new column with results.

Execution Sample

Apache Spark

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

def greet(name):
    return f"Hello, {name}!"

greet_udf = udf(greet, StringType())
df = spark.createDataFrame([('Alice',), ('Bob',)], ['name'])
df2 = df.withColumn('greeting', greet_udf(df.name))
df2.show()

This code defines a UDF that adds a greeting to each name in the DataFrame.

Execution Table

Step	Action	Input	UDF Output	DataFrame State
1	Define function greet	name string	N/A	No DataFrame change
2	Register greet as UDF	greet function	UDF object created	No DataFrame change
3	Create DataFrame	[('Alice',), ('Bob',)]	N/A	DataFrame with column 'name' and 2 rows
4	Apply UDF to 'name' column	Row 1: 'Alice'	'Hello, Alice!'	New column 'greeting' added to row 1
5	Apply UDF to 'name' column	Row 2: 'Bob'	'Hello, Bob!'	New column 'greeting' added to row 2
6	Show DataFrame	N/A	N/A	DataFrame shows columns 'name' and 'greeting' with values

💡 All rows processed, UDF applied to each, DataFrame transformed with new column

Variable Tracker

Variable	Start	After Step 2	After Step 3	After Step 6
greet	Function not defined	Function defined	Function defined	Function defined
greet_udf	None	UDF object created	UDF object created	UDF object created
df	None	None	DataFrame with 2 rows and 'name' column	DataFrame unchanged
df2	None	None	None	DataFrame with 'name' and 'greeting' columns

Key Moments - 3 Insights

Why do we need to register the Python function as a UDF before using it in Spark?

Spark runs code distributed across machines and needs to know how to apply your Python function on its data. Registering as a UDF wraps your function so Spark can run it on each row, as shown in execution_table steps 2 and 4.

Does the UDF change the original DataFrame or create a new one?

The UDF does not change the original DataFrame. It creates a new DataFrame with the added column, as seen in execution_table step 6 where df2 has the new 'greeting' column.

What happens if the UDF is applied to a null or missing value?

If the input is null, the UDF receives None and should handle it. Otherwise, Spark will return null for that row's output. This is important to avoid errors during execution.

Visual Quiz - 3 Questions

Look at the execution table, what is the UDF output when the input name is 'Bob'?

A 'Hello, Bob!'

B 'Hello, Alice!'

C 'Bob'

D None

💡 Hint

Check execution_table row 5 where input is 'Bob' and see the UDF output column.

At which step does the DataFrame get a new column added?

A Step 3

B Step 6

C Step 4

D Step 2

💡 Hint

Look at execution_table rows 4 and 5 where the UDF is applied and new column values appear.

If we skip registering the function as a UDF, what will happen when applying it to the DataFrame?

A Spark will apply the function automatically

B Spark will throw an error because it doesn't recognize the function

C The DataFrame will remain unchanged

D The function will run but output will be null

💡 Hint

Refer to key_moments question about why registering as UDF is necessary.

Concept Snapshot

UDFs let you run your own Python code on Spark DataFrames.
Define a Python function.
Register it as a UDF with a return type.
Apply it to DataFrame columns.
Spark runs it on each row, creating new columns.
Original DataFrame stays unchanged.

Full Transcript

User Defined Functions (UDFs) in Apache Spark allow you to apply your own Python functions to DataFrame columns. First, you define a Python function that does what you want, like adding a greeting to a name. Then, you register this function as a UDF with Spark, telling it the return type. After that, you apply the UDF to a DataFrame column. Spark runs your function on each row's value and creates a new column with the results. The original DataFrame does not change; instead, a new DataFrame with the added column is created. This process lets you customize data processing beyond built-in Spark functions.

Practice

1. What is the main purpose of a User Defined Function (UDF) in Apache Spark?

easy

A. To create new Spark sessions

B. To manage cluster resources

C. To optimize Spark SQL queries automatically

D. To apply custom logic to each row of a DataFrame

Solution

1. Step 1: Understand UDF role in Spark

   A UDF lets you add your own logic to process each row in a DataFrame, beyond built-in functions.

2. Step 2: Compare options

   The other options relate to Spark setup or optimization, not custom row logic.

3. Final Answer:

   To apply custom logic to each row of a DataFrame -> Option D

4. Quick Check:

   UDF = custom row logic [OK]

Hint: UDFs add your own row-wise logic in Spark [OK]

Common Mistakes:

• Confusing UDFs with Spark session management
• Thinking UDFs optimize queries automatically
• Assuming UDFs manage cluster resources

2. Which of the following is the correct way to register a UDF in PySpark with a return type of IntegerType?

easy

A. spark.udf.register('my_udf', func, 'IntegerType')

B. spark.udf.register('my_udf', func)

C. spark.udf.register('my_udf', func, IntegerType())

D. spark.udf.create('my_udf', func, IntegerType())

Solution

1. Step 1: Recall UDF registration syntax

   The correct method is spark.udf.register with function name, function, and return type as a DataType object.

2. Step 2: Evaluate options

   spark.udf.register('my_udf', func, IntegerType()) uses IntegerType() correctly. spark.udf.register('my_udf', func) misses return type. spark.udf.create('my_udf', func, IntegerType()) uses a non-existent method. spark.udf.register('my_udf', func, 'int') uses a string instead of IntegerType().

3. Final Answer:

   spark.udf.register('my_udf', func, IntegerType()) -> Option C

4. Quick Check:

   Register UDF with return type object [OK]

Hint: Always specify return type as DataType object when registering UDF [OK]

Common Mistakes:

• Omitting the return type parameter
• Using incorrect method names like create
• Passing return type as string instead of DataType

3. Given the following PySpark code, what will be the output of df.select('name', 'age_plus_one').show()?

from pyspark.sql.functions import udf
from pyspark.sql.types import IntegerType

def add_one(age):
    return age + 1

add_one_udf = udf(add_one, IntegerType())
df = spark.createDataFrame([('Alice', 30), ('Bob', 25)], ['name', 'age'])
df = df.withColumn('age_plus_one', add_one_udf('age'))

medium

A. [[Alice, 31], [Bob, 26]]

B. [[Alice, 30], [Bob, 25]]

C. Error: UDF missing return type

D. [[Alice, 1], [Bob, 1]]

Solution

1. Step 1: Understand UDF logic

   The UDF adds 1 to the 'age' column values: 30 + 1 = 31, 25 + 1 = 26.

2. Step 2: Check DataFrame transformation

   The new column 'age_plus_one' will have values 31 and 26 for Alice and Bob respectively.

3. Final Answer:

   [[Alice, 31], [Bob, 26]] -> Option A

4. Quick Check:

   age + 1 = 31, 26 [OK]

Hint: UDF applies function row-wise, transforming each value [OK]

Common Mistakes:

• Forgetting to specify return type causing errors
• Assuming original age remains unchanged
• Confusing UDF output with input data

4. Identify the error in the following PySpark UDF code snippet:

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

def greet(name):
    return 'Hello ' + name

greet_udf = udf(greet)
df = spark.createDataFrame([('John',), ('Jane',)], ['name'])
df = df.withColumn('greeting', greet_udf('name'))
df.show()

medium

A. UDF missing return type specification

B. Function greet has wrong parameter type

C. DataFrame creation syntax is incorrect

D. withColumn method is used incorrectly

Solution

1. Step 1: Check UDF registration

   The UDF is created without specifying the return type, which is required.

2. Step 2: Verify other code parts

   Function parameter and DataFrame creation are correct. withColumn usage is proper.

3. Final Answer:

   UDF missing return type specification -> Option A

4. Quick Check:

   Always specify return type when creating UDF [OK]

Hint: Always specify return type when defining UDF [OK]

Common Mistakes:

• Omitting return type causes runtime errors
• Assuming return type is optional
• Misunderstanding function parameter types

5. You want to create a UDF that returns the length of a string column but safely handles null values by returning 0 instead of failing. Which of the following UDF definitions correctly implements this behavior?

hard

A. def safe_length(s):
return len(s)
safe_length_udf = udf(safe_length, IntegerType())

B. def safe_length(s):
if s is None:
return 0
else:
return len(s)
safe_length_udf = udf(safe_length, IntegerType())

C. def safe_length(s):
if s == '':
return 0
return len(s)
safe_length_udf = udf(safe_length, StringType())

D. def safe_length(s):
return s.length()
safe_length_udf = udf(safe_length, IntegerType())

Solution

1. Step 1: Handle null values safely

   def safe_length(s): if s is None: return 0 else: return len(s) safe_length_udf = udf(safe_length, IntegerType()) checks if input is None and returns 0, preventing errors.

2. Step 2: Check return type and logic

   def safe_length(s): if s is None: return 0 else: return len(s) safe_length_udf = udf(safe_length, IntegerType()) returns IntegerType and uses len(s) correctly. Others either miss null check or use wrong return type or method.

3. Final Answer:

   def safe_length(s): if s is None: return 0 else: return len(s) safe_length_udf = udf(safe_length, IntegerType()) -> Option B

4. Quick Check:

   Check nulls first, then return length [OK]

Hint: Check for None before calling len() in UDF [OK]

Common Mistakes:

• Not handling None causing runtime errors
• Returning wrong data type from UDF
• Using string methods incorrectly in Python