Skip to main content

04.4 PySpark Aggregations

PySpark aggregation architecture

Architect goal: Understand groupBy().agg(), distributed aggregation, partial reduction, shuffle, cardinality, skew, exact versus approximate metrics, and production aggregation design.

1. Why aggregation matters​

Aggregation converts many input rows into fewer analytical rows: revenue by customer, orders by country, daily sales, average flight delay, distinct customers, and top products.

from pyspark.sql import functions as F

result = (
    df
    .groupBy("country")
    .agg(
        F.sum("amount").alias("total_amount"),
        F.count("*").alias("order_count"),
        F.avg("amount").alias("avg_amount")
    )
)

Current PySpark groupBy() returns GroupedData and supports column names and expressions; current Spark also supports column ordinals. citeturn482604search2

2. Aggregation is distributed​

Conceptually:

Input partitions
      ↓
Local / partial aggregation
      ↓
Shuffle by grouping key
      ↓
Final aggregation
      ↓
Result

Grouping requires rows belonging to the same key to meet. A grouped aggregation therefore commonly contains an Exchange/shuffle boundary.

3. Group cardinality is critical​

One billion rows grouped by country may produce a few hundred groups. One billion rows grouped by a highly unique event_id can produce hundreds of millions of groups.

Ask:

  1. How many input rows?
  2. How many distinct grouping keys?
  3. How evenly are keys distributed?
  4. How large is aggregation state?
  5. Does the business require exact results?

4. Core aggregate functions​

Common functions include sum, count, countDistinct, avg, min, max, first, last, collect_list, collect_set, approx_count_distinct, approx_percentile, and approx_top_k. Databricks currently documents these and additional aggregate functions. citeturn482604search4turn482604search9

5. COUNT semantics​

These are different:

F.count("*")
F.count("customer_id")
F.countDistinct("customer_id")

count(*) counts rows, count(column) counts non-NULL values, and count distinct counts unique non-NULL values. citeturn482604search1

6. Multiple metrics together​

result = df.groupBy("country").agg(
    F.count("*").alias("orders"),
    F.sum("amount").alias("revenue"),
    F.avg("amount").alias("avg_order_value"),
    F.min("amount").alias("min_order"),
    F.max("amount").alias("max_order")
)

Compute compatible metrics together rather than repeatedly rebuilding the same dataset.

7. Conditional aggregation​

result = df.groupBy("country").agg(
    F.sum(F.when(F.col("status") == "COMPLETED", 1).otherwise(0)).alias("completed_orders"),
    F.sum(F.when(F.col("status") == "CANCELLED", 1).otherwise(0)).alias("cancelled_orders")
)

8. Aggregating expressions​

df.groupBy("country").agg(
    F.sum(F.col("quantity") * F.col("unit_price")).alias("gross_sales")
)

9. Multiple grouping dimensions and grain​

daily_sales = (
    df.groupBy("country", "order_date")
      .agg(F.sum("amount").alias("revenue"), F.count("*").alias("orders"))
)

Before writing groupBy, state: one output row represents what?

Most aggregation defects are actually grain-definition defects.

10. Global aggregation​

df.groupBy().agg(
    F.sum("amount").alias("total_revenue"),
    F.count("*").alias("total_orders")
)

Useful for control totals, reconciliation, and pipeline audit metrics.

11. NULL behavior​

Aggregates generally ignore NULL input values for numeric expressions. For 100, NULL, 200: SUM is 300, AVG is 150, count(amount) is 2, while count(*) is 3.

Do not blindly convert NULL to zero; missing and zero may have different business meanings.

12. Aggregation after a join​

A duplicate key in an upstream dimension can multiply fact rows before aggregation and inflate SUM.

Before aggregating after a join:

  • Validate join-key uniqueness.
  • Understand join cardinality.
  • Check row counts before and after the join.
  • Reconcile totals against a trusted source.

13. Partial aggregation​

Partition 1: A=10, A=20, B=5
Partition 2: A=7, B=8

Local combine:
P1 → A=30, B=5
P2 → A=7, B=8

Shuffle:
A → 30 + 7
B → 5 + 8

Final:
A=37
B=13

Spark can reduce rows locally before shuffle. It is therefore inaccurate to think every original input row must cross the network.

14. Exchange and physical plan​

Run:

result.explain("formatted")

A typical grouped aggregation has a conceptual shape:

Scan
  ↓
Partial HashAggregate
  ↓
Exchange
  ↓
Final HashAggregate

Exact operators vary with Spark version, runtime, source, statistics, and query shape.

15. High-cardinality grouping​

Grouping on a nearly unique key can create huge aggregation state and reduce the benefit of local combining.

Symptoms can include:

  • Memory pressure
  • Spill
  • Long tasks
  • Large shuffle volume
  • Garbage collection pressure

16. Aggregation skew​

If one key owns a disproportionate share of rows, one shuffle partition can become much larger than others.

Symptoms:

  • One or a few tasks run much longer.
  • Spill is concentrated in certain tasks.
  • Stage completion waits for stragglers.
  • Executor pressure is uneven.

Skew is a data-distribution problem.

17. Exact distinct​

df.groupBy("country").agg(
    F.countDistinct("customer_id").alias("unique_customers")
)

Exact distinct can be substantially more expensive than additive metrics because uniqueness must be tracked.

18. Approximate distinct​

df.groupBy("country").agg(
    F.approx_count_distinct("customer_id").alias("estimated_customers")
)

Databricks documents approx_count_distinct as an HLL++-based estimate; its default relative standard deviation corresponds to roughly 5% expected error and can be configured. citeturn482604search0

Use approximation for appropriate dashboards, exploration, profiling, and capacity planning—not when an exact metric is required.

19. Approximate percentiles and top-K​

Databricks provides approx_percentile and approx_top_k for large-scale approximate analytics. citeturn482604search3turn482604search6

Typical uses include p95/p99 latency, distribution analysis, most frequent error codes, and top event types.

20. Pivot​

pivoted = (
    df.groupBy("country")
      .pivot("quarter")
      .agg(F.sum("amount"))
)

Pivot changes row values into columns. It can become expensive when the pivot domain is large; bounded, explicitly known values are preferable for production designs.

21. collect_list and collect_set​

df.groupBy("customer_id").agg(
    F.collect_list("product_id").alias("products")
)

Large groups can create very large in-memory aggregation state. Do not use these as a substitute for a normalized child dataset when group sizes can become huge.

22. Aggregation and BI design​

Detailed Silver facts
        ↓
Business aggregation
        ↓
Gold analytical table
        ↓
BI / SQL / semantic layer

Example Gold grain: one row per route + operating date, with measures such as departures, passengers, revenue, average delay, and cancellation rate.

23. Reduce before expensive work​

filtered = (
    df.select("country", "amount", "status")
      .filter(F.col("status") == "COMPLETED")
)

result = filtered.groupBy("country").agg(
    F.sum("amount").alias("revenue")
)

Projection and filtering can reduce the data that reaches the aggregation stage.

24. Avoid premature repartitioning​

Do not automatically call repartition(key) before groupBy(key). The aggregation itself has distribution requirements.

First inspect the physical plan and runtime metrics. Repartitioning is a deliberate tuning decision, not a ritual.

25. Production Gold example​

gold = (
    spark.table("main.silver.orders")
    .select("order_date", "country", "customer_id", "amount", "status")
    .filter(F.col("status") == "COMPLETED")
    .groupBy("order_date", "country")
    .agg(
        F.sum("amount").alias("revenue"),
        F.count("*").alias("orders"),
        F.countDistinct("customer_id").alias("unique_customers"),
        F.avg("amount").alias("avg_order_value")
    )
)

Output grain: one row per order_date + country.

Put that grain in the technical specification.

26. Production monitoring​

For a large aggregation inspect:

  1. Input row count.
  2. Distinct grouping-key count.
  3. Stage duration.
  4. Shuffle read/write.
  5. Spill.
  6. Task-duration distribution.
  7. Retries.
  8. Executor memory pressure.
  9. Output row count.

A slow aggregation may be caused by skew, excessive cardinality, too much input, an upstream many-to-many join, poor partition distribution, or large aggregation state—not by the aggregate function itself.

27. Hands-on lab — customer revenue​

Create an orders DataFrame with order_id, customer_id, country, order_date, amount, and status.

Build progressively:

  1. Revenue by country.
  2. Revenue + order count + average order value.
  3. Revenue by country and month.
  4. Exact unique customers.
  5. Approximate unique customers.

Compare results and plans.

28. Hands-on lab — execution​

result = df.groupBy("country").agg(
    F.sum("amount").alias("revenue")
)
result.explain("formatted")

Identify the scan, partial aggregation, Exchange, and final aggregation. Then use the Spark UI to locate the shuffle stage.

29. Hands-on lab — skew​

Create deliberately skewed data where one key owns most rows.

Compare it with balanced data and inspect:

  • Task duration
  • Shuffle distribution
  • Spill
  • Stragglers

The objective is to recognize skew from runtime evidence rather than guess.

30. Architect decision framework​

RequirementPreferred approach
Exact totalSUM
Exact row countCOUNT(*)
Exact unique customersCOUNT DISTINCT
Large-scale approximate cardinalityAPPROX COUNT DISTINCT
Percentile at scaleAPPROX PERCENTILE
Top frequent values at scaleAPPROX TOP K
Small bounded pivotPIVOT
Huge per-group listReconsider COLLECT_LIST
Financial reportingExact + reconciliation
Repeated dashboard aggregationConsider pre-aggregation/materialization

Databricks distinguishes additive aggregates such as SUM and COUNT from non-additive measures such as COUNT DISTINCT for rollup/materialization decisions. citeturn482604search13

31. Interview questions​

  1. What happens when groupBy() executes?
  2. What is GroupedData?
  3. COUNT(*) vs COUNT(column)?
  4. Why can groupBy cause a shuffle?
  5. What is partial aggregation?
  6. Why is high-cardinality grouping expensive?
  7. How does skew affect aggregation?
  8. When would you use approximate aggregation?
  9. How would you design a 50 TB daily aggregate?
  10. How would you validate a Gold aggregate?
  11. How do additive and non-additive measures affect pre-aggregation?

32. Cheat sheet​

from pyspark.sql import functions as F

df.groupBy("country").agg(
    F.sum("amount").alias("revenue"),
    F.count("*").alias("orders"),
    F.avg("amount").alias("avg_amount")
)

F.countDistinct("customer_id")
F.approx_count_distinct("customer_id")

F.sum(
    F.when(F.col("status") == "COMPLETED", 1).otherwise(0)
)

df.groupBy("country").pivot("quarter").agg(
    F.sum("amount")
)

result.explain("formatted")

33. Key takeaways​

  1. Aggregation changes dataset grain.
  2. Group cardinality matters as much as input size.
  3. Distributed aggregation commonly involves shuffle.
  4. Partial aggregation can reduce data before shuffle.
  5. COUNT variants have different NULL semantics.
  6. Skew can create aggregation stragglers.
  7. Exact and approximate distinct solve different business problems.
  8. Large collect_list/collect_set operations can create huge state.
  9. Always define output grain before writing groupBy.
  10. Enterprise aggregation requires correctness, reconciliation, and performance evidence.

Next​

04.5 PySpark Joins — join types, cardinality, broadcast joins, shuffle joins, semi/anti joins, NULL-safe joins, skew, duplicate explosion, and enterprise join troubleshooting.

Official references​