Memory, caching, and Adaptive Query Execution

6 h
requires SPK-07

From your world. Redshift WLM memory/spill-to-disk and query queues; diagnosing OOM

Make AQE coalescing visible in the plan

Adaptive Query Execution re-optimizes a query at runtime using real shuffle statistics. Its most common win is coalescing. It collapses many tiny post-shuffle partitions into fewer partitions of the right size. This turns the static spark.sql.shuffle.partitions count into an adaptive count. You will prove it. You will produce an aggregation whose final executed plan shows the AQE coalescing.

Task

Implement coalesced_aggregate(spark, order_items). It must do these three things:

Enable AQE and shuffle-partition coalescing for this query.
Set a high initial spark.sql.shuffle.partitions (so there are small partitions to coalesce). Also set spark.sql.adaptive.coalescePartitions.minPartitionNum = 1 (so coalescing can reduce the count to a few partitions at this tiny scale).
Return the per-company sum of line_amount as columns company_id, total.

It is good practice to restore the session configs afterward, but it is not required.

How it is graded

Data (DuckDB oracle): equals SELECT company_id, sum(line_amount) FROM order_items GROUP BY company_id (float tolerance).
Plan — AQE on: after an action, the final plan contains AdaptiveSparkPlan with isFinalPlan=true and a coalesced AQEShuffleRead. This proves that AQE coalesced the shuffle partitions. A solution that leaves AQE off has no such node.

Starter code

solution.py

"""SPK-08 / 02-aqe-coalesce — starter. See README.md."""


def coalesced_aggregate(spark, order_items):
    # TODO: enable AQE + coalescePartitions, set a high shuffle.partitions and
    # minPartitionNum=1, then return per-company sum(line_amount) as
    # company_id, total. The final plan must show AQE coalescing.
    raise NotImplementedError

Hints

Hint 1

Enable AQE for this query: spark.conf.set('spark.sql.adaptive.enabled','true') and 'spark.sql.adaptive.coalescePartitions.enabled','true'. Set a high initial 'spark.sql.shuffle.partitions' (e.g. 64) so there are small post-shuffle partitions for AQE to coalesce.

Hint 2

Return the per-company sum of line_amount. The test runs an action and reads the FINAL plan: it must contain AdaptiveSparkPlan (isFinalPlan=true) and a coalesced AQEShuffleRead.

Hint 3

Also set 'spark.sql.adaptive.coalescePartitions.minPartitionNum','1' so coalescing can reduce to few partitions at tiny scale. Return columns company_id, total.

Run it locally

uv run de start SPK-08/02-aqe-coalesce
# edit workspace/SPK-08/02-aqe-coalesce/
uv run de test