Unit testing Spark transformations in Apache Spark - Time & Space Complexity
When we test Spark transformations, we want to know how the time to run tests grows as data grows.
We ask: How does the test execution time change when input data size increases?
Analyze the time complexity of the following Spark transformation test code.
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
input_df = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "value"])
result_df = input_df.filter(input_df.id > 1).select("value")
assert result_df.count() == 1
This code creates a small DataFrame, applies a filter and select transformation, then checks the count.
Look at what repeats when the input grows.
- Primary operation: The filter and select run on each row of the input DataFrame.
- How many times: Once per row, so as many times as there are rows (n).
As input rows increase, the number of operations grows roughly the same.
| Input Size (n) | Approx. Operations |
|---|---|
| 10 | About 10 filter checks and selects |
| 100 | About 100 filter checks and selects |
| 1000 | About 1000 filter checks and selects |
Pattern observation: The work grows directly with the number of rows.
Time Complexity: O(n)
This means the test time grows linearly as the input data size grows.
[X] Wrong: "Unit tests run instantly no matter how big the data is."
[OK] Correct: Even tests run transformations on all rows, so bigger data means more work and longer test time.
Understanding how test time grows helps you write tests that stay fast and reliable as data grows.
"What if we cache the DataFrame before filtering? How would the time complexity change?"
Practice
Solution
Step 1: Understand unit testing purpose
Unit testing focuses on checking small parts of code to find errors early.Step 2: Relate to Spark transformations
Testing Spark transformations means verifying each step works correctly before full runs.Final Answer:
To check small parts of Spark code for errors early -> Option AQuick Check:
Unit testing = Early error detection [OK]
- Confusing unit testing with performance tuning
- Thinking unit tests deploy code
- Assuming unit tests create visualizations
Solution
Step 1: Understand DataFrame comparison
Comparing DataFrames requires checking their collected data, not just printing or counts.Step 2: Identify correct assertion
Using assert with collected rows compares actual data content correctly.Final Answer:
assert actual_df.collect() == expected_df.collect() -> Option BQuick Check:
Compare collected data for equality [OK]
- Comparing printed output instead of data
- Using show() which returns None
- Comparing only row counts without data
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?Solution
Step 1: Understand the filter condition
The filter "id > 1" keeps rows where id is greater than 1, so only id=2 remains.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'.Final Answer:
[Row(id=2, value='b')] -> Option DQuick Check:
Filter id > 1 keeps only id=2 [OK]
- Including rows with id=1 mistakenly
- Confusing collect output format
- Assuming filter excludes id=2
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?
Solution
Step 1: Understand DataFrame equality
Spark DataFrames do not support direct '==' comparison; it checks object identity, not content.Step 2: Correct way to compare DataFrames
Use methods like collect() or exceptAll() to compare data content, not '=='.Final Answer:
DataFrames cannot be compared directly with '==' -> Option AQuick Check:
Use collect() for DataFrame equality, not '==' [OK]
- Assuming '==' compares data content
- Ignoring schema differences
- Forgetting to initialize Spark session
Solution
Step 1: Prepare input with duplicates
To test duplicate removal, input must contain duplicate rows on 'id'.Step 2: Apply transformation and compare output
After transformation, output should have duplicates removed; compare with expected DataFrame without duplicates.Final Answer:
Create input DataFrame with duplicates, apply transformation, assert output equals expected DataFrame without duplicates -> Option CQuick Check:
Test removes duplicates by comparing expected output [OK]
- Testing with no duplicates in input
- Comparing only row counts, not content
- Expecting empty output after removal
Ready to prove it?
Pass the test and this pattern moves unknown
