UDFs (User Defined Functions) in Apache Spark

Introduction

UDFs let you add your own custom logic to process data in Spark. They help when built-in functions are not enough.

You want to apply a special calculation to a column in a Spark DataFrame.

You need to transform data in a way Spark does not support by default.

You want to reuse a custom operation on many rows of data.

You want to keep your data processing code clean by separating custom logic.

You want to apply Python or Scala functions inside Spark SQL queries.

Syntax

Apache Spark

from pyspark.sql.functions import udf
from pyspark.sql.types import <ReturnType>

# Define a Python function
def my_function(value):
    # custom logic
    return modified_value

# Register the function as a UDF with return type
my_udf = udf(my_function, <ReturnType>())

# Use the UDF in DataFrame operations
df = df.withColumn('new_column', my_udf(df['existing_column']))

Replace with a Spark data type like StringType(), IntegerType(), etc.

UDFs work row by row on DataFrame columns.

Examples

This example shows using a UDF on an empty DataFrame. It runs without error and returns no rows.

Apache Spark

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

def add_one(x):
    return x + 1 if x is not None else None

add_one_udf = udf(add_one, IntegerType())

# Edge case: empty DataFrame
empty_df = spark.createDataFrame([], 'value INT')
empty_df = empty_df.withColumn('plus_one', add_one_udf(empty_df['value']))
empty_df.show()

This example shows a UDF used on a DataFrame with one row, adding a greeting message.

Apache Spark

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

def greet(name):
    if name:
        return 'Hello ' + name
    else:
        return 'Hello Stranger'

greet_udf = udf(greet, StringType())

# DataFrame with one row
one_row_df = spark.createDataFrame([('Alice',)], ['name'])
one_row_df = one_row_df.withColumn('greeting', greet_udf(one_row_df['name']))
one_row_df.show()

This example shows handling null values inside a UDF and applying it to multiple rows.

Apache Spark

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

def mark_end(text):
    if text:
        return text + ' END'
    else:
        return None

mark_end_udf = udf(mark_end, StringType())

# DataFrame with multiple rows including null
multi_df = spark.createDataFrame([('start',), (None,), ('middle',)], ['text'])
multi_df = multi_df.withColumn('marked', mark_end_udf(multi_df['text']))
multi_df.show()

Sample Program

This program creates a Spark DataFrame with numbers, defines a UDF to double each number, and applies it. It prints the DataFrame before and after applying the UDF.

Apache Spark

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

# Create Spark session
spark = SparkSession.builder.appName('UDF Example').getOrCreate()

# Define a Python function to double a number
def double_value(number):
    if number is None:
        return None
    return number * 2

# Register the function as a UDF
double_value_udf = udf(double_value, IntegerType())

# Create a DataFrame with sample data
data = [(1,), (5,), (None,), (10,)]
df = spark.createDataFrame(data, ['number'])

print('Original DataFrame:')
df.show()

# Use the UDF to create a new column with doubled values
df = df.withColumn('doubled', double_value_udf(df['number']))

print('DataFrame after applying UDF:')
df.show()

# Stop Spark session
spark.stop()

Important Notes

Time complexity: UDFs run once per row, so O(n) where n is number of rows.

Space complexity: Minimal extra space, mostly for function call overhead.

Common mistake: forgetting to handle null values inside the UDF causes errors.

Use UDFs when built-in Spark functions cannot do the needed operation efficiently.

Summary

UDFs let you add custom row-wise logic to Spark DataFrames.

Always specify the return type when registering a UDF.

Handle nulls carefully inside your UDF to avoid errors.

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