Testing PySpark Code Locally

Structure transformations as pure DataFrame functions, run them on a local SparkSession in pytest, and compare results reliably.

PySparkIntermediate2 min readUpdated 11 Oct 2026

~/notes/pyspark/testing-pyspark-code

Make transforms testable

Keep reading and writing at the edges. Put logic in functions that take and return DataFrames.

def daily_sales(orders: DataFrame) -> DataFrame:
    return (orders.filter("status = 'paid'")
                  .groupBy(F.to_date("ts").alias("business_date"))
                  .agg(F.sum("amount").alias("net_sales")))

One local SparkSession per test run

import pytest
from pyspark.sql import SparkSession
 
@pytest.fixture(scope="session")
def spark():
    s = (SparkSession.builder.master("local[2]")
         .config("spark.sql.shuffle.partitions", "2")   # fast for tiny data
         .getOrCreate())
    yield s
    s.stop()

Starting Spark takes seconds, so share one session across all tests.

Build input with an explicit schema

def test_daily_sales_ignores_unpaid(spark):
    rows = [("O1", "paid", 100.0, "2026-01-05 10:00:00"),
            ("O2", "cancelled", 50.0, "2026-01-05 11:00:00")]
    df = spark.createDataFrame(rows, "order_id string, status string, amount double, ts string") \
              .withColumn("ts", F.to_timestamp("ts"))
    result = {r.business_date.isoformat(): r.net_sales for r in daily_sales(df).collect()}
    assert result == {"2026-01-05": 100.0}

Collecting a tiny result into a Python dict makes the comparison independent of row order.

Finished this note?

LinkedIn X WhatsApp

Make transforms testable

One local SparkSession per test run

Build input with an explicit schema

Related notes