A Better Feedback Loop for PySpark ETL: Test It in a Container

"It works on my machine" is a frustrating way to find out that an ETL job depends on a particular Python version, Java runtime, or Spark installation.

Containers give us a practical way to make that feedback loop more predictable. In this post, we'll look at a small PySpark word-count job and use its Docker setup to run an end-to-end test. The example is deliberately small: the goal is to focus on the development and testing workflow, not on building a production-scale pipeline.

The approach here uses Docker Compose to run the test suite inside a container. It does not use the Python Testcontainers library. Testcontainers is another good option when tests need to create and tear down dependencies dynamically; for this example, a single Compose service is enough.

The ETL job

The job reads text files, counts words with Spark, and writes the result as Parquet:

from pyspark.sql import SparkSession


class WordCountJob:
    def __init__(self, spark: SparkSession):
        self._spark = spark

    def execute(self, src: str, dst: str):
        self._spark \
            .sparkContext \
            .textFile(src) \
            .flatMap(lambda line: line.split(" ")) \
            .map(lambda word: (word, 1)) \
            .reduceByKey(lambda a, b: a + b) \
            .map(lambda t: (str(t[0]), int(t[1]))) \
            .toDF(["word", "count"]) \
            .write \
            .parquet(dst)

WordCountJob receives a SparkSession rather than constructing one internally. That small design choice makes the job easier to reuse: the caller controls the Spark runtime, while execute focuses on the data flow.

The stages are straightforward:

1. Read the input path with Spark's textFile.
2. Split each line into words.
3. Emit (word, 1) pairs and aggregate counts with reduceByKey.
4. Convert the result to a DataFrame and write it as Parquet.

This is a compact example of Python acting as the glue around a data-processing engine: Python describes the transformations, while Spark executes them.

Test the output, not just the plumbing

The test starts Spark in local mode, runs the real job against the repository's sample input, and reads the resulting Parquet file back:

from pyspark import SparkContext
from pyspark.sql import SparkSession


spark = SparkSession(SparkContext("local", "PySpark Word Count"))
output_path = "/tmp/golden/output.parquet"

WordCountJob(spark).execute(
    src="/code/raw/input.txt",
    dst=output_path,
)

actual = spark.read.parquet(output_path).collect()

That is more useful than checking only that execute returned without raising an exception. It exercises the input read, Spark transformations, Parquet write, and Parquet read in one test.

One improvement worth making in a production test is to avoid depending on Spark's row order. The current example compares a collected list of rows; Spark does not promise a stable order unless the query explicitly sorts. Compare mappings instead:

output_path = "/tmp/golden/output.parquet"
actual_counts = {
    row.word: row.count
    for row in spark.read.parquet(output_path).collect()
}

assert actual_counts == {
    "one": 4,
    "two": 1,
    "three": 3,
    "four": 1,
    "five": 1,
    "six": 1,
}

It's also a good idea to use a unique temporary output directory per test and stop the Spark session during teardown. That keeps tests independent, especially when they run in parallel.

Give the test a known runtime

The Compose file builds a service called jobs_executor and mounts the checkout at /code:

services:
  jobs_executor:
    build:
      context: .
      dockerfile: container/Dockerfile
    volumes:
      - ${PWD}:/code

The mount is important. It makes the input data and test files available inside the container without baking every source change into a new image. The test can use the same path every time: /code/raw/input.txt.

The Dockerfile installs Python, Java, Scala, and the Python requirements. The requirements file includes pytest and pyspark, so developers and CI use the same test runner and Spark package inside the image. In other words, the container is not just where the code happens to run; it is part of the test environment.

Run the repository's CI script from the project root:

./script/ci

The script builds the Compose image and then runs pytest in a one-off container. The GitHub Actions workflow uses that same script on pushes, so local and CI test commands follow the same path.

Why use a container for ETL development?

This setup helps in a few practical ways:

• Fewer machine-specific setup steps. A developer does not need to install Spark and its runtime dependencies directly on their laptop just to run the test.
• A more representative integration test. The test exercises file input, Spark execution, and Parquet output, rather than mocking the central parts of the job.
• A repeatable CI entry point. Local development and the workflow both call script/ci.
• A short feedback loop. A small local-mode job can be tested without connecting to a shared Spark cluster.

This is still a local Spark test, not a miniature production cluster. It will not prove that a job behaves correctly under a distributed workload, with production credentials, or against every external system. If the pipeline depends on a database, object store, or message broker, add an appropriate test service or use a tool such as Testcontainers to provision dependencies for the test.

Make the environment repeatable, too

Containers only help with reproducibility when their dependencies are compatible and controlled. In this repository, requirements.txt allows any PySpark version, while the Dockerfile installs default-jdk. During validation of this repository, that combination resolved to PySpark 4.2.1 with Java 11, and Spark failed to start because its classes require a newer Java runtime.

Before relying on this setup—or publishing it as a copy-and-paste tutorial—pin PySpark and install a Java version supported by that PySpark release. Pinning the Python dependencies also prevents a future package release from silently changing the test environment.

Other useful next steps for a real ETL project include:

• Add cases for empty input, repeated whitespace, and malformed records.
• Write to a unique temporary path and clean it up after each test.
• Separate transformation logic from input and output details when the job grows.
• Keep local-mode tests fast, and add a smaller number of integration tests for external systems.

The takeaway

An ETL test container does not need to be complicated. Here, one Compose service gives the job a consistent Python-and-Java environment, mounts the project data, and runs an end-to-end test against a real local Spark session.

That makes it easier to change a pipeline and get meaningful feedback before deploying it. Start with one representative job, make the test deterministic, and add the dependencies your pipeline actually needs.

This post was inspired by the practical data-engineering framing in Python for Data Engineering; the focus here is the containerized development and testing loop.

This is the github repo with the code.