Grouping and Aggregation in PySpark
Group, aggregate, and pivot — then read the query plan to see the shuffle that groupBy introduces and why it dominates the job's runtime.
Aggregation is where Spark starts moving data between machines, and where job runtimes start being decided. This lesson connects the API to what actually happens.
The data
from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.appName("agg").master("local[*]").getOrCreate()
sales = spark.createDataFrame(
[
("UK", "electronics", "alice", 1200.00, "2026-01-15"),
("UK", "books", "alice", 45.50, "2026-01-16"),
("UK", "electronics", "bob", 890.00, "2026-01-16"),
("US", "electronics", "carol", 2100.00, "2026-01-17"),
("US", "books", "carol", 32.00, "2026-01-17"),
("US", "clothing", "dave", 150.75, "2026-01-18"),
("DE", "electronics", "erin", 670.00, "2026-01-18"),
("DE", "books", "erin", None, "2026-01-19"),
],
["country", "category", "customer", "amount", "sale_date"],
)
sales.show()+-------+-----------+--------+------+----------+
|country| category|customer|amount| sale_date|
+-------+-----------+--------+------+----------+
| UK|electronics| alice|1200.0|2026-01-15|
| UK| books| alice| 45.5|2026-01-16|
| UK|electronics| bob| 890.0|2026-01-16|
| US|electronics| carol|2100.0|2026-01-17|
| US| books| carol| 32.0|2026-01-17|
| US| clothing| dave|150.75|2026-01-18|
| DE|electronics| erin| 670.0|2026-01-18|
| DE| books| erin| NULL|2026-01-19|
+-------+-----------+--------+------+----------+Note the null in the last row — it matters shortly.
Grouping
sales.groupBy("country").agg(
F.count("*").alias("sales"),
F.round(F.sum("amount"), 2).alias("revenue"),
F.round(F.avg("amount"), 2).alias("avg_sale"),
F.max("amount").alias("largest"),
).orderBy(F.desc("revenue")).show()+-------+-----+-------+--------+-------+
|country|sales|revenue|avg_sale|largest|
+-------+-----+-------+--------+-------+
| US| 3|2282.75| 760.92| 2100.0|
| UK| 3| 2135.5| 711.83| 1200.0|
| DE| 2| 670.0| 670.0| 670.0|
+-------+-----+-------+--------+-------+Germany shows 2 sales but an average of 670 — the null was skipped by avg and sum but
counted by count("*"). That is correct SQL behaviour and also a reporting bug waiting to
happen.
The count trap
sales.groupBy("country").agg(
F.count("*").alias("rows"),
F.count("amount").alias("non_null_amounts"),
F.countDistinct("customer").alias("customers"),
).show()+-------+----+----------------+---------+
|country|rows|non_null_amounts|customers|
+-------+----+----------------+---------+
| UK| 3| 3| 2|
| US| 3| 3| 2|
| DE| 2| 1| 1|
+-------+----+----------------+---------+count("*") gives 2 for Germany; count("amount") gives 1. Whenever a report’s totals do
not reconcile, this is the first thing to check.
Multiple grouping keys
sales.groupBy("country", "category").agg(
F.round(F.sum("amount"), 2).alias("revenue")
).orderBy("country", "category").show()+-------+-----------+-------+
|country| category|revenue|
+-------+-----------+-------+
| DE| books| NULL|
| DE|electronics| 670.0|
| UK| books| 45.5|
| UK|electronics| 2090.0|
| US| books| 32.0|
| US| clothing| 150.75|
| US|electronics| 2100.0|
+-------+-----------+-------+sum of an all-null group is null, not zero. Coalesce if zero is the answer you want:
sales.groupBy("country", "category").agg(
F.coalesce(F.round(F.sum("amount"), 2), F.lit(0.0)).alias("revenue")
).orderBy("country", "category").show(2)+-------+-----------+-------+
|country| category|revenue|
+-------+-----------+-------+
| DE| books| 0.0|
| DE|electronics| 670.0|
+-------+-----------+-------+Pivoting
sales.groupBy("country").pivot("category").agg(
F.coalesce(F.round(F.sum("amount"), 2), F.lit(0.0))
).show()+-------+-----+--------+-----------+
|country|books|clothing|electronics|
+-------+-----+--------+-----------+
| DE| 0.0| 0.0| 670.0|
| UK| 45.5| 0.0| 2090.0|
| US| 32.0| 150.75| 2100.0|
+-------+-----+--------+-----------+pivot has to know the distinct values before it can build the columns, so without a list it
runs an extra job to find them. On wide data, supply the values:
sales.groupBy("country").pivot("category", ["books", "electronics"]).agg(
F.sum("amount")
).show()+-------+-----+-----------+
|country|books|electronics|
+-------+-----+-----------+
| DE| NULL| 670.0|
| UK| 45.5| 2090.0|
| US| 32.0| 2100.0|
+-------+-----+-----------+One job instead of two, and clothing is excluded because you did not ask for it.
Seeing the shuffle
sales.groupBy("country").agg(F.sum("amount")).explain()== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[country#0], functions=[sum(amount#3)])
+- Exchange hashpartitioning(country#0, 200), ENSURE_REQUIREMENTS
+- HashAggregate(keys=[country#0], functions=[partial_sum(amount#3)])
+- Scan ExistingRDD[country#0,category#1,customer#2,amount#3,sale_date#4]Exchange hashpartitioning is the shuffle — data written to disk, sent across the network,
and re-read. It is by far the most expensive line in most plans.
Notice Spark aggregates twice. partial_sum runs on each partition before the shuffle, so
only one partial total per key per partition crosses the network rather than every row. This
is why sum is cheap and countDistinct is not — distinct counts cannot be partially
combined the same way.
The 200 is spark.sql.shuffle.partitions. On this eight-row DataFrame that means 200 tasks
for three groups:
print("partitions after groupBy:",
sales.groupBy("country").agg(F.sum("amount")).rdd.getNumPartitions())partitions after groupBy: 200Adaptive Query Execution coalesces these at runtime in Spark 3.x — that is what
isFinalPlan=false means. Confirm it by looking at the plan after execution:
agg = sales.groupBy("country").agg(F.sum("amount"))
agg.collect()
agg.explain()== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=true
+- == Final Plan ==
*(2) HashAggregate(keys=[country#0], functions=[sum(amount#3)])
+- AQEShuffleRead coalesced
+- ShuffleQueryStage 0AQEShuffleRead coalesced — 200 partitions became one. AQE is on by default and is the main
reason Spark 3 needs far less manual tuning than Spark 2.
Aggregating without grouping
sales.agg(
F.count("*").alias("rows"),
F.round(F.sum("amount"), 2).alias("total"),
F.countDistinct("customer").alias("customers"),
).show()+----+-------+---------+
|rows| total|customers|
+----+-------+---------+
| 8|5088.25| 5|
+----+-------+---------+One row for the whole DataFrame. This still shuffles — everything has to reach one place.
Collecting values per group
sales.groupBy("country").agg(
F.collect_set("customer").alias("customers"),
F.collect_list("category").alias("categories"),
).show(truncate=False)+-------+--------------+---------------------------------+
|country|customers |categories |
+-------+--------------+---------------------------------+
|DE |[erin] |[electronics, books] |
|UK |[bob, alice] |[electronics, books, electronics]|
|US |[carol, dave] |[electronics, books, clothing] |
+-------+--------------+---------------------------------+collect_set deduplicates, collect_list keeps everything including order. Both build the
whole array in memory on one executor, so a group with millions of rows will run that executor
out of memory. Use them for small groups only.
Practice
1. Find each country's highest-spending customer.
sales.groupBy("country", "customer").agg(F.sum("amount").alias("spent")) \
.groupBy("country").agg(F.max(F.struct("spent", "customer")).alias("top")) \
.select("country", "top.customer", "top.spent").show()+-------+--------+------+
|country|customer| spent|
+-------+--------+------+
| DE| erin| 670.0|
| UK| alice|1245.5|
| US| carol|2132.0|
+-------+--------+------+max over a struct compares by the first field, so this returns the customer attached to the
maximum. A window function does the same thing more readably — that is the next lesson.
2. Why does avg differ from sum / count("*") for Germany?
avg: 670.0
sum / count("*"): 335.0avg divides by the count of non-null values (1), while count("*") counts rows (2). Both
are defensible; they answer different questions. Decide deliberately whether a missing amount
means “zero” or “unknown”, and write the aggregation to match.
3. Set spark.sql.shuffle.partitions to 4 and re-check the partition count.
spark.conf.set("spark.sql.shuffle.partitions", 4)
print(sales.groupBy("country").agg(F.sum("amount")).rdd.getNumPartitions())4Lowering it removes scheduling overhead on small data. On large data, too low a value gives tasks that will not fit in executor memory. AQE handles most of this now — tune it manually only when you have measured a problem.
4. Compare the plans for count("customer") and countDistinct("customer").
count shows partial_count before the exchange; countDistinct shows an extra
Expand/Aggregate stage and often a second shuffle, because distinct values cannot be
partially combined — every value must reach one place to be deduplicated. On high-cardinality
columns use approx_count_distinct, which uses HyperLogLog and is orders of magnitude cheaper
for a ~2% error.
Next: joining DataFrames, and how to avoid shuffling both sides.