Tricks to Boost Your Spark Pipeline Performance
By: Yu Zhou
Tags: Data Engineering, Pipeline, AWS, Hadoop, Spark

0. What does my Spark pipeline look like?

Before I go over my three methods, I want to describe my Spark pipeline first. My Spark pipeline includes thousands of independent Spark applications. Each Spark application intakes a customer’s data in S3, makes a long chain of transformations, and finally writes outputs into S3. These Spark applications are all in Pyspark in AWS EMR. We concurrently run a hundred Spark applications from these thousands. On bad days, it takes up to 24 hours for the whole work. It’s worthwhile to point out that one customer’s data can much bigger than another customer’s data, therefore these independent Spark applications actually have different resource footprint.

1. Tidy Up Pipeline Output

The number one well-known strategy for performance optimization is reducing the size of data in your application: ingest only necessary raw data, and keep intermediate data frames compact. In one performance tuning sprint, I replaced one of our raw data source with another source that had lower granularity, I saw 50% performance improvement easily. If an alternative data source is not available, you can use input = spark.read.parquet("fs://path/file.parquet").select(...) to limit reading to only useful columns. Reading fewer data into memory will speed up your application.

It should be equally obvious that writing less output into your destination directory also improves performance easily. Because the last step often gets less attention, the last step where you write outputs is likely not the most compact. In our data pipeline, we were very liberal towards our output. There were columns and rows that were only informative the best; we created many unnecessary data. In one performance tuning sprint, I reduced our output size (by 40%), and I saw a 2-hour performance improvement. If you consider the cost of data storage, and the cost for downstream processes to handle unnecessary rows, tidy up your pipeline output spreads performance benefits to the whole system and company. Don’t forget an easy win there.

2. Balance Workload via Randomization

The number of cores for executors, the size of memory for executors are the very basics for Spark memory management. There is a long list of parameters: overhead memory size, default partition number, and more. You can refer to these two documents for details: spark.org, AWS EMR. To find a good resource configuration for a Spark application, you need solid engineering knowledge and experience. One insight I learned from playing with these knobs is that networking kills performance. By reducing spark.sql.shuffle.partitions , I once saw 4-hour performance improvement. Although high partitions improve resilience, low partitions allow clusters to avoid sending data around, hence saving networking resource for in-memory calculation.

Most writings on Spark performance tuning concerns one application, however, we usually do not care much about individual application performance. In our Spark pipeline, we have thousands of Spark applications, we want the total run time to stay low, we do not want to tune and assign every single application its own perfect resource configuration. Without per-application resource configuration, we create three buckets: small, medium, and large. Customers with large data volume get a large memory configuration, customers with small data volume get a small memory configuration. With different “sizes” of applications in the pipeline, we ran into issues such as a long-tail pipeline (one last application takes an extra long time), larger applications blocking smaller applications, and other issues. Although we could use delicate engineering to address these issues, I found randomization a quicker way to go. In one performance tuning sprint, I simply randomized the thousands of applications and observed a 4-hour performance improvement. Randomization of the application not only made the workload evenly distributed (easier cluster-level resource provisioning) but also well mixed different sizes of applications together (less idleness).

3. Replace Joins with Window Functions

Inside our Spark pipeline, we manipulate raw input data by creating and adding new columns. In this process, we inevitably make input data frames larger and create intermediate data frames. withColumn() is a common pyspark.sql function we use to create new columns, here links to its official Spark document. We create intermediate tables because we handle different business logic in each intermediate table, then in a later step, we would join the tables to get a final table. This pattern is all good and familiar. However, when the chain of data manipulation grows longer and longer, a simple join between intermediate tables later in the process can easily create memory error. Familiar techniques such as persist() to cache intermediate data does not even help.

In one performance tuning sprint, I decided to avoid joins because of consistent memory problems. I instead used Window functions to create new columns that I would otherwise accomplish via joins. I not only avoid memory problem but also I saw a 33% performance improvement. Joins involved two tables, however, the Window function only involved a few columns. Needless to say, writing Window functions were less straightforward than joins, but replacing joins with Window functions certainly helped simplify the computation graph behind the scenes.