Unittesting Apache Spark Applications
By: Maria Karanasou
Tags: Pyspark, Python, Apache Spark, Data Science, Towards Data Science

Unittesting Apache Spark Applications

A PySpark case

Unittesting Spark applications is not that straight-forward. For most of the cases you’ll probably need an active spark session, which means that your test cases will take a long time to run and that perhaps we’re tiptoeing around the boundaries of what can be called a unit test. But, it is definitely worth doing it.

So, should I?

Well, yes! Testing your software is always a good thing, and it will most likely save you from many headaches, plus, you’ll be forced to have your code implemented in smaller bits and pieces that’ll be easier to test, thus, gain in readability and simplicity.

Okay, then what do I need to do this?

Well, I’d say we can start with pip install spark-testing-base and work our way from there. We’ll also need pyspark (of course) and unittest (unittest2) and pytest for this — even though pytest is a personal preference.

Spark testing base is a collection of base classes to help with spark testing. For this example we’ll be using the base SQLTestCase which inherits from SparkTestingBaseReuse , that creates and reuses a SparkContext.

On SparkSession and SparkContext :

From personal experience (using the currently latest spark version 2.4.+), I’ve found that I needed to make some minor adjustments to the SQLTestCase , which is a test case I use quite a lot in my current project. So, here’s an example of the adjustments I’ve made to suit my needs:

Example of a base spark test case, based on Spark testing base’ s SQLTestCase

To sum up the changes I’ve made:

I added a configuration to have the timezone set to UTC for consistency. Timezone consistency is something very basic to have throughout your code, so please make sure you always set spark.sql.session.timeZone

for consistency. Timezone consistency is something very basic to have throughout your code, so please make sure you always set Another important thing to set in the configuration is the spark.sql.shuffle.partitions to something reasonable for the machine that will be running the tests, like <= cores * 2 . If we don’t do that, then spark will use the default value, which is 200 partitions, and it will unnecessarily and inevitably slow down the whole process. <= cores * 2 is a general good rule, not only for the tests.

to something reasonable for the machine that will be running the tests, like . If we don’t do that, then spark will use the default value, which is partitions, and it will unnecessarily and inevitably slow down the whole process. is a general good rule, not only for the tests. Also added a method to sort the dataframes to be compared before the comparison. There is a compareRDDWithOrder method in one of the base classes, but I think it is easier to work with dataframes.

method in one of the base classes, but I think it is easier to work with dataframes. The schema_nullable_helper method should be used with caution , as it may end up sabotaging your test case, depending on what you need to test. The use case for this is for when you create dataframes without specifying a schema (which is currently deprecated), because spark tries to infer the data types, sometimes you have inconsistencies for the Nullable flag between the two dataframes to be compared, depending on the data used to create them. This method will update one of the two dataframes’ schema to what the other’s schema is regarding nullables only.

method should be , as it may end up sabotaging your test case, depending on what you need to test. The use case for this is for when you create dataframes without specifying a schema (which is currently deprecated), because spark tries to infer the data types, sometimes you have inconsistencies for the Nullable flag between the two dataframes to be compared, depending on the data used to create them. This method will update one of the two dataframes’ schema to what the other’s schema is regarding only. And lastly, I added a slightly adjusted version of the setUp for the appName and the config . The session instantiation is also different in the latest pyspark version. (There is a pending release for the support of 2.2.+ and 2.3.+ spark versions still open here and here, so, we’ll be subclassing to work around this)