Unit testing Spark transformations in Apache Spark - Step-by-Step Execution

Concept Flow - Unit testing Spark transformations

Write transformation function

↓

Create test input DataFrame

↓

Apply transformation function

↓

Collect or compare output DataFrame

↓

Assert output matches expected

↓

Test Pass/Fail

Unit testing Spark transformations means writing a function, applying it to test data, and checking if the output is as expected.

Execution Sample

Apache Spark

def add_one(df):
    return df.withColumn('value_plus_one', df['value'] + 1)

input_df = spark.createDataFrame([(1,), (2,), (3,)], ['value'])
output_df = add_one(input_df)
output_df.show()

This code adds 1 to the 'value' column in a Spark DataFrame and shows the result.

Execution Table

Step	Action	Input DataFrame	Transformation Applied	Output DataFrame
1	Create input DataFrame	[{'value':1}, {'value':2}, {'value':3}]	None	[{'value':1}, {'value':2}, {'value':3}]
2	Call add_one function	[{'value':1}, {'value':2}, {'value':3}]	Add column 'value_plus_one' = value + 1	[{'value':1, 'value_plus_one':2}, {'value':2, 'value_plus_one':3}, {'value':3, 'value_plus_one':4}]
3	Show output DataFrame	[{'value':1, 'value_plus_one':2}, {'value':2, 'value_plus_one':3}, {'value':3, 'value_plus_one':4}]	Display rows	Output displayed as table
4	Compare output with expected	Output DataFrame	Check equality with expected DataFrame	Test Pass if equal
5	Test ends	N/A	N/A	Test result: Pass

💡 Test ends after output matches expected DataFrame

Variable Tracker

Variable	Start	After Step 1	After Step 2	After Step 3	Final
input_df	None	[{'value':1}, {'value':2}, {'value':3}]	[{'value':1}, {'value':2}, {'value':3}]	[{'value':1}, {'value':2}, {'value':3}]	[{'value':1}, {'value':2}, {'value':3}]
output_df	None	None	[{'value':1, 'value_plus_one':2}, {'value':2, 'value_plus_one':3}, {'value':3, 'value_plus_one':4}]	[{'value':1, 'value_plus_one':2}, {'value':2, 'value_plus_one':3}, {'value':3, 'value_plus_one':4}]	[{'value':1, 'value_plus_one':2}, {'value':2, 'value_plus_one':3}, {'value':3, 'value_plus_one':4}]

Key Moments - 3 Insights

Why do we create a small input DataFrame instead of using the full dataset?

We use a small input DataFrame to keep tests fast and focused on the transformation logic, as shown in execution_table step 1.

How do we check if the transformation worked correctly?

We compare the output DataFrame to an expected DataFrame to see if they match, as in execution_table step 4.

Why do we not run the full Spark job in unit tests?

Unit tests focus on small parts (transformations) to catch errors early and run quickly, avoiding full job overhead.

Visual Quiz - 3 Questions

Look at the execution table, what is the output DataFrame after step 2?

A [{'value':1}, {'value':2}, {'value':3}]

B [{'value':1, 'value_plus_one':2}, {'value':2, 'value_plus_one':3}, {'value':3, 'value_plus_one':4}]

C Empty DataFrame

D DataFrame with only 'value_plus_one' column

💡 Hint

Check the 'Output DataFrame' column in execution_table row for step 2

At which step do we compare the output DataFrame to the expected DataFrame?

A Step 1

B Step 3

C Step 4

D Step 5

💡 Hint

Look for the step mentioning 'Compare output with expected' in execution_table

If the input DataFrame had an extra column, how would the output DataFrame change after step 2?

A It would have the extra column plus 'value_plus_one'

B It would only have 'value_plus_one' column

C It would drop all columns

D It would be empty

💡 Hint

Transformation adds a new column but keeps existing ones, see execution_table step 2

Concept Snapshot

Unit testing Spark transformations:
- Write a function for your transformation
- Create a small input DataFrame
- Apply the function to input
- Compare output DataFrame to expected
- Assert equality to pass test

Full Transcript

Unit testing Spark transformations means writing a small function that changes a DataFrame, then testing it with a small example DataFrame. We apply the function, get the output, and check if it matches what we expect. This helps catch mistakes early and keeps tests fast. We track variables like input_df and output_df step by step. Key moments include why we use small data, how we compare outputs, and why we avoid full jobs in unit tests. The execution table shows each step clearly, from creating input to test pass.

Practice

1. What is the main purpose of unit testing Spark transformations?

easy

A. To check small parts of Spark code for errors early

B. To improve the speed of Spark jobs

C. To deploy Spark applications automatically

D. To visualize Spark data outputs

Solution

1. Step 1: Understand unit testing purpose

   Unit testing focuses on checking small parts of code to find errors early.

2. Step 2: Relate to Spark transformations

   Testing Spark transformations means verifying each step works correctly before full runs.

3. Final Answer:

   To check small parts of Spark code for errors early -> Option A

4. Quick Check:

   Unit testing = Early error detection [OK]

Hint: Unit tests catch errors early in small code parts [OK]

Common Mistakes:

• Confusing unit testing with performance tuning
• Thinking unit tests deploy code
• Assuming unit tests create visualizations

2. Which of the following is the correct way to compare actual and expected DataFrames in a Spark unit test?

easy

A. print(actual_df) == print(expected_df)

B. assert actual_df.collect() == expected_df.collect()

C. actual_df.show() == expected_df.show()

D. actual_df.count() != expected_df.count()

Solution

1. Step 1: Understand DataFrame comparison

   Comparing DataFrames requires checking their collected data, not just printing or counts.

2. Step 2: Identify correct assertion

   Using assert with collected rows compares actual data content correctly.

3. Final Answer:

   assert actual_df.collect() == expected_df.collect() -> Option B

4. Quick Check:

   Compare collected data for equality [OK]

Hint: Use collect() to compare DataFrame contents in tests [OK]

Common Mistakes:

• Comparing printed output instead of data
• Using show() which returns None
• Comparing only row counts without data

3. Given the following Spark code in a unit test:

input_df = spark.createDataFrame([(1, 'a'), (2, 'b')], ['id', 'value'])
transformed_df = input_df.filter("id > 1")
result = transformed_df.collect()

What is the value of result?

medium

A. [Row(id=1, value='a')]

B. []

C. [Row(id=1, value='a'), Row(id=2, value='b')]

D. [Row(id=2, value='b')]

Solution

1. Step 1: Understand the filter condition

   The filter "id > 1" keeps rows where id is greater than 1, so only id=2 remains.

2. Step 2: Collect the filtered rows

   Collect returns a list of Row objects matching the filter, here only one row with id=2 and value='b'.

3. Final Answer:

   [Row(id=2, value='b')] -> Option D

4. Quick Check:

   Filter id > 1 keeps only id=2 [OK]

Hint: Filter keeps rows matching condition before collect [OK]

Common Mistakes:

• Including rows with id=1 mistakenly
• Confusing collect output format
• Assuming filter excludes id=2

4. You wrote this unit test code:

expected_df = spark.createDataFrame([(1, 'x')], ['id', 'value'])
actual_df = spark.createDataFrame([(1, 'x')], ['id', 'value'])
assert expected_df == actual_df

Why does this test fail?

medium

A. DataFrames cannot be compared directly with '=='

B. The schemas of DataFrames are different

C. The data in DataFrames is different

D. Spark session is not initialized

Solution

1. Step 1: Understand DataFrame equality

   Spark DataFrames do not support direct '==' comparison; it checks object identity, not content.

2. Step 2: Correct way to compare DataFrames

   Use methods like collect() or exceptAll() to compare data content, not '=='.

3. Final Answer:

   DataFrames cannot be compared directly with '==' -> Option A

4. Quick Check:

   Use collect() for DataFrame equality, not '==' [OK]

Hint: Never use '==' to compare DataFrames directly [OK]

Common Mistakes:

• Assuming '==' compares data content
• Ignoring schema differences
• Forgetting to initialize Spark session

5. You want to unit test a Spark transformation that removes duplicate rows based on the 'id' column, keeping the first occurrence. Which approach correctly tests this transformation?

hard

A. Create input DataFrame without duplicates, apply transformation, assert output equals input

B. Create input DataFrame with duplicates, apply transformation, assert output count equals input count

C. Create input DataFrame with duplicates, apply transformation, assert output equals expected DataFrame without duplicates

D. Create input DataFrame with duplicates, apply transformation, assert output is empty

Solution

1. Step 1: Prepare input with duplicates

   To test duplicate removal, input must contain duplicate rows on 'id'.

2. Step 2: Apply transformation and compare output

   After transformation, output should have duplicates removed; compare with expected DataFrame without duplicates.

3. Final Answer:

   Create input DataFrame with duplicates, apply transformation, assert output equals expected DataFrame without duplicates -> Option C

4. Quick Check:

   Test removes duplicates by comparing expected output [OK]

Hint: Test duplicate removal by comparing to expected no-duplicate DataFrame [OK]

Common Mistakes:

• Testing with no duplicates in input
• Comparing only row counts, not content
• Expecting empty output after removal