Adaptive Query Execution & Cost-Based Optimization

Catalyst plans a query before it runs. It guesses sizes, and a guess can be wrong. Adaptive Query Execution, or AQE, fixes the plan while the query runs.

You met AQE in earlier chapters, for example "Joins in Spark" and "Data Skew & Salting". This chapter shows how it works, and how to read what it did.

You will learn five things.

1. How AQE uses real sizes at the border between two stages.
2. Three features: merge small partitions, switch the join, split skewed partitions.
3. What AQE cannot fix.
4. How to read isFinalPlan=true and the Initial Plan.
5. How the cost-based optimizer (CBO) and AQE work together, and which setting fits which symptom.

We use the shop again. The paid revenue by customer is Asha 65.00, Carla 15.50 and Dev 80.00.

A shuffle moves rows between machines, so each key lands in one place. You read about it in "The Shuffle: Spark's Performance Killer". A shuffle splits a query into stages.

A stage is a group of tasks that run without a shuffle.

When a stage ends, Spark has written its shuffle files. At that moment it knows the real size of every partition. AQE uses these facts.

Spark runs a query with AQE in four steps.

1. Spark cuts the plan at each shuffle. Each piece is a query stage.
2. Spark runs every stage that does not wait for another stage.
3. When a stage finishes, Spark reads its real sizes, in bytes and rows.
4. Spark re-optimizes the rest of the plan with those sizes. Then it runs the next stages.

Catalyst has only estimates when it plans. AQE has measurements, but only after a shuffle has run. So AQE can change the rest of the plan, not the part that already ran.

AQE is on by default since Spark 3.2. Its main setting is spark.sql.adaptive.enabled. Check your version if you are not sure.

AQE has many parts. Three matter most. Each one has a setting and a default.

The next three sections show one feature each. Each one shows a simplified plan from explain().

Settings and defaults come from Spark 3.2 and later. Names change a little between versions, so check yours.

After a shuffle, Spark makes spark.sql.shuffle.partitions partitions. The default is 200. On small data, most of them are empty or tiny.

AQE merges small partitions that sit next to each other, for example partitions 0, 1 and 2. This is called coalescing.

It merges up to a target size. The target is at most spark.sql.adaptive.advisoryPartitionSizeInBytes, 64 MB by default.

But spark.sql.adaptive.coalescePartitions.parallelismFirst is true by default. Then Spark aims for about one partition per core, with a target size of at least 1 MB.

Look at paid revenue by customer. Only customers 1, 3 and 4 have paid orders.

df = orders.filter("status = 'paid'").groupBy("customer_id").sum("amount")
df.explain()      # before the run
df.collect()
df.explain()      # after the run: same df

# the first explain() prints
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[customer_id], functions=[sum(amount)])
   +- Exchange hashpartitioning(customer_id, 200)
      +- HashAggregate(keys=[customer_id], functions=[partial_sum(amount)])
         +- Filter (status = paid)
            +- FileScan parquet orders

# the second explain() prints
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
   *(2) HashAggregate(keys=[customer_id], functions=[sum(amount)])
   +- AQEShuffleRead coalesced
      +- ShuffleQueryStage 0
         +- Exchange hashpartitioning(customer_id, 200)
            +- *(1) HashAggregate(keys=[customer_id], functions=[partial_sum(amount)])
               +- *(1) Filter (status = paid)
                  +- FileScan parquet orders
+- == Initial Plan ==
   (the same tree as before the run)

The Exchange line still says 200. That is the plan Spark made before the run. AQEShuffleRead sits above it and merges the small pieces when it reads them.

Merging also means fewer small output files. A groupBy followed by a write makes fewer, bigger files.

A sort-merge join shuffles both tables. A broadcast join copies the small table to every executor.

In "Joins in Spark", you saw AQE switch this join after a filter. Here you read the plan that shows the switch.

Spark picks the join before the run, from size estimates. After a filter, the real size can be much smaller than the estimate. AQE sees the real size after the shuffle and can switch to a broadcast join.

Join orders to the customers from India (country = IN).

df = orders.join(customers.filter("country = 'IN'"), "customer_id") \
           .select("order_id", "name", "amount")
df.explain()      # before the run
df.collect()
df.explain()      # after the run: same df

# the first explain() prints
AdaptiveSparkPlan isFinalPlan=false
+- Project [order_id, name, amount]
   +- SortMergeJoin [customer_id], [customer_id], Inner
      :- Sort
      :  +- Exchange hashpartitioning(customer_id, 200)
      :     +- FileScan parquet orders
      +- Sort
         +- Exchange hashpartitioning(customer_id, 200)
            +- Filter (country = IN)
               +- FileScan parquet customers

# the second explain() prints (Final Plan only)
+- Project [order_id, name, amount]
   +- BroadcastHashJoin [customer_id], [customer_id], Inner, BuildRight
      :- AQEShuffleRead local
      :  +- ShuffleQueryStage 0
      :     +- Exchange hashpartitioning(customer_id, 200)
      :        +- FileScan parquet orders
      +- BroadcastQueryStage 2
         +- BroadcastExchange
            +- AQEShuffleRead local
               +- ShuffleQueryStage 1
                  +- Exchange hashpartitioning(customer_id, 200)
                     +- Filter (country = IN)
                        +- FileScan parquet customers

The join line changed from SortMergeJoin to BroadcastHashJoin. The Sort lines are gone.

BuildRight means the right side, customers, is the copied table.

Both Exchange lines are still there. Both shuffles already ran when AQE saw the sizes. So the saving is partial: Spark skips the sorts and the network fetch of the big side.

AQEShuffleRead local is the local shuffle reader. It reads shuffle files on the same machine, so rows do not cross the network. It is on by default (spark.sql.adaptive.localShuffleReader.enabled).

You met skew in "Data Skew & Salting". One partition is much bigger than the others.

AQE splits it when it is over 5 times the median and over 256 MB. The median is the middle partition size.

You see it in the final plan. The join line gets a marker, and the read line says skewed.

# explain() AFTER the run: Final Plan only
+- SortMergeJoin(skew=true) [customer_id], [customer_id], Inner
   :- Sort [customer_id ASC NULLS FIRST]
   :  +- AQEShuffleRead skewed
   :     +- ShuffleQueryStage 0
   +- Sort [customer_id ASC NULLS FIRST]
      +- AQEShuffleRead coalesced
         +- ShuffleQueryStage 1

No marker means AQE did not split anything.

Maybe no partition passed both limits. Maybe the join was a broadcast join. Maybe the join type cannot be split, such as a full outer join.

The chapter "Data Skew & Salting" lists the settings and the cases AQE misses.

Before Spark 3.2, the plan says CustomShuffleReader instead of AQEShuffleRead.

AQE learns only from finished shuffles. That sets three limits.

1. The first stage is planned from estimates. AQE cannot fix a bad scan or a bad filter placed before any shuffle.
2. A plan with no shuffle has no query stages, so AQE has nothing to measure.
3. AQE cannot see inside a Python UDF. A slow UDF stays slow.

Look at the second limit. A filter and a select on orders have no Exchange.

orders.filter("status = 'paid'").select("order_id", "amount").explain()

*(1) Project [order_id, amount]
+- *(1) Filter (status = paid)
   +- FileScan parquet orders

The plan you print is not always the plan that ran. Before the run, explain() shows isFinalPlan=false. That plan can still change.

After df.collect(), call df.explain() again. Now it shows isFinalPlan=true. The plan has two parts, as the earlier example showed.

Use collect() on the same DataFrame. count() and show() run a new plan, so df.explain() still shows isFinalPlan=false.

Compare the two parts line by line. A new BroadcastHashJoin, a coalesced read or skew=true tells you what AQE did.

Match the scan lines first. Then match the Exchange lines.

The Spark UI shows the same in the SQL tab. Open the query.

The graph shows the final plan when the query ends. While it runs, the graph can change.

You met the cost-based optimizer (CBO) in "Spark SQL & the Catalyst Optimizer". Statistics are facts about a table, such as the row count. CBO needs them from ANALYZE TABLE, and spark.sql.cbo.enabled is false by default.

The two work at different times.

CBO still matters for a query with many joins. AQE does not reorder joins. Only CBO picks the order by cost, before the run.

To use it, collect statistics and turn on two settings.

# orders and customers must be saved tables, not temp views
spark.sql("ANALYZE TABLE orders COMPUTE STATISTICS FOR COLUMNS customer_id, status")
spark.sql("ANALYZE TABLE customers COMPUTE STATISTICS FOR COLUMNS customer_id, country")
spark.conf.set("spark.sql.cbo.enabled", "true")
spark.conf.set("spark.sql.cbo.joinReorder.enabled", "true")   # false by default

Dynamic partition pruning is a third helper. Here, a partition means one folder of a table, not a shuffle partition.

Some tables are stored in folders, one folder for each value of a column. If the big table uses the join key for its folders, Spark can skip folders. It reads only the folders that the filtered small table needs.

By default, this works only when Spark copies the small side in a broadcast join. The setting is spark.sql.optimizer.dynamicPartitionPruning.enabled, true by default.

Start from the symptom, not from the setting. Find your symptom in the left column.

Change one setting at a time. Then compare the Final Plan with the Initial Plan.

Four points to keep.

• AQE reads real sizes at each shuffle and re-optimizes the rest of the plan.
• It merges small partitions, switches to a broadcast join, and splits skewed partitions.
• It cannot fix a bad first plan, a plan with no shuffle, or a slow Python UDF.
• To see the Final Plan, call explain() again on the same DataFrame after collect().

Answer these without looking back.

1. Why does a plan with no Exchange have no AdaptiveSparkPlan line?
2. In the join example, why is the saving only partial after the switch?
3. Which line in the Final Plan shows that partitions were merged?
4. Which optimizer can reorder a join of five tables, and what does it need first?