Kakao Publishes Case Study on Optimizing Spark AQE with Coalescing Post Shuffle Partitions

Kakao's Applied Analytics Team published a case study utilizing Coalescing Post Shuffle Partitions, a core feature of Adaptive Query Execution (AQE) introduced in Spark 3.0. This approach dynamically adjusts the number of partitions based on runtime statistics to overcome the limitations of existing rule-based and cost-based optimizations.

Experimental Results and Performance Improvements

In an experiment with spark.sql.shuffle.partitions set to 10000 and AQE enabled, the number of partitions decreased to 5000, 358, and 118 across three Shuffle stages, respectively. Notably, during the second HashAggregate stage, Peak Memory was reduced by more than half, from 634.8GB to 332.0GB, significantly improving memory efficiency.

Key Configurations and Mechanism

AQE adjusts partition sizes to be close to spark.sql.adaptive.advisoryPartitionSizeInBytes (default 64MB). However, when spark.sql.adaptive.coalescePartitions.parallelismFirst is set to its default value of true, it prioritizes parallelism and operates based on the minimum partition size (minPartitionSize, default 1MB). The official Spark documentation recommends setting this value to false to align the final partition size with the Advisory Size.

Recommended Configuration

To effectively utilize AQE, it is recommended to first set spark.sql.shuffle.partitions to a sufficiently large value, run the job to determine the actual number of Shuffle partitions, and then reset the configuration based on that data. Additionally, explicitly specifying spark.sql.adaptive.coalescePartitions.initialPartitionNum allows for more precise control over the initial number of partitions.