How to Write Unit Tests for PySpark Applications?

Unit testing is one of the most important practices in software development, and it’s just as essential when working with big data applications like PySpark. When working with big data using PySpark, it's common to write transformation logic for processing data. But how can we be sure our logic is correct in every scenario?

The simple and powerful answer: Unit Testing.

Unit testing helps you validate your PySpark code early, before it’s deployed into production. When your logic is processing millions or even billions of records, even the smallest bug can result in incorrect reports, financial loss, or data quality issues.

In PySpark, unit testing focuses on isolated testing of transformation functions without depending on external data sources or the Spark context. Testing is a crucial part of writing reliable PySpark code. Whether you're working with DataFrames or complex transformation logic, ensuring your code works as expected helps avoid costly bugs in production.

In this article, you will learn how to write unit tests for your PySpark applications, with real-world examples, best practices, and clear explanations before every code block.

Key Points-

Unit testing is critical for preventing data quality issues in PySpark applications.
PySpark transformations should be tested using mock DataFrames, not real data sources.
You can clean phone numbers using regexp_replace() to keep only digits.
SparkSession should be created once per test class using setUpClass() for efficiency.
Test cases should compare expected vs. actual DataFrames using assertDataFrameEqual.
withColumn() lets you apply transformations to DataFrame columns.
Python's unittest is a reliable framework to structure and run PySpark tests.
Clear and consistent test data is essential for valid testing.
Real-world problems (like cleaning phone numbers) make excellent practice examples.
Well-tested PySpark code is easier to maintain, debug, and scale in production environments.

What is Unit Testing in PySpark?

Unit testing in PySpark means testing individual transformation functions or pieces of logic using mock data. You’re not testing full pipelines or reading/writing real files — you're just checking that your function does what it's supposed to do on small, in-memory test data.

Why Unit Testing is Important in PySpark?

Python provides a built-in testing framework called unittest — it's simple, powerful, and works well with PySpark.

With unittest, you can:

Create reusable test cases
Automate testing
Validate DataFrame transformations easily
Improve code reliability and maintainability

When working with PySpark, we often deal with complex ETL pipelines, data cleaning, and transformations. If we don’t test these properly, even a small bug can cause incorrect results in production.

Benefits of Unit Testing in PySpark:

Catch errors early
Validate business logic
Easy debugging
Confidence in code quality
Easy maintenance for future changes

Real-World Use Case: Cleaning Phone Numbers in Customer Data

Let’s say you are working with customer data.
The phone numbers in your dataset are messy; they have different formats, extra characters, and inconsistent styles.

Our goal is standardize all phone numbers to the below format.


# Input phone number and desired phone number
Input  →"(123) 456-7890"
Output → "1234567890"

Use Transformation Function to Clean Phone Numbers

We can create a function clean_phone_number() that uses PySpark regexp_replace() to remove all non-numeric characters from a phone number string, keeping only the digits.


# main_transformation.py
from pyspark.sql.functions import regexp_replace, col

# Function to clean phone numbers
def clean_phone_number(df, column_name):
    return df.withColumn(column_name, regexp_replace(col(column_name), "[^0-9]", ""))

The code above:

col(column_name) is used to reference the phone number column.
The regexp_replace(..., "[^0-9]", "") This functionreplaces a string in a column of a data frame using regular expressions. In this example, it removes any characters that are not digits.
Finally, withColumn() applies this transformation to the DataFrame, returning a new one with cleaned phone numbers.

Create a SparkSession for Testing

Before writing any tests, we need to set up a SparkSession. This acts as the entry point to all PySpark functionality.

In unittest, we can use setUpClass() and tearDownClass() methods to create and stop SparkSession only once for all tests (this makes the testing faster).


# Initialize SparkSession for testing
import unittest
from pyspark.sql import SparkSession

class PySparkTestCase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.spark = SparkSession.builder \
            .appName("PySpark Unit Test Example") \
            .getOrCreate()

    @classmethod
    def tearDownClass(cls):
        cls.spark.stop()

In the code above,

setUpClass() runs once before all tests. It creates a shared SparkSession.
tearDownClass() runs once after all tests. It shuts down the Spark session.

Write a Test Case for the Function

We can write a test case to check if our clean_phone_number() function works properly.

We’ll:

Create a sample input DataFrame (Messy phone numbers)
Create an expected output DataFrame (cleaned phone numbers)
Apply our function and compare the results


# Messy phone numbers for sample data frame
sample_data = [
    {"name": "Sai", "phone": "(123) 456-7890"},
    {"name": "Ram", "phone": "123.456.7890"},
    {"name": "Geetha", "phone": "(123)-456-7890"},
    {"name": "Vani", "phone": "123 456 7890"}
]

Expected Output After Transformation

All phone numbers should only have digits.


# Cleaned phone numbers for expected data frame
expected_data = [
    {"name": "Sai", "phone": "1234567890"},
    {"name": "Ram", "phone": "1234567890"},
    {"name": "Geetha", "phone": "1234567890"},
    {"name": "Vani", "phone": "1234567890"}
]

Final Unit Test Case

You can use the helper functions to check DataFrame equality. Let's write tests for our data frame transformation function clean_phone_number().


# Final unit test case
from pyspark.testing.utils import assertDataFrameEqual

# Importing your function
from main_transformation import clean_phone_number  

class TestTransformation(PySparkTestCase):
    def test_clean_phone_number(self):

        # Create input and expected DataFrames
        original_df = self.spark.createDataFrame(sample_data)
        expected_df = self.spark.createDataFrame(expected_data)
 
        # Apply transformation 
        transformed_df = clean_phone_number(original_df, "phone")
       
        # Assert they are equal
        assertDataFrameEqual(transformed_df, expected_df)

In the code above,

createDataFrame() turns our sample dictionaries into Spark DataFrames.
clean_phone_number() is applied to clean the phone numbers.
assertDataFrameEqual() checks if both DataFrames are the same:
- Same schema
- Same column names
- Same values (ignoring row order)
If the data doesn't match, the test will fail with a detailed error.

Running the Unit Test for PySpark

Store your test code in a file named something like test_transformation.py and then run it using Python’s built-in test runner.


# Run the test using unit test
python -m unittest test_transformation.py

Expected Output after Running Test

If everything is correct, you'll get:

Frequently Asked Questions on Write Unit Tests for PySpark

Unit testing in PySpark involves testing individual transformation functions or logic in isolation using small, mock DataFrames. It ensures that each function works as expected before integrating it into large data pipelines or deploying to production.

Because PySpark applications process huge datasets, even small logic errors can lead to major data quality issues or incorrect business outcomes. Unit testing catches bugs early, validates transformations, and increases confidence in code quality.

Python’s built-in unittest module is most commonly used. It works well with PySpark and supports features like reusable test classes, assertions, and setup/teardown methods for managing SparkSession.

You can use assertDataFrameEqual from pyspark.testing.utils. For example:

Use setUpClass() in a base test class to create a SparkSession only once for all tests. For example:
class PySparkTestCase(unittest.TestCase): @classmethod def setUpClass(cls): cls.spark = SparkSession.builder.appName("UnitTest").getOrCreate() @classmethod
def tearDownClass(cls):
cls.spark.stop()

You can manually run test methods inside notebook cells. But for full test automation, it’s better to write .py files and run them using:
python -m unittest test_file.py

A good project structure is:
my_project/
├── main_transformation.py
├── tests/
│ └── test_transformation.py
Keep your test files in a separate folder like tests/ and name them starting with test_ to make them discoverable by unittest or pytest.

Include test data with nulls, empty strings, or unexpected types: For example:
sample_data = [{"name": "Alice", "phone": None}]
Make sure your transformation handles these gracefully, and write separate test cases for each edge scenario.

Yes. Avoid reading real files in unit tests. Instead, mock the data with spark.createDataFrame() using Python dictionaries. Use integration tests for end-to-end file I/O testing.

Chain the functions in a single test case if they depend on each other: For example:
df = clean_phone_number(df, "phone")df = normalize_name_column(df, "name")
But generally, it's best to test each function in isolation to keep unit tests simple and focused.

Conclusion

In this article, you’ve learned that unit testing in PySpark is an essential practice for building clean, reliable, and production-ready data pipelines. By writing tests for transformation functions like clean_phone_number(), you can catch bugs early, reduce production issues, and maintain high-quality code.

Here’s what we covered:

Wrote a transformation function to clean phone numbers using regexp_replace().
Created a reusable base test class that initializes a SparkSession for testing.
Built sample input and expected output data using mock dictionaries.
Wrote a test case to validate the transformation using assertions.
Ran the unit test using unittest to ensure everything works as expected.

By applying these practices, you’ll write PySpark code that’s easier to test, debug, and maintain — saving time and preventing costly mistakes in production environments.

Happy Learning!!

Reference

https://spark.apache.org/docs/latest/api/python/getting_started/testing_pyspark.html