Testing Glue Pyspark jobs
By: Vincent Claes
Tags: AWS, Testing, Spark, Pyspark, Glue

Testing Glue Pyspark jobs

Setup your environment so that your Glue PySpark job reads from and writes to a mocked S3 bucket, thanks to moto server.

Photo by Scott Sanker on Unsplash

The challenge

A typical use case for a Glue job is;

you read data from S3;

you do some transformations on that data;

you dump the transformed data back to S3.

When writing a PySpark job, you write your code and tests in Python and you use the PySpark library to execute your code on a Spark cluster. But how do I let both Python and Spark communicate with the same mocked S3 Bucket?

In this article, I’ll show you how you can setup a mocked S3 bucket that you can access from your python process as well as from the Spark cluster.

The test environment

Prerequisites for creating a Glue job

We are using Glue 1.0, which means Python 3.6.8, Spark/PySpark 2.4.3 and Hadoop 2.8.5.

make sure;

note: Glue uses Hadoop 2.8.5, but for simplicity we use Hadoop 2.7 because it’s shipped with Spark 2.4.3.

Python dependencies

pipenv --python 3.6

pipenv install moto[server]

pipenv install boto3

pipenv install pyspark==2.4.3

PySpark code that uses a mocked S3 bucket

If you have followed the above steps, you should be able to run successfully the following script: ¹ ² ³

import os

import signal

import subprocess import boto3

from pyspark.sql import DataFrame

from pyspark.sql import SparkSession

# start moto server, by default it runs on localhost on port 5000.

process = subprocess.Popen(

"moto_server s3", stdout=subprocess.PIPE,

shell=True, preexec_fn=os.setsid

) # create an s3 connection that points to the moto server.

s3_conn = boto3.resource(

"s3", endpoint_url="http://127.0.0.1:5000"

)

# create an S3 bucket.

s3_conn.create_bucket(Bucket="bucket") # configure pyspark to use hadoop-aws module.

# notice that we reference the hadoop version we installed.

os.environ[

"PYSPARK_SUBMIT_ARGS"

] = '--packages "org.apache.hadoop:hadoop-aws:2.7.3" pyspark-shell' # get the spark session object and hadoop configuration.

spark = SparkSession.builder.getOrCreate()

hadoop_conf = spark.sparkContext._jsc.hadoopConfiguration() # mock the aws credentials to access s3.

hadoop_conf.set("fs.s3a.access.key", "dummy-value")

hadoop_conf.set("fs.s3a.secret.key", "dummy-value")

# we point s3a to our moto server.

hadoop_conf.set("fs.s3a.endpoint", "http://127.0.0.1:5000")

# we need to configure hadoop to use s3a.

hadoop_conf.set("fs.s3.impl", "org.apache.hadoop.fs.s3a.S3AFileSystem") # create a pyspark dataframe.

values = [("k1", 1), ("k2", 2)]

columns = ["key", "value"]

df = spark.createDataFrame(values, columns) # write the dataframe as csv to s3.

df.write.csv("s3://bucket/source.csv") # read the dataset from s3

df = spark.read.csv("s3://bucket/source.csv") # assert df is a DataFrame

assert isinstance(df, DataFrame) # shut down the moto server.

os.killpg(os.getpgid(process.pid), signal.SIGTERM) print("yeeey, the test ran without errors.")

Copy and paste the above code to a file called “pyspark-mocked-s3.py” and execute:

pipenv shell

python pyspark-mocked-s3.py

The output will look like:



* Running on (glue-test-1) bash-3.2$ python pyspark-mocked-s3.py* Running on http://127.0.0.1:5000/ (Press CTRL+C to quit) ... 127.0.0.1 - - [28/Nov/2019 20:54:59] "HEAD /bucket/source.csv/part-00005-0f74bb8c-599f-4511-8bcf-8665c6c77cc3-c000.csv HTTP/1.1" 200 -

127.0.0.1 - - [28/Nov/2019 20:54:59] "GET /bucket/source.csv/part-00005-0f74bb8c-599f-4511-8bcf-8665c6c77cc3-c000.csv HTTP/1.1" 206 - yeeey, the test ran without errors.

Writing a test case

The principles showed in the above script are applied in a more structured way in my repo testing-glue-pyspark-jobs.

In this repo, you will find a Python file, test_glue_job.py. This file is an example of a test case for a Glue PySpark job. It combines the above logic with the principles outlined in an article I wrote about testing serverless services. Have a look at the test case and follow the steps in the readme to run the test. For convenience, I have added the test case below.

Good luck!

¹ https://stackoverflow.com/a/50242383/1771155

² https://gist.github.com/tobilg/e03dbc474ba976b9f235

³ https://github.com/spulec/moto/issues/1543#issuecomment-429000739