Skip to main content
Join the official 2020 Python Developers SurveyStart the survey!

Assertion functions to test Spark Collections like DataFrames.

Project description

A collection of assertion functions to test Spark Collections like DataFrames

As you develop Spark applications, you can eventually end up writing methods that apply transformations over Spark DataFrames. In order to test the results, you can create pandas DataFrames and use the test functions provided by pandas as pyspark does not provide any functions to assist with testing.

spark-test provides testing functions similar to pandas but geared towards Spark Collections.

Let’s say you have a function to apply some transformations on a Spark DataFrame (the full code for this example can be found in tests/

def transform(df):
    Fill nulls with 0, sum 10 to Age column and only return distinct rows

    df =
    df = df.withColumn('Age', df['Age'] + 10)
    df = df.distinct()

    return df

We can then write a test case with as many test inputs as we need and test the results with assert_dataframe_equal:

from spark_test.testing import assert_dataframe_equal

def test_transform(spark, transform):

    input_df = spark.createDataFrame(
        [['Tom', 25], ['Tom', 25], ['Charlie', 24], ['Dan', None]],
        schema=['Name', 'Age']

    expected = spark.createDataFrame(
        [['Tom', 35], ['Charlie', 34], ['Dan', 0]],
        schema=['Name', 'Age']
    result = transform(input_df)

    assert_frame_equal(expected, result)

Of course, tests are more interesting when they fail so let’s introduce a bug in our transform function:

def bugged_transform(df):
    Fill nulls with 0, sum 10 to Age column and only return distinct rows

    df =  # Whoops! Should be 0!
    df = df.withColumn('Age', df['Age'] + 10)
    df = df.distinct()

    return df

Passing both functions to our test using pytest.mark.parametize yields the following output with a nice message on what failed:

$ pytest tests/
============================= test session starts =============================
platform linux -- Python 3.7.3, pytest-5.0.0, py-1.8.0, pluggy-0.12.0
rootdir: /home/tfarias/repos/spark-test
collected 2 items

tests/ .F                                                [100%]

================================== FAILURES ===================================
_______________________ test_transform[bugged_transform] ________________________

            assert left_d[key] == right_d[key], msg.format(
>               field=key, l_value=left_d[key], r_value=right_d[key]
E           AssertionError: Values for Age do not match:
E           Left=10
E           Right=11


Distributed under the MIT License.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for spark-test, version 0.2.8
Filename, size File type Python version Upload date Hashes
Filename, size spark-test-0.2.8.tar.gz (4.9 kB) File type Source Python version None Upload date Hashes View

Supported by

Pingdom Pingdom Monitoring Google Google Object Storage and Download Analytics Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page