Most of us have been burned by a Spark job that passed local development but failed silently in production, corrupting data or dropping records we didn't know were gone. Testing Spark pipelines: real-world lessons and examples is not a luxury—it's the foundation of confidence in your data code. In this article, we'll walk through the patterns we've built across trading platforms, e-commerce warehouses, and SaaS analytics to catch bugs before they reach production.
Why Spark testing is different
Unlike typical application code, Spark jobs process large distributed datasets and often run on schedules we can't interrupt. A bug in a transformation may not raise an exception; it might silently drop 0.5% of your records or compute the wrong aggregation for a specific edge case. By the time analytics or finance teams notice, the damage is already in the warehouse.
Testing Spark requires a different mindset. We can't rely solely on unit tests of pure functions because Spark's lazy evaluation and distributed execution create blind spots. We need small, fast unit tests for transform logic, shared session management to avoid startup overhead, and integration tests that run the full pipeline on realistic sample data.
Unit test transforms — Small DataFrames with pytest
The fastest feedback loop comes from testing individual transformation functions with pytest and small in-memory DataFrames. We create minimal test data, run our transform, and validate the output.
Here's a real example: a revenue pipeline that applies discounts and calculates taxes.
from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, IntegerType
def apply_discount_and_tax(df, discount_rate=0.1, tax_rate=0.08):
"""Apply discount to price, then add tax."""
return (
df
.withColumn("discounted_price", F.col("price") * (1 - F.lit(discount_rate)))
.withColumn("tax", F.col("discounted_price") * F.lit(tax_rate))
.withColumn("final_price", F.col("discounted_price") + F.col("tax"))
)
# Test code
def test_apply_discount_and_tax():
spark = SparkSession.builder.appName("test").master("local").getOrCreate()
schema = StructType([
StructField("product_id", StringType(), False),
StructField("price", DoubleType(), False),
])
test_data = [
("SKU001", 100.0),
("SKU002", 50.0),
]
df_input = spark.createDataFrame(test_data, schema=schema)
df_output = apply_discount_and_tax(df_input, discount_rate=0.1, tax_rate=0.08)
result = df_output.collect()
# For SKU001: discount = 100 * 0.9 = 90, tax = 90 * 0.08 = 7.2, final = 97.2
assert result[0]["final_price"] == 97.2
assert result[1]["final_price"] == 48.6This test runs in milliseconds and catches logic errors immediately. The key is to keep test DataFrames small and focused on one scenario at a time. We'll add more test functions for edge cases: null values, zero amounts, negative prices, and boundary conditions.
Shared session fixture — Create one SparkSession per test run
Creating a new SparkSession for every test is wasteful. SparkSession initialization takes time, especially in CI/CD. Instead, we use pytest fixtures to share a single session across all tests in a module or test run.
import pytest
from pyspark.sql import SparkSession
@pytest.fixture(scope="session")
def spark():
"""Create one SparkSession for the entire test session."""
session = (
SparkSession.builder
.appName("test-suite")
.master("local[*]")
.config("spark.sql.shuffle.partitions", "4")
.config("spark.driver.memory", "2g")
.getOrCreate()
)
yield session
session.stop()
@pytest.fixture(autouse=True)
def reset_spark_cache(spark):
"""Clear the cache between tests to avoid cross-test pollution."""
yield
spark.catalog.clearCache()
def test_discount_logic(spark):
# Use the shared spark session
schema = StructType([...])
df = spark.createDataFrame([...], schema=schema)
result = apply_discount_and_tax(df)
assert len(result.collect()) == 2
def test_tax_calculation(spark):
# Same spark session, no overhead
schema = StructType([...])
df = spark.createDataFrame([...], schema=schema)
result = apply_discount_and_tax(df, tax_rate=0.15)
# Assert tax was calculated correctly
passUsing scope="session" means Spark starts once and runs all tests. The reset_spark_cache fixture clears cached DataFrames between tests, preventing state leakage. This reduces test runtime from minutes to seconds on a typical laptop.
Compare DataFrames — chispa or assertDataFrameEqual
Comparing two DataFrames for equality is trickier than it sounds. Row order may differ, floating-point comparisons need tolerance, and null handling matters. We have two solid tools: chispa (third-party, more features) and assertDataFrameEqual (built into Spark 3.5+).
Here's how we use assertDataFrameEqual in production:
from pyspark.testing import assertDataFrameEqual
def test_revenue_aggregation(spark):
"""Verify daily revenue is summed correctly."""
schema = StructType([
StructField("date", StringType(), False),
StructField("region", StringType(), False),
StructField("revenue", DoubleType(), False),
])
input_data = [
("2024-01-15", "US-EAST", 1000.0),
("2024-01-15", "US-WEST", 500.0),
("2024-01-15", "US-EAST", 200.0),
("2024-01-16", "US-EAST", 800.0),
]
df_input = spark.createDataFrame(input_data, schema=schema)
# Aggregate by date and region
df_output = (
df_input
.groupBy("date", "region")
.agg(F.sum("revenue").alias("total_revenue"))
.orderBy("date", "region")
)
expected_schema = StructType([
StructField("date", StringType(), False),
StructField("region", StringType(), False),
StructField("total_revenue", DoubleType(), False),
])
expected_data = [
("2024-01-15", "US-EAST", 1200.0),
("2024-01-15", "US-WEST", 500.0),
("2024-01-16", "US-EAST", 800.0),
]
df_expected = spark.createDataFrame(expected_data, schema=expected_schema)
# assertDataFrameEqual handles row ordering and null values
assertDataFrameEqual(df_output, df_expected)If we were using chispa instead, the call would be assert_df_equality(df_output, df_expected). Both tools ignore row order by default and provide clear error messages showing which rows differ.
Integration tests — Run jobs on sample data end to end
Unit tests catch logic bugs, but integration tests catch orchestration bugs: missing columns, wrong data types, schema mismatches, and bugs that only appear when multiple transforms run in sequence.
An integration test runs the entire pipeline (or a realistic slice) on sample data and validates the final output. Here's a simplified e-commerce example:
def test_daily_sales_pipeline_integration(spark, tmp_path):
"""Run the full sales pipeline and validate output."""
# Create input data (simulating raw order events)
orders_schema = StructType([
StructField("order_id", StringType(), False),
StructField("customer_id", StringType(), False),
StructField("amount", DoubleType(), False),
StructField("status", StringType(), False),
StructField("created_at", StringType(), False),
])
orders_data = [
("ORD001", "CUST001", 99.99, "completed", "2024-01-15T10:00:00Z"),
("ORD002", "CUST002", 49.99, "completed", "2024-01-15T11:30:00Z"),
("ORD003", "CUST001", 150.00, "pending", "2024-01-15T12:00:00Z"),
("ORD004", "CUST003", 75.50, "completed", "2024-01-15T14:00:00Z"),
]
df_orders = spark.createDataFrame(orders_data, schema=orders_schema)
# Stage 1: Filter completed orders
df_completed = df_orders.filter(F.col("status") == "completed")
# Stage 2: Calculate daily totals by customer
df_daily_customer = (
df_completed
.withColumn("date", F.col("created_at").cast("date"))
.groupBy("date", "customer_id")
.agg(F.sum("amount").alias("daily_total"))
.orderBy("date", "customer_id")
)
# Stage 3: Write to temporary output
output_path = str(tmp_path / "output")
df_daily_customer.coalesce(1).write.mode("overwrite").parquet(output_path)
# Stage 4: Read back and validate
df_result = spark.read.parquet(output_path)
expected_schema = StructType([
StructField("date", StringType(), False),
StructField("customer_id", StringType(), False),
StructField("daily_total", DoubleType(), False),
])
expected_data = [
("2024-01-15", "CUST001", 99.99),
("2024-01-15", "CUST002", 49.99),
("2024-01-15", "CUST003", 75.50),
]
df_expected = spark.createDataFrame(expected_data, schema=expected_schema)
assertDataFrameEqual(df_result, df_expected)This test validates that the pipeline reads data, applies filters, performs aggregations, writes to storage, and reads back correctly. We use tmp_path (a pytest fixture) to create a temporary directory, so tests don't interfere with each other or with real data.
Common mistakes
- Forgetting to handle nulls. Your transform may work perfectly for clean data but fail or behave unexpectedly when nulls appear. Always test with null values in key columns.
- Relying on row order. Spark makes no guarantees about row order unless you explicitly sort. Tests that depend on a specific order are brittle. Use assertDataFrameEqual, which ignores order.
- Testing only happy paths. A production bug often stems from an edge case: empty DataFrames, single-row results, or boundary values. Test those explicitly.
- Creating a SparkSession in every test. It's slow and wasteful. Use pytest fixtures with session scope.
- Mocking too aggressively. Mocking Spark operations defeats the purpose of testing. Use real in-memory DataFrames instead.
- Ignoring performance in integration tests. A test that takes 5 minutes to run will be skipped by developers. Keep sample data small enough to run in under 30 seconds.
Key takeaways
- Unit test individual transforms with small in-memory DataFrames and pytest to catch logic errors early.
- Share a single SparkSession across all tests using pytest fixtures with session scope to reduce overhead.
- Use assertDataFrameEqual or chispa to compare DataFrames reliably, handling row order, nulls, and floating-point precision automatically.
- Run integration tests on realistic sample data that exercises the full pipeline from input to output, catching orchestration and schema bugs that unit tests miss.
- Always test edge cases: nulls, empty DataFrames, boundary values, and single-row results.
- Keep integration tests fast (under 30 seconds) so developers actually run them before committing.
Comments
Loading comments…