groupBy() looks simple in PySpark, but Aggregations scale and can trigger one of Spark’s most expensive operations: a shuffle.

If you’re preparing for PySpark interviews or working with Spark in production, follow me for practical Data Engineering, AI/ML, and PySpark content.

📩 Subscribe by email to get new tutorials and interview questions delivered straight to your inbox.

Now let’s understand what really happens when you write:

df.groupBy("department").count()

It looks like a single line of code about Aggregations.

But behind that line, Spark may need to:

Read Data
   ↓
Partition Data
   ↓
Shuffle Records
   ↓
Group Same Keys
   ↓
Aggregate
   ↓
Return / Write Result

Understanding this execution model is much more valuable than simply memorizing aggregation syntax.

📚 Complete Apache Spark Tutorial Progress

APACHE SPARK LEARNING PATH

1. ✓ What is Apache Spark?
2. ✓ Why Spark is Faster than Hadoop
3. ✓ Spark Architecture
4. ✓ Driver vs Executor
5. ✓ Cluster Managers
6. → RDD vs DataFrame vs Dataset
7. → SparkSession
8. ✓ Reading Data
9. ✓ Lazy Evaluation
10. ✓ DAG
11. ✓ Stages and Tasks
12. ✓ Shuffle
13. ✓ Narrow vs Wide Transformations
14. → Writing Data
15. → Data Transformations
16. →Window Functions
17. Aggregations

What Is an Aggregation?

Aggregations combine multiple rows and produce a summarized result.

For example, suppose we have:

+------+-----------+------+
| name |department |salary|
+------+-----------+------+
|Alice |Engineering|120000|
|Bob   |Engineering|95000 |
|Carol |Finance    |110000|
|David |Finance    |85000 |
|Eva   |HR         |75000 |
+------+-----------+------+

We might want to calculate:

• number of employees per department
• average salary
• maximum salary
• minimum salary
• total salary
• number of distinct employees

For example:

Engineering → 2 employees
Finance     → 2 employees
HR          → 1 employee

This is an aggregation.

1. The Most Common Aggregation Pattern

The most common PySpark pattern is:

df.groupBy("department").count()

Result:

+-----------+-----+
|department |count|
+-----------+-----+
|Engineering|2    |
|Finance    |2    |
|HR         |1    |
+-----------+-----+

Conceptually:

groupBy()
   ↓
Identify grouping key
   ↓
Bring same keys together
   ↓
Calculate aggregate

2. groupBy() Alone Doesn’t Calculate Anything

This is an important Spark concept.

Consider:

grouped = df.groupBy("department")

This doesn’t give you the final aggregation.

You still need an aggregation operation:

grouped.count()

or:

grouped.sum("salary")

or:

grouped.avg("salary")

The pattern is:

groupBy()
    ↓
Define groups
    ↓
Aggregation function
    ↓
Calculate result

3. Counting Rows

The simplest aggregation is count().

df.groupBy("department").count()

This answers:

How many records exist for each department?

Example:

Engineering → 2
Finance     → 2
HR          → 1

4. Sum

To calculate total salary per department:

df.groupBy("department").sum("salary")

Example:

Engineering → 215000
Finance     → 195000
HR          → 75000

You can rename the resulting column:

df.groupBy("department") \
  .sum("salary") \
  .withColumnRenamed("sum(salary)", "total_salary")

But for production code, agg() is often cleaner when multiple metrics are required.

5. Average

Calculate average salary:

df.groupBy("department").avg("salary")

Result conceptually:

Engineering → 107500
Finance     → 97500
HR          → 75000

6. Minimum and Maximum

Minimum salary:

df.groupBy("department").min("salary")

Maximum salary:

df.groupBy("department").max("salary")

These are useful for profiling and business analytics.

7. Using agg()

For real-world pipelines, you often need several metrics simultaneously.

For example:

from pyspark.sql import functions as F
result = df.groupBy("department").agg(
    F.count("*").alias("employee_count"),
    F.sum("salary").alias("total_salary"),
    F.avg("salary").alias("average_salary"),
    F.min("salary").alias("minimum_salary"),
    F.max("salary").alias("maximum_salary")
)

Output:

+-----------+--------------+------------+--------------+-------------+-------------+
|department |employee_count|total_salary|average_salary|minimum_salary|maximum_salary|
+-----------+--------------+------------+--------------+-------------+-------------+
|Engineering|2             |215000      |107500        |95000        |120000       |
|Finance    |2             |195000      |97500         |85000        |110000       |
|HR         |1             |75000       |75000         |75000        |75000        |
+-----------+--------------+------------+--------------+-------------+-------------+

This is one of the most useful patterns to know for interviews.

8. Multiple Grouping Columns

You aren’t limited to one grouping column.

Suppose you want:

department + country

You can write:

df.groupBy(
    "department",
    "country"
).agg(
    F.count("*").alias("employee_count"),
    F.avg("salary").alias("average_salary")
)

The grouping key becomes:

(department, country)

For example:

Engineering + US
Engineering + UK
Finance     + US
Finance     + UK

Spark groups records based on the combination.

9. Conditional Aggregations

Production analytics frequently requires conditional metrics.

Suppose we want to count employees earning more than $100,000.

df.groupBy("department").agg(
    F.sum(
        F.when(F.col("salary") > 100000, 1)
         .otherwise(0)
    ).alias("high_salary_count")
)

Conceptually:

salary > 100000
       ↓
      1
salary <= 100000
       ↓
      0
       ↓
     SUM

This allows you to create business-specific metrics without creating multiple intermediate DataFrames.

10. count() vs countDistinct()

These are not the same.

Count rows

F.count("*")

counts records.

Count distinct values

F.countDistinct("customer_id")

counts unique customers.

Suppose:

customer_id
-----------
101
101
102
103
103

Then:

count(*)              = 5
countDistinct(customer_id) = 3

This distinction is frequently tested in interviews.

11. Aggregating Distinct Values

You can also collect distinct values.

For example:

df.groupBy("department").agg(
    F.collect_set("country").alias("countries")
)

Possible output:

Engineering → [US, UK]
Finance     → [US]
HR          → [UK]

Be careful with operations such as collect_set() when groups can contain a very large number of distinct values.

The result itself can become large.

12. collect_list() vs collect_set()

These functions are commonly confused.

collect_list()

Keeps duplicates.

F.collect_list("product")

Input:

A
A
B

Result:

[A, A, B]

collect_set()

Removes duplicates.

F.collect_set("product")

Result:

[A, B]

Use these carefully because the aggregated collection can become large.

13. Aggregations and NULL Values

NULL handling matters in production.

Consider:

salary
------
100000
120000
NULL

An average aggregation generally ignores NULL values rather than treating NULL as zero.

Therefore:

F.avg("salary")

is conceptually based on:

100000
120000

not:

100000
120000
0

If your business logic requires NULL to mean zero, make that explicit:

F.avg(
    F.coalesce(F.col("salary"), F.lit(0))
)

But don’t automatically convert NULL to zero.

A NULL can mean:

“The value is unknown.”

That is semantically different from:

“The value is zero.”

14. Aggregations Can Trigger a Shuffle

This is where aggregation becomes a performance topic.

Consider:

df.groupBy("department").agg(
    F.sum("salary")
)

Imagine the data initially looks like:

Executor 1:
Engineering
Finance
Executor 2:
HR
Engineering
Executor 3:
Finance
HR

Spark needs to combine records belonging to the same grouping key.

Conceptually:

Engineering ──┐
Engineering ──┼──→ Same partition
Engineering ──┘
Finance ──────┐
Finance ──────┼──→ Same partition
Finance ──────┘
HR ───────────┐
HR ───────────┼──→ Same partition
HR ───────────┘

This redistribution is a shuffle.

15. Why Shuffle Makes Aggregations Expensive

A shuffle can involve:

CPU
 ↓
Serialization
 ↓
Network transfer
 ↓
Memory
 ↓
Disk spill
 ↓
Deserialization

Therefore:

df.groupBy("customer_id").sum("amount")

can become expensive when the dataset is large.

This is why aggregation performance is closely connected to:

• partitioning
• shuffle
• data skew
• memory
• network I/O

16. The Two-Phase Nature of Aggregation

One of Spark’s important optimization ideas is that aggregation can often be performed partially before data is shuffled.

Suppose:

Partition 1:
A → 10
A → 20
Partition 2:
A → 30
A → 40

Instead of immediately sending every record across the network, Spark can conceptually perform local aggregation:

Partition 1:
A → 30
Partition 2:
A → 70

Then shuffle the partial results:

A → 30
A → 70

and combine:

A → 100

Conceptually:

Local Aggregation
       ↓
Partial Results
       ↓
Shuffle
       ↓
Final Aggregation

This can significantly reduce the amount of data that needs to move across the network.

The exact physical strategy depends on the query and Spark’s execution plan.

17. Why Aggregation Before Shuffle Matters

Imagine:

Original records:
1,000,000,000

After local aggregation:

Partial records:
10,000,000

The exact reduction depends heavily on the grouping keys and distribution, but the principle is important:

Reduce data locally before moving it across the network whenever the execution strategy allows it.

This is one reason Spark’s aggregation operators can be much more efficient than naively moving every record.

18. Aggregation and Data Skew

Consider:

customer_id = 1001 → 500 million records
customer_id = 1002 → 10,000 records
customer_id = 1003 → 8,000 records

If you aggregate by:

df.groupBy("customer_id").sum("amount")

the key 1001 can dominate the computation.

Conceptually:

Partition 1 → 10 million records
Partition 2 → 8 million records
Partition 3 → 12 million records
Partition 4 → 500 million records  ← Straggler

This can lead to:

• one slow task
• executor memory pressure
• spilling
• long stage duration

This is why data skew is one of the next critical concepts after understanding aggregations.

19. Aggregation Strategy: Filter Before Grouping

Suppose you only care about active customers.

Instead of:

df.groupBy("customer_id").agg(
    F.sum("amount")
).filter(...)

when logically equivalent, consider filtering earlier:

df.filter(
    F.col("status") == "ACTIVE"
).groupBy("customer_id").agg(
    F.sum("amount").alias("total_amount")
)

Conceptually:

All Data
   ↓
Filter
   ↓
Less Data
   ↓
GroupBy
   ↓
Shuffle Less Data

This can reduce:

• CPU
• memory
• network traffic
• shuffle volume

The key condition is that moving the filter earlier must preserve the intended semantics.

20. Select Only Required Columns

Suppose your input contains:

100 columns

but your aggregation requires:

customer_id
amount

Avoid unnecessarily carrying the other 98 columns.

For example:

df.select(
    "customer_id",
    "amount"
).groupBy("customer_id").agg(
    F.sum("amount").alias("total_amount")
)

This can reduce the amount of data processed and shuffled.

21. Aggregation with Dates

A common production requirement is time-based aggregation.

Suppose you have:

transaction_date
customer_id
amount

You can derive a month:

monthly = (
    df
    .withColumn(
        "month",
        F.date_trunc("month", "transaction_date")
    )
    .groupBy("month")
    .agg(
        F.sum("amount").alias("total_sales"),
        F.countDistinct("customer_id").alias("unique_customers")
    )
)

Result:

+-------------------+-----------+----------------+
|month              |total_sales|unique_customers|
+-------------------+-----------+----------------+
|2026-01-01 00:00:00|...        |...             |
|2026-02-01 00:00:00|...        |...             |
+-------------------+-----------+----------------+

This pattern appears frequently in:

• financial analytics
• retail
• insurance
• telecom
• customer analytics
• reporting pipelines

22. Aggregations with Window Functions Are Different

A common interview trap:

Is a window function the same as a groupBy() aggregation?

No.

groupBy() reduces multiple rows into fewer rows.

For example:

Input: 100 rows
       ↓
groupBy(customer_id)
       ↓
10 customers

A window function can calculate an aggregate while retaining the original rows.

For example:

from pyspark.sql.window import Window
window = Window.partitionBy("department")
df.withColumn(
    "department_avg_salary",
    F.avg("salary").over(window)
)

Input:

Alice → Engineering → 120000
Bob   → Engineering → 95000

Output can retain both rows:

Alice → Engineering → 120000 → 107500
Bob   → Engineering → 95000  → 107500

Think:

groupBy()
   ↓
Reduce rows
Window
   ↓
Keep rows + calculate over related rows

This distinction is extremely important.

23. Aggregation vs Window Function

24. Aggregation with Multiple Metrics

Production reporting often requires several metrics in one pass.

For example:

summary = (
    df.groupBy("customer_id")
      .agg(
          F.count("*").alias("transaction_count"),
          F.sum("amount").alias("total_amount"),
          F.avg("amount").alias("average_amount"),
          F.min("amount").alias("minimum_amount"),
          F.max("amount").alias("maximum_amount"),
          F.countDistinct("product_id").alias("unique_products")
      )
)

This produces a compact customer-level summary.

Rather than building many separate aggregation pipelines, grouping related metrics into one aggregation can make the pipeline easier to understand and can allow Spark to plan them together.

25. agg() with Expressions

You aren’t limited to simple column references.

For example:

df.groupBy("department").agg(
    F.round(F.avg("salary"), 2).alias("avg_salary")
)

Or:

df.groupBy("department").agg(
    (
        F.sum("salary") / F.countDistinct("employee_id")
    ).alias("salary_per_employee")
)

You can compose Spark SQL functions into expressions.

Prefer built-in functions whenever possible rather than moving the aggregation logic into Python.

26. A Production-Style Example

Suppose we have an e-commerce transaction dataset:

transaction_id
customer_id
country
product_id
amount
status
transaction_date

We want:

Monthly sales and unique customers for successful transactions.

A clean PySpark implementation:

from pyspark.sql import functions as F
result = (
    df
    .filter(F.col("status") == "SUCCESS")
    .select(
        "customer_id",
        "amount",
        "transaction_date"
    )
    .withColumn(
        "month",
        F.date_trunc("month", "transaction_date")
    )
    .groupBy("month")
    .agg(
        F.sum("amount").alias("total_sales"),
        F.countDistinct("customer_id").alias("unique_customers"),
        F.count("*").alias("transaction_count")
    )
)

Notice the sequence:

Read
 ↓
Filter
 ↓
Select required columns
 ↓
Derive month
 ↓
Group
 ↓
Aggregate

This is a good example of building a pipeline where unnecessary data is reduced before the aggregation.

27. How to Debug a Slow Aggregation

Suppose:

df.groupBy("customer_id").agg(
    F.sum("amount")
).show()

takes 45 minutes.

Don’t immediately increase executor count.

Investigate:

Step 1 — Inspect the plan

df.groupBy("customer_id").agg(
    F.sum("amount")
).explain("formatted")

Step 2 — Check Spark UI

Look for:

• shuffle read
• shuffle write
• stage duration
• task duration
• spill
• failed tasks

Step 3 — Check key distribution

df.groupBy("customer_id") \
  .count() \
  .orderBy(F.desc("count")) \
  .show(20)

Step 4 — Check input volume

Did today’s data increase from:

100 GB

to:

2 TB?

Step 5 — Check partitions

Determine whether the data is appropriately partitioned for the workload.

28. Aggregation and Adaptive Query Execution

Modern Spark can use Adaptive Query Execution (AQE) to adapt certain aspects of query execution based on runtime statistics.

This can help with situations such as:

• post-shuffle partition sizing
• skewed joins
• certain join strategy decisions

Conceptually:

Initial Plan
     ↓
Execute
     ↓
Runtime Statistics
     ↓
Adaptive Optimization
     ↓
Continue Execution

This is particularly useful because real-world data distributions aren’t always what the optimizer initially expects.

29. Common Aggregation Mistakes

Mistake 1 — Grouping unnecessarily

Bad:

df.groupBy("customer_id").count()

when you actually need a global count:

df.count()

Don’t create a distributed grouping operation when a simpler operation solves the problem.

Mistake 2 — Carrying unnecessary columns

Avoid carrying dozens of unused columns into expensive aggregation pipelines.

Mistake 3 — Ignoring skew

A single extremely frequent key can dominate aggregation time.

Mistake 4 — Confusing count() and countDistinct()

These answer different questions.

Mistake 5 — Collecting aggregated results blindly

Even an aggregated DataFrame can be large.

For example:

df.groupBy("customer_id").count().collect()

could still produce millions of rows.

Mistake 6 — Using Python UDFs unnecessarily

Prefer native Spark SQL functions for aggregation logic where possible.

30. ⚡ Performance Checklist

Before deploying a large aggregation pipeline, ask:

□ Did I filter unnecessary rows early?
□ Did I select only required columns?
□ Could the grouping key be heavily skewed?
□ How much data will be shuffled?
□ Is the partition count appropriate?
□ Am I accidentally using collect()?
□ Can Spark use built-in expressions?
□ Do I need caching?
□ Have I inspected explain("formatted")?
□ Have I checked the Spark UI?
□ Is AQE enabled and appropriate for this workload?
□ Is the output creating too many files?

31. Aggregation Mental Model

Remember this:

                  AGGREGATION
                       │
                       ↓
                  groupBy(key)
                       │
                       ↓
              Same keys together
                       │
                       ↓
                    Shuffle
                       │
                       ↓
             Partial / Final Agg
                       │
                       ↓
                  Result

And from the performance perspective:

More Input
    ↓
More Data to Process
    ↓
More Data to Shuffle
    ↓
More Memory / Network / Disk
    ↓
Potentially Slower Job

Therefore:

Filter early
Select early
Understand skew
Control partitions
Inspect the plan

32. Common Interview Questions

1. What is aggregation in Spark?

Aggregation combines multiple records based on one or more grouping keys and calculates summary metrics such as count, sum, average, minimum, or maximum.

2. Does groupBy() cause a shuffle?

It can. Grouping data across partitions generally requires records with the same key to be brought together, which is commonly implemented using a shuffle.

3. Why is groupBy() expensive?

Because grouping large distributed datasets can involve network redistribution, serialization, memory usage, and disk spill.

4. What is the difference between count() and countDistinct()?

count() counts records or non-null values depending on the expression, while countDistinct() counts unique values.

5. What is the difference between groupBy() and a window function?

groupBy() generally reduces rows into groups. Window functions calculate values over related rows while retaining the individual rows.

6. How would you optimize a slow aggregation?

I would inspect the Spark UI and physical plan, check shuffle volume, partition count, data skew, input size, filtering, column pruning, and whether the grouping strategy can be improved.

7. What is data skew in aggregation?

Data skew occurs when some grouping keys have disproportionately more records than others, causing uneven partition sizes and potentially creating slow straggler tasks.

8. Why should you filter before aggregation?

If the filter is logically safe to push earlier, it reduces the amount of data that reaches the aggregation and potentially reduces shuffle, CPU, memory, and network usage.

9. Why prefer built-in aggregation functions over Python UDFs?

Spark understands built-in expressions and can optimize them as part of the query plan, while Python UDFs can introduce additional execution and serialization overhead.

10. Can aggregation happen partially before shuffle?

Yes. Spark can use partial/local aggregation strategies for supported aggregations, reducing the amount of intermediate data that needs to be shuffled.

33. One Interview Scenario

Interviewer:

You have a 2 TB transaction dataset and need to calculate total transaction amount by customer. The job is taking 90 minutes. What would you do?

Weak answer:

Increase the number of executors.

Strong answer:

First, I would inspect the Spark UI and physical plan to determine whether the bottleneck is shuffle, skew, spilling, insufficient parallelism, or something else. Since the aggregation is by customer, I’d check the customer-key distribution for skew. I’d also filter unnecessary records and select only required columns before the aggregation, verify the number of shuffle partitions, and check whether AQE is helping. If the grouping key is highly skewed, I’d investigate an appropriate skew-handling strategy. I would then benchmark the change rather than assuming a larger cluster will solve the problem.

That answer demonstrates production thinking.

34. The Most Important Concept

If you remember only one thing from this article, remember this:

Aggregation is not just a groupBy() operation. At scale, it is a data-distribution problem.

When you write:

df.groupBy("customer_id").sum("amount")

think:

                    2 TB
                     │
                     ↓
              Group by customer
                     │
                     ↓
                  SHUFFLE
                     │
          ┌──────────┼──────────┐
          ↓          ↓          ↓
       Partition  Partition  Partition
          │          │          │
          └──────────┼──────────┘
                     ↓
              Final Aggregation

Your ability to reason about what happens between those boxes is what makes you strong at Spark.

Key Takeaways

• Aggregations summarize data across groups.
• groupBy() defines grouping keys.
• agg() allows multiple metrics in one operation.
• Common functions include count, sum, avg, min, max, and countDistinct.
• Aggregations can require a shuffle.
• Shuffle can become expensive because of network, CPU, memory, serialization, and disk spill.
• Spark can perform partial/local aggregation before shuffle for supported operations.
• Data skew can create extremely slow aggregation tasks.
• Filtering and column pruning before aggregation can reduce data movement.
• groupBy() and window functions solve different problems.
• Avoid unnecessary collect().
• Prefer built-in Spark functions over unnecessary Python UDFs.
• Use Spark UI and explain("formatted") to investigate slow aggregations.
• AQE can improve certain runtime execution decisions.

Continue Learning Apache Spark

You’ve now moved from:

Transformations
      ↓
Actions
      ↓
Aggregations

The next question is:

How does Spark physically distribute all this data across the cluster?

That’s where partitioning becomes critical.

Next Article → Partitioning in Apache Spark

We’ll explore:

• What exactly is a partition?
• How many partitions should you have?
• How does one partition relate to a task?
• What is repartition()?
• What is coalesce()?
• How does partitioning affect shuffle?
• How can bad partitioning make a Spark job slower?
• How should you think about partition size in production?

Further Reading

1. ✓ What is Apache Spark?
2. ✓ Why Spark is Faster than Hadoop
3. ✓ Spark Architecture
4. ✓ Driver vs Executor
5. ✓ Cluster Managers
6. → RDD vs DataFrame vs Dataset
7. → SparkSession
8. ✓ Reading Data
9. ✓ Lazy Evaluation
10. ✓ DAG
11. ✓ Stages and Tasks
12. ✓ Shuffle
13. ✓ Narrow vs Wide Transformations
14. → Writing Data
15. → Data Transformations
16. →Window Functions
17. Aggregations