Testing PySpark DataFrame transformations | by Eriks Dombrovskis | Towards Data Science

If you like tests — not writing a lot of them and their usefulness then you have come to the right place. I mostly write Spark code using Scala but I see that PySpark is becoming more and more dominant. Unfortunately I often see less tests when it comes to developing Spark code with Python. I think unit testing PySpark code is even easier than Spark-Scala code(unit testing Spark-Scala). Lets go trough an example and see how easy it is to unit test DataFrame transformations.
We will imagine a simple dataset but which is so big that it is more convenient to test our transformations in our integrated development environment (IDE) on much smaller data size.
Not to complicate too much lets imagine that the employee dataset id column contains the employee age which we need to extract to a separate column called age. The age of the employee is described by the number before the dash (“-”) symbol.
To do so we need to extract the number before the dash symbol. It could be done in two ways — either using regular expression (regex) or splitting the column value by the dash symbol. Lets use regex in this example to extract the age:
\d+(?=-)
\d matches any digit character (0–9).
+ is a quantifier to match 1 or more proceeding this type of characters (digits).
(?=-) matches a group after the main expression without including it in the result after the dash symbol.
To use this regex in our transformation we will write a simple function that takes a DataFrame and name of the employee id column but returns a brand new DataFrame with age column:
def extract_age_func(input_df: DataFrame, id_col: str):    pattern = '\d+(?=-)'    return input_df.withColumn('age', regexp_extract(col(id_col), pattern, 0))
Now moving on to testing — we will be using one of the more popular and Python standard library package unittest. The great thing about unittest is its ease to create a reusable SparkSession. Lets at first write some SparkSession initialization code that we can reuse for all our tests:
class PySparkTestCase(unittest.TestCase):    """Set-up of global test SparkSession"""    @classmethod    def setUpClass(cls):        cls.spark = (SparkSession                     .builder                     .master("local[1]")                     .appName("PySpark unit test")                     .getOrCreate())    @classmethod    def tearDownClass(cls):        cls.spark.stop()
Here we created class PySparkTestCase() that extends unittest.TestCase class and has two class methods:
setUpClass() — “Hook method for setting up class fixture before running tests in the class.” Be sure not to use setUp instead of setUpClass to not initialize a new SparkSession per test. Here we set up SparkSession to run locally on one thread.
tearDownClass() — “Hook method for deconstructing the test fixture after testing it.” When tests have been run we specify to stop our SparkSession.
Before the start of writing our tests we have to understand what needs to be tested when doing DataFrame transformations. From a DataFrame point of view there are two things — DataFrame schema test and DataFrame data test. Lets create helper functions that can accomplish this for us:
def test_schema(df1: DataFrame, df2: DataFrame, check_nullable=True):    field_list = lambda fields: (fields.name, fields.dataType, fields.nullable)    fields1 = [*map(field_list, df1.schema.fields)]    fields2 = [*map(field_list, df2.schema.fields)]    if check_nullable:        res = set(fields1) == set(fields2)    else:        res = set([field[:-1] for field in fields1]) == set([field[:-1] for field in fields2])    return res
and
def test_data(df1: DataFrame, df2: DataFrame):    data1 = df1.collect()    data2 = df2.collect()    return set(data1) == set(data2)
test_schema() — takes two DataFrames and compares if there are differences between them schema wise. If schemas match the function return a True else False. Additionally there is flag whether to check column nullability as this is not always needed and sometimes can get tedious to manage.
test_data() — also takes two DataFrames and checks whether data in those DataFrames match — returns True if matches and False if does not match.
Great, we now have helper functions to check DataFrame equality. Lets write tests for our DataFrame transformation function extract_age_func():
class SimpleTestCase(PySparkTestCase):    def test_dataparser_schema(self):        input_df = self.spark.createDataFrame(            data=[['Jan', 'Janson', 'jj@email.com', '20-504123'],                  ['Jen', 'Jenny', 'jen@email.com', '55-357378'],                  ['Bill', 'Bill', 'bill@email.com', '79-357378']],            schema=['first_name', 'last_name', 'email', 'id'])        transformed_df = extract_age_func(input_df, "id")        expected_df = self.spark.createDataFrame(            data=[['Jan', 'Janson', 'jj@email.com', '20-504123', '20'],                  ['Jen', 'Jenny', 'jen@email.com', '55-357378', '55'],                  ['Bill', 'Bill', 'bill@email.com', '79-357378', '79']],            schema=['first_name', 'last_name', 'email', 'id', 'age'])        self.assertTrue(test_schema(transformed_df, expected_df))
The approach is very simple — we create an input DataFrame right in our test case and run it trough our transformation function to compare it to our expected DataFrame. Similarly for data comparison test:
def test_dataparser_data(self):    input_df = self.spark.createDataFrame(        data=[['Jan', 'Janson', 'jj@email.com', '20-504123'],              ['Jen', 'Jenny', 'jen@email.com', '55-357378'],              ['Bill', 'Bill', 'bill@email.com', '79-357378']],        schema=['first_name', 'last_name', 'email', 'id'])    transformed_df = extract_age_func(input_df, "id")    expected_df = self.spark.createDataFrame(        data=[['Jan', 'Janson', 'jj@email.com', '20-504123', '20'],              ['Jen', 'Jenny', 'jen@email.com', '55-357378', '55'],              ['Bill', 'Bill', 'bill@email.com', '79-357378', '79']],        schema=['first_name', 'last_name', 'email', 'id', 'age'])    self.assertTrue(test_data(transformed_df, expected_df))
Here you can also create the input DataFrame from a separate data file like .csv or .txt if its more convenient. Or as I usually do — create in test case but be careful of not to put sensitive information that you will be pushing to some code repository.
Now lets run these tests and check if our age extraction transformation works as expected. I will be using PyCharm for running these tests:
Excellent, looks like our age extraction transformation works as expected. The great thing about running these tests in our IDE is that we can also set breakpoints in our code and see our transformations or tests execute step by step:
This is all you need to start testing your PySpark transformations. As you can see the time to test locally in your IDE would be sufficiently less than packaging, sending it to your test environment and waiting for your cluster to spin up. And the debugging functionality of and IDE is so convenient that if you want to you can dissect your transformation step by step.
I hope this gives a good base to start writing your own tests. Good luck!
Code available on GitHub.