AQE optimisations
In prev blog I discussed internal of AQE and how AQE nodes converge to a final physical plan. In this blog I will some of the optimisations performed by AQE.
AQE perform multiple optimisation to spark plan(some of them are during physical planning and some of them during logical plan reoptimisation) e.g.
Dynamically coalescing shuffle partitions: simplifies or even avoids tuning the number of shuffle partitions.
Dynamically switching join strategies: SMJ to BHJ conversion
Dynamically optimizing skew joins: Optimise skew joins etc.
These optimisations are nothing but spark rules which gets applied in a batch to spark plan and new plan gets generated.
@transient private val queryStageOptimizerRules: Seq[Rule[SparkPlan]] = Seq(
PlanAdaptiveDynamicPruningFilters(this),
ReuseAdaptiveSubquery(context.subqueryCache),
OptimizeSkewInRebalancePartitions,
CoalesceShufflePartitions(context.session),
// `OptimizeShuffleWithLocalRead` needs to make use of 'AQEShuffleReadExec.partitionSpecs'
// added by `CoalesceShufflePartitions`, and must be executed after it.
OptimizeShuffleWithLocalRead
) ++ context.session.sessionState.adaptiveRulesHolder.queryStageOptimizerRulesThese are list of rules which gets applied to spark plan during createQueryStage function.
Let’s discuss about Dynamically coalescing shuffle partitions spark rule:
If you have been working with Spark for some time, you might get familiar with the number, 200. There seem always 200 tasks running no matter how small is your data. As you can see from the following example, 200 tasks are there to process a very small dataset, with less than1kb of data for each task.
The reason behind the 200 tasks is that the number of shuffle partitions is fixed during a spark job execution with the default value as 200. The number of shuffle partitions determines the number of buckets into which a mapper task writes and the number of output partitions on the reducer side.
The number of shuffle partitions could affect Spark performance significantly. Too less shuffle partitions would increase memory pressure on executors and can lead to disk spills. On the other hand too few partitions would increase overhead of task management in driver. Therefore choosing right partition count is important and given a stage consists of multiple filter and aggregate operations, tuning partition count manually for one stage might not work for other stage.
This is where AQE’s dynamic coalescing feature helps. AQE splits a Spark Job into multiple query stages and re-optimise the query plan of downstream query stages based on the runtime statistics collected from the completed upstream query stages. The CoalesceShufflePartitoins rule is the AQE optimizer rule created for dynamically configuring the shuffle partition number. This rule coalesces contiguous shuffle partitions according to a target size of the output partitions. For the example used above, if we enable AQE and rerun the same query, we can see there is only one task created this moment. As the entire dataset used in this example is very small, only one partition is required.
To coalesce shuffle partitions, the CoalesceShufflePartitoins rule first needs to know the target size of the coalesced partition, which is determined based on the advisoryTargetSize (spark.sql.adaptive.advisoryPartitionSizeInBytes, default 64MB), the minPartitionSize (spark.sql.adaptive.coalescePartitions.minPartitionSize, default 1MB), and the minNumPartitions (spark.sql.adaptive.coalescePartitions.minPartitionNum, if it is not set, fall back to Spark default parallelism).
Get Subham Singhal’s stories in your inbox
Join Medium for free to get updates from this writer.
maxTargetSize = total shuffle input data size from the map output statistics / minNumPartitions
if advisoryTargetSize > advisoryTargetSize then targetSize = maxTargetSize
if targetSize < minPartitionSize then targetSize = minPartitionSize
With the target size ready, the CoalesceShufflePartitions rule can start to coalesce the partitions according to the target size and to define the coalesced partition specs (CoalescedPartitionSpec) which will be used to create shuffle reader later.
Let’s take an example of shuffle with 5 partition (shuffle.partitions = 5)
When the CoalesceShufflePartitions rule is applied, it gets the size statistics of all shuffle partitions from the MapOutputStatistics collected at the map stage. In our example, we have two shuffles, each of which has five shuffle partitions:
Let’s assume targetSize = 64Mb.
CoalesceShufflePartitions will first run through all shuffle partitions and get final size of all partitions. Here we have 5 partitions with final sizes as
50mb, 20mb, 30mb, 10mb and 80mb.
CoalesceShufflePartitions will start coalsecing paritions until size is < targetSize(64mb). With this we come up with 3 partitions group:
[50mb, 20 + 30 + 10 = 60mb, 80mb]
At this point, the CoalesceShufflePartitions rule created a sequence of CoalescedPartitionSpec objects, each of which defines the spec for one of the output partition. These partition specs will be used by the shuffle reader to output partitions accordingly. The CoalesceShufflePartitions rule creates an AQEShuffleReadExec operator which wraps the current shuffle query stage and creates the ShuffledRowRDD with the CoalescedPartitoinSpecs defined above. For each partition of the ShuffledRowRDD, which is defined by one CoalescedPartitionSpec, a shuffle reader is created to read shuffle blocks from the range of map output partitions defined by the startReducerIndex and endReducerIndex of the CoalescedPartitionSpec.
Reference:
