Spark AQE — A Detailed Guide with Examples

A Practical Guide for Spark AQE

Source: Databricks

Spark AQE, or Adaptive Query Execution, is a feature introduced in Apache Spark 3.0 to enhance query performance by dynamically adjusting execution plans at runtime based on the statistics collected during execution.

This contrasts with traditional static query optimization, where the execution plan is determined before query execution and remains unchanged throughout.

In this blog, we will learn about Spark AQE with examples and use cases.

Key Features of Spark AQE

Here are some key aspects of Spark AQE:

1. Dynamic Partition Pruning: AQE can prune partitions dynamically, which means it can skip reading unnecessary partitions during query execution, reducing I/O and speeding up query processing.
2. Dynamic Join Optimization: AQE can optimize join operations dynamically by choosing the best join strategy based on the sizes of the tables involved. For instance, it can switch between broadcast join and shuffle join at runtime.
3. Optimizing Skewed Joins: AQE can detect skewed data and handle it by splitting the skewed partitions into smaller tasks, thereby balancing the load and improving overall performance.
4. Coalescing Shuffle Partitions: AQE can dynamically coalesce small shuffle partitions into larger ones to reduce the overhead associated with processing many small tasks, which improves the efficiency of the execution.
5. Handling Runtime Statistics: AQE leverages runtime statistics, such as the size of data processed, to make informed decisions about optimizing query execution plans dynamically.

AQE Framework — Adaptive Query Execution

This section breaks down the AQE workflow and its key features in a simple and digestible manner.

• Materialization Points: Breaks in the execution pipeline where data is fully processed (materialized), typically during shuffles or broadcast exchanges. These points divide a query into query stages.
• Query Stages: Subsections of a query, bounded by materialization points, that can be independently optimized and executed.

AQE Workflow: Step-by-Step

Here is a simple step-by-step guide:

1. Initial Execution:

• The AQE framework starts by executing leaf stages, which are stages that do not depend on any other stages.
• As these stages finish, they produce intermediate results (materialized data).

2. Reoptimization:

• Once a stage finishes materialization, AQE updates the query plan with new statistics gathered from the completed stage.
• The optimizer uses these statistics to reoptimize the query plan, applying logical and physical optimization rules, including AQE-specific ones.

3. Subsequent Execution:

• The newly optimized query plan identifies new stages ready for execution (stages where all dependencies have been materialized).
• This cycle of execution and re-optimization continues until the entire query is completed.

Key Features of AQE in Spark 3.0

Here are the key features of the AQE Framework:

Dynamically Coalescing Shuffle Partitions:

• What It Is: Merging small shuffle partitions into larger ones to reduce the overhead of managing many small tasks.
• Why It Matters: Improves efficiency and reduces processing time by minimizing the number of tasks.

Dynamically Switching Join Strategies:

• What It Is: Selecting the most efficient join strategy (e.g., switching from a shuffle join to a broadcast join) based on the sizes of the tables involved at runtime.
• Why It Matters: Enhances performance by choosing the optimal join method for the given data sizes, which can vary significantly.

Dynamically Optimizing Skew Joins:

• What It Is: Handling skewed data (where some partitions are much larger than others) by splitting these large partitions into smaller ones, balancing the load across tasks.
• Why It Matters: Prevents performance bottlenecks caused by uneven data distribution, ensuring more balanced and efficient execution.

Why AQE is Important?

• Adaptability: AQE makes Spark more adaptable by allowing it to react to actual runtime data characteristics, rather than relying solely on static plans.
• Performance: By optimizing queries dynamically, AQE can significantly improve performance and resource utilization.
• Efficiency: Helps in making better use of computing resources, leading to faster query execution times and reduced costs.

Dynamically Coalescing Shuffle Partitions

Here is a section to make you understand the core concepts:

Understanding Shuffle in Spark:

• Shuffle: A process that moves data across the network to redistribute it for downstream operations, such as joins or aggregations. It’s an expensive operation due to network I/O.
• Partitions: Dividing data into chunks for parallel processing. The number of partitions impacts performance.

Challenges with Shuffle Partitions:

Too Few Partitions:

• Large data chunks in each partition.
• Tasks may need to spill data to disk (especially during sorts or aggregations), slowing down the query.

Too Many Partitions:

• Small data chunks in each partition.
• Numerous small network data fetches, causing inefficient I/O and burdening the task scheduler.

Solution: Dynamic Partition Coalescing:

• Start with a large number of shuffle partitions.
• Combine small adjacent partitions into larger ones at runtime using shuffle file statistics.

Example:

• Query: SELECT max(i) FROM tbl GROUP BY j
• Initial shuffle partitions: 5 (based on initial setting).
• Without AQE: 5 tasks, even if 3 partitions are very small.
• With AQE: Combine 3 small partitions into 1, reducing the number of tasks from 5 to 3.

Dynamically Switching Join Strategies

Join Strategies in Spark:

• Broadcast Hash Join: Most performant if one side of the join fits in memory.
• Planned if the estimated size of a join relation is below the broadcast-size threshold.
• Challenges: Size estimation can be inaccurate due to selective filters or complex operations.

Source: Databricks

Solution: Dynamic Join Strategy Switching:

• Replans join strategy at runtime based on the accurate size of join relations.
• Example: Initial plan for a sort-merge join, but AQE finds one side is small enough to fit in memory.
• Converts to broadcast hash join for better performance.

Source: Databricks

Further Optimization:

• Convert regular shuffle to localized shuffle.
• Localized shuffle reads data on a per mapper basis instead of a per reducer basis, reducing network traffic.

Dynamically Optimizing Skew Joins

Understanding Data Skew:

• Data Skew: Uneven distribution of data among partitions.
• Severe skew affects performance, especially in joins.

Solution: Skew Join Optimization:

• Detects skew automatically from shuffle file statistics.
• Splits skewed partitions into smaller subpartitions for balanced processing.

Example:

Joining Table A and Table B.

• Table A has a significantly larger partition (A0) compared to others.
• Without AQE: 4 tasks running sort-merge join, one task takes much longer.
• With AQE: Split A0 into smaller parts, resulting in 5 balanced tasks.
• Each task is completed in roughly the same time, improving overall performance.

Enabling Adaptive Query Execution (AQE) in Apache Spark is straightforward and involves setting a few configuration options. Here’s how you can enable AQE:

Steps to Enable AQE in Spark

1. Set the Spark Configuration: You need to set the configuration options in your Spark application to enable AQE. This can be done either programmatically in your Spark code or via the Spark configuration file.

Programmatically in Spark Code:

from pyspark.sql import SparkSession

# Create a Spark session with AQE enabled
spark = SparkSession.builder \
    .appName("MyApp") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .config("spark.sql.adaptive.join.enabled", "true") \
    .getOrCreate()

# Your Spark code here

Via Spark Configuration File

If you’re using a configuration file (e.g., spark-defaults.conf), add the following lines:

spark.sql.adaptive.enabled=true
spark.sql.adaptive.coalescePartitions.enabled=true
spark.sql.adaptive.skewJoin.enabled=true
spark.sql.adaptive.join.enabled=true

Key Configuration Parameters:

• spark.sql.adaptive.enabled: Main switch to enable AQE (default is false).
• spark.sql.adaptive.coalescePartitions.enabled: Enable dynamic coalescing of shuffle partitions.
• spark.sql.adaptive.skewJoin.enabled: Enable dynamic skew join optimization.
• spark.sql.adaptive.join.enabled: Enable dynamic switching of join strategies.

Additional Optional Configurations

You might want to tweak additional settings to fine-tune AQE behavior based on your specific needs:

• spark.sql.adaptive.shuffle.targetPostShuffleInputSize: Target post-shuffle partition size (default is 64MB).
• spark.sql.adaptive.localShuffleReader.enabled: Enable localized shuffle read (default is true).
• spark.sql.adaptive.coalescePartitions.minPartitionNum: Minimum number of shuffle partitions after coalescing.
• spark.sql.adaptive.advisoryPartitionSizeInBytes: Advisory partition size in bytes (default is 64MB).

Example of Additional Configurations

spark = SparkSession.builder \
    .appName("MyApp") \
    .config("spark.sql.adaptive.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.enabled", "true") \
    .config("spark.sql.adaptive.skewJoin.enabled", "true") \
    .config("spark.sql.adaptive.join.enabled", "true") \
    .config("spark.sql.adaptive.shuffle.targetPostShuffleInputSize", "128MB") \
    .config("spark.sql.adaptive.localShuffleReader.enabled", "true") \
    .config("spark.sql.adaptive.coalescePartitions.minPartitionNum", "2") \
    .config("spark.sql.adaptive.advisoryPartitionSizeInBytes", "128MB") \
    .getOrCreate()

Conclusion

Adaptive Query Execution (AQE) in Spark 3.0 represents a significant advancement in query optimization by reducing reliance on static statistics. Traditional cost-based optimization in Spark often faced challenges balancing the overhead of stats collection with the accuracy of estimations.

Detailed statistics, like column histograms, are expensive to collect but necessary for improving selectivity and cardinality estimation and detecting data skew.

AQE effectively mitigates the need for such detailed statistics and the associated manual tuning effort. By dynamically adjusting query plans at runtime based on real-time data characteristics, AQE enhances query performance without requiring extensive prior knowledge of the data.

This adaptability makes SQL query optimization more robust, particularly in the presence of arbitrary UDFs and unpredictable dataset changes, such as sudden data size fluctuations or frequent data skew.

References

[embed]How to Speed up SQL Queries with Adaptive Query Execution Learn more about the new Spark 3.0 feature Adaptive Query Execution and how to use it to accelerate SQL query execution…www.databricks.com

[embed]Performance Tuning For some workloads, it is possible to improve performance by either caching data in memory or by turning on some…spark.apache.org