PySpark Performance Tuning: 10 Best Practices for Faster Jobs

Apache Spark has become the de facto standard for large-scale data processing, and Python has grown into one of the most popular languages for data science and machine learning. Combining the two, PySpark enables you to leverage the power of Spark and Python to build scalable, distributed applications to transform massive datasets.

However, getting optimal performance out of PySpark takes some effort. Out of the box, Spark jobs can be slow and resource-hungry if you don‘t configure them properly. Tuning Spark and structuring your PySpark application with best practices in mind is critical to achieving peak speed and efficiency.

In this post, we‘ll dive into 10 key areas for PySpark performance optimization, from Spark configurations to query tuning to caching strategies. By the end, you‘ll have a solid foundation to build high-performance PySpark applications. Let‘s get started!

1. Tune Spark Configurations for Your Cluster and Workload

One of the first things to look at is your Spark configuration. Spark has dozens of configuration properties that control everything from memory allocation to shuffle behavior to compression. While the defaults are sufficient to get started, to get the most performance out of Spark, you need to tune them for your specific cluster setup and workload.

Some key properties to look at include:

spark.executor.instances – Number of executors to launch. More executors = more parallelism, but too many can cause excessive overhead.

spark.executor.memory – Memory to allocate to each executor. Increase for memory-intensive workloads, but be careful not to set too high and cause Out of Memory errors.

spark.executor.cores – Number of cores per executor. Set based on number of cores per node in cluster.

spark.sql.shuffle.partitions – Number of partitions to use when shuffling data. Increase to boost parallelism, but set too high and tasks become bottlenecked on scheduling.

spark.sql.autoBroadcastJoinThreshold – Maximum size in bytes for a table to broadcast to all worker nodes when performing a join. Increase to enable more broadcast joins and avoid shuffling.

Ideal values for these depend on your cluster size, node specs, and whether you‘re running in standalone mode or on a cluster manager like YARN or Kubernetes. Experiment with different values and monitor the Spark UI to find the sweet spot.

2. Use Optimal Data Storage Formats and Compression

The format and compression you use for your data has a major impact on read and write speeds. For most analytical workloads on Spark, columnar file formats like Apache Parquet and Apache ORC offer significant benefits over row-oriented formats like CSV or JSON:

• Columnar storage enables reading only needed columns, not entire rows
• Columnar formats support flexible compression schemes and encoding to minimize storage
• Predicate pushdown can dramatically reduce the amount of data scanned

Avoid dumping data into uncompressed CSV or JSON files if performance is important. And while whole-file compression like gzip is better than no compression, Spark can‘t read columnar-compressed Parquet or ORC data without fully decompressing the files, killing performance. Let Parquet or ORC handle compression on its own.

3. Optimize Data Partitioning and File Size

In addition to storage format, how you lay out your data files can impact query speeds. Two key considerations are partitioning and file size.

Partitioning divides the data into folders based on the values of one or more columns. When queries filter on a partition column, Spark can avoid reading entire data files and only scan the relevant partitions. Partitioning works best when:

• Partition columns are frequently used in query filters
• Data has high cardinality in partition columns (many distinct values)
• Partitioning scheme produces evenly-sized partitions

Over-partitioning is a common pitfall that produces thousands of small files, which can slow down queries as Spark has to scan many more files. Aim for partitions of at least 128 MB in size. If partitions are too small, consider bucketing instead, which groups data into larger files within each partition.

4. Cache Frequently-Used DataFrames

One of the most powerful features of Spark is its ability to cache data in memory for lightning-fast access. Whenever you have a DataFrame that‘s queried multiple times, caching can avoid redundant disk I/O and re-computation.

To cache a DataFrame, simply call the cache() method:

df.cache()

Spark supports several different storage levels for caching, from memory-only to memory-and-disk to disk-only. Memory is the fastest option, but also the most limited. SparkSession‘s setCheckpointDir() method lets you set a directory for Spark to spill cached data to disk if it exceeds memory limits.

Be strategic about what data you cache. Focus on frequently-used DataFrames or tables that are expensive to recompute. And don‘t forget to uncache DataFrames when you‘re done with them to free up memory:
df.unpersist()

5. Minimize Shuffles in Queries and Pipelines

Shuffles are one of the most expensive operations in Spark. They occur when data needs to be re-partitioned across the cluster, typically for join or aggregation operations. During a shuffle, data is written to disk, transferred across the network, and read back into memory in a different arrangement.

Avoiding or minimizing shuffles is one of the best ways to speed up Spark jobs. Some techniques include:

• Using broadcast joins for joining large and small tables, rather than repartition joins
• Filtering data early in the query, before joins and aggregations, to reduce the data shuffled
• Repartitioning DataFrames with coalescent partitioning columns to avoid shuffles in downstream operations
• Using window functions for per-partition aggregations and analytics

Examining the Spark UI and query plans can help identify shuffles in your jobs. Aim to push shuffles down the DAG as far as possible to minimize the cumulative impact.

6. Parallelize Work with Optimal Partitioning

Spark‘s magic comes from parallelizing work across a cluster. But to fully utilize the cluster, you need to ensure your data is well-partitioned. Under-partitioning leads to a handful of tasks doing most of the work while other executors sit idle. Over-partitioning means more time spent scheduling and less on actual computation.

Aim for total partitions of 3-4 times the number of cores in your cluster. For smaller datasets, reduce partitioning to avoid having many more partitions than total cores.

When doing joins or aggregations, consider explicitly repartitioning to ensure even distribution of the data. Use coalescent partitioning columns to avoid data skew. And beware of using non-key columns or high-cardinality columns for partitioning, as that can lead to tiny or empty partitions.

7. Use DataFrames and Spark SQL, Not RDDs

Since Spark 2.0, DataFrames and Spark SQL have become the primary APIs for Spark. While you can still use RDDs in PySpark, DataFrames and SQL should be your go-to for most workloads.

Unlike RDDs, DataFrames are optimized out of core, using a query optimizer to generate efficient query plans. The query optimizer can:

• Pushdown predicates to filter data early and minimize I/O
• Automatically select optimal join algorithms (e.g. broadcast vs. shuffle hash join)
• Prune unneeded columns early to reduce data processed
• Cache and reuse computations without manual intervention

That‘s not to say RDDs never have a place – they offer finer-grained control over partitioning and computation. But for most data transformation and analytical workloads, DataFrames and SQL will be faster and require less hand-tuning.

8. Beware the Perils of UDFs

User-defined functions are extremely useful for extending Spark SQL with custom logic. However, they can also kill performance if used inappropriately.

Spark UDFs are black boxes to the optimizer – it can‘t peek inside them to optimize accordingly. More importantly, Python UDFs require data to be serialized between the JVM and Python worker on every row, adding substantial overhead.

If possible, try to express your logic using DataFrame and SQL operations instead of UDFs. When UDFs are unavoidable:

• Use Spark‘s built-in functions as much as possible within UDFs
• Consider Pandas UDFs, which operate on Pandas Series instead of individual values and are much faster
• Window UDFs can optimize per-partition operations

9. Monitor Performance with the Spark UI

The Spark UI is your best friend for monitoring and debugging job performance. It provides a wealth of information, including:

• A DAG visualization showing the stages and tasks in each job
• Metrics on task scheduling delays, task runtimes, and data I/O
• Timing of SQL queries and a visualization of the query plan
• Environment and configuration details

Keep the Spark UI open whenever developing or tuning Spark applications. It‘s invaluable for identifying bottlenecks, stragglers, and resource contention.

10. Leverage Spark 3.x Performance Optimizations

Finally, it‘s worth upgrading to the latest Spark 3.x release if you haven‘t already. Spark 3.0 and 3.1 introduced several major performance optimizations, including:

• Adaptive Query Execution to dynamically coalesce shuffle partitions
• Dynamic partition pruning to avoid scanning unneeded partitions
• Join hints to specify join strategy (e.g. broadcast, merge, shuffle hash)
• ANSI SQL compliance for easier porting of SQL queries to Spark

Spark 3.2 further improved performance with optimizations like partition-aware file scan planning and binary file vectorization.

If you‘re still on Spark 2.x, consider upgrading to take advantage of these enhancements. The Spark SQL Migration Guide can help you identify any breaking changes in SQL syntax between versions.

Go Faster with PySpark!

We‘ve covered a lot of ground in this post, but these 10 areas are some of the most impactful for PySpark performance. Putting these practices into action can dramatically speed up your Spark jobs and let you iterate faster.

However, PySpark performance tuning is a complex topic. The "best" practices in any situation depend heavily on factors like your data, cluster setup, and application requirements. Treat these as starting points for your own exploration, not ironclad rules.

The best way to optimize your PySpark jobs is to leverage tools like the Spark UI, experiment with different techniques, and measure the results. Over time, you‘ll build an intuition for what works best for your setup and workloads.

Now get out there and make your Spark jobs fly! And if you have any other PySpark performance tips, share them in the comments below.