Understanding Adaptive Query Optimization (On Vs Off)

Understanding Adaptive Query Optimization (On Vs Off)

Did you (might) know?

1. Irrespective of the choice of language used in Spark(be it Scala or Python or SQL), it all gets converted to Java byte code at the end
2. AQE (Adaptive Query Execution) feature is on by default in Spark 3.x onwards, it had to be turned-on in Spark 2.x

I am sure you have seen this flow of Spark’s query execution / optimization flow in different flavors and formats. In this blog, we will see how this query execution plan differs when AQE is Off and when it is On.

Following diagram shows the view of query optimization done by Spark Driver when AQE is off.

When AQE is OFF

User runs a query involving an action (in this case: count()).

Tip: groupby() is a wide transformation and hence it causes a shuffle of partition from one task to another

1. When a query is submitted to Spark, driver creates ‘Unresolved logical plan’ first
2. It gets analyzed by driver program with use of metadata catalog to verify schema of the dataframe submitted along with the query. To see if there is any schema related issue, if in case the flow stops at this point with an error.
3. If its past schema check, Analyzed logical plan is created and dataframe is enriched with types for the columns involved.
4. Output from this stage is fed into Catalyst catalog, it creates an optimized logical plan using the metadata catalog.
5. Rules are applied at this stage. For example, if there are two dataframes that is joined and there is a filter condition followed by, filter is done first followed by join to reduce the number of rows that needs to be joined. As join is a costly wide transformation operation
6. Further optimizations like evaluating which join strategy (sort merge join, broadcast join, etc..) to be applied based on the size of dataframes, data skipping (if it can be and it its needed), predicate pushdown (moving filters early in the stage to reduce rows to process later)
7. Finally, we have number of physical plans are created, which is submitted to Cost model to select the cheapest cost consuming physical plan
8. This plan gets converted into RDD (Resilient Distributed Dataset) in the form of Java byte code. You could see this stage in your DAG as a last step with the name ‘WholeStageCodeGen’

Following diagram is almost the same as the previous one, just a difference is with AQE feature on, Runtime Statistics created at the end of query optimization flow and it is used to check if the plan can further be optimized utilizing the metadata catalog based on characteristics of dataframe, cluster configuration (like spark shuffle partition count: default 200 but user could override it). Also it has other capabilities mentioned here.

When AQE is ON

First DAG in the following image shows execution without AQE, you could see SortMergeJoin is used.

In the second one, with AQE on, you could see BroadcastHashJoin being used and a BroadCastExchange, it reduces data shuffling by broadcasting small tables.

In the final one, with localShuffleReader enabled, BroadCastExchange changed to CustomShuffleReader while the rest remains the same.

Note on localShuffleReader: When true and ‘spark.sql.adaptive.enabled’ is true, Spark tries to use local shuffle reader to read the shuffle data when the shuffle partitioning is not needed, for example, after converting sort-merge join to broadcast-hash join. (Reference : Spark documentation)

Reference: partner.databricks.com

Following is a simple DAG visualization from Databricks, with AQE is on. You could see the AQEShuffleRead in Stage 39 as first step.

With AQE enabled

Following is a query Execution plan from a join query, (read from Initial plan to all the way up in each section).

First box (Before the execution: Compile time) shows SortMergeJoin in Current Plan and note IsFinalPlan=false, meaning AQE has not finalized on physical plan yet.

Second box (During the execution : Runtime) still SortMergeJoin in Current Plan still IsFinalPlan=false, we have runtime statistics calculated at this stage as compared to Compile time statistics in previous box. (Runtime Statistics from AQE).

Third box (After the execution: AQE kicked-in) is where AQE has selected the final plan and you could see AQE in option, one improvement being BrostCastHashJoin is used instead of SortMergeJoin in previous stages.

Reference: Adaptive query execution — Azure Datiabricks | Microsoft Learn

Tip: You can get this query execution plan using df.explain(mode)

Key takeaway is that Spark creates ‘Runtime Statistics’ at the end of query optimization flow when AQE is on, this statistics is used to further optimize the logical plan.