How aggregations execute

Partial and final aggregation

HashAggregate(keys=[country], functions=[sum(amount)])              <- final: merge partial sums
+- Exchange hashpartitioning(country, 200)                           <- shuffle partial results
   +- HashAggregate(keys=[country], functions=[partial_sum(amount)]) <- partial: per task
      +- FileScan parquet orders [country,amount]

If a tasktask: The work for one partition in one stage, run on one CPU core. Learn more → holds a million orders from 50 countries, the partial aggregation sends 50 rows to the shuffleshuffle: Moving rows between machines so that all rows with the same key end up together. Needed by joins, groupBy and sorting, and usually the most expensive step of a job. Learn more → instead of a million. This is the same idea as a MapReduce combiner, and it is why a groupBy on a hot key is usually far less painful than a join on one: the hot key is collapsed to one row per task before the shuffle.

Partial aggregation only helps when each task sees repeated keys. Grouping by a nearly unique column (an order id) does the hashing work twice and shrinks nothing.

The three aggregation operators

Distinct aggregates

A COUNT(DISTINCT user_id) per country cannot be pre-aggregated into a single number, because the same user can appear in many tasks. Spark rewrites it into two levels of aggregation:

1. 1
   Aggregate by (country, user_id), partial then final, with a shuffle on both columns. This deduplicates users.
2. 2
   Aggregate the result by country, counting rows, with a second shuffle.

With several distinct aggregates on different columns, such as COUNT(DISTINCT user_id) and COUNT(DISTINCT product_id) together, Spark adds an Expand operator that copies every input row once per distinct group before the first shuffle, so the data shuffled grows with the number of distinct columns.

orders.groupBy("country").agg(F.countDistinct("user_id").alias("users"))
orders.groupBy("country").agg(F.approx_count_distinct("user_id", 0.02).alias("users"))

SELECT country, COUNT(DISTINCT user_id) AS users FROM orders GROUP BY country;
SELECT country, approx_count_distinct(user_id, 0.02) AS users FROM orders GROUP BY country;

approx_count_distinct uses HyperLogLog++ sketches, which merge like ordinary partial aggregates: one shuffle, fixed memory per group, and a relative standard error you choose (5% by default).

Spotting trouble in a plan

• SortAggregate on a large input: an extra sort; see if a rewrite gets a HashAggregate.
• Expand under an aggregation: several distinct aggregates multiplying rows.
• Spill in the aggregation stagestage: A group of steps Spark can run without moving data between machines. A new stage starts at every shuffle. Learn more →: too many keys per task for memory; more shuffle partitionspartition: A chunk of a DataFrame's rows. Spark processes each partition as one task, so partitions decide how much work runs in parallel. Learn more → help.
• collect_list on a hot key: object hash aggregation cannot shrink it; one huge array lands in one task.

Common mistakes

• ✗Assuming every groupBy is equally cheap — Operator choice, distinct aggregates and object buffers change the cost a lot.
• ✗Many COUNT(DISTINCT) on different columns in one query — The Expand operator multiplies the shuffled rows.
• ✗Using exact distinct counts for dashboards — approx_count_distinct is one shuffle and fixed memory, with a chosen error.
• ✗Salting a groupBy with only sum and count — Partial aggregation already handles hot keys for these; salting helps joins and non-decomposable aggregates more.