groupBy + agg

What it does

groupBy on its own returns a GroupedData object, not a DataFrame. Calling agg (or a shortcut like count()) finishes the job and returns a DataFrame with the grouping columns followed by the aggregates.

Step by step

1. 1
   Rows are split into groups by the key, here dept.
2. 2
   Each aggregate is computed per group: count counts rows, avg averages non-null salaries.
3. 3
   The output has one row per group. Output order is not guaranteed: add orderBy if you need one.

Run the example

from pyspark.sql import functions as F
result = (employees
    .groupBy("dept")
    .agg(
        F.count("*").alias("people"),
        F.round(F.avg("salary")).alias("avg_salary"),
        F.max("salary").alias("top_salary"),
    )
    .orderBy("dept"))

SELECT dept,
       COUNT(*)            AS people,
       ROUND(AVG(salary))  AS avg_salary,
       MAX(salary)         AS top_salary
FROM employees
GROUP BY dept
ORDER BY dept

Farid's department is null, and he still gets a group: groupBy treats all nulls in a key as one group. In ascending order that null group sorts first.

Three ways to count

For a quick estimate on huge data, F.approx_count_distinct uses HyperLogLog++ and is far cheaper than an exact distinct count, with a default relative error of about 5%.

Filtering on an aggregate

SQL has HAVING for conditions on aggregates. In PySpark you just filter after the aggregation; Catalyst produces the same plan.

(employees.groupBy("dept")
    .agg(F.count("*").alias("n"))
    .filter(F.col("n") >= 2))

SELECT dept, COUNT(*) AS n
FROM employees
GROUP BY dept
HAVING COUNT(*) >= 2

Under the hood: partial aggregation

An aggregation needs all rows of a key on one executor, which means a shuffle. But Spark does not ship every row. It aggregates in two phases:

1. Map side

   Each partition computes partial results per key

2. →

   Shuffle

   Only one partial row per key per partition moves

3. →

   Reduce side

   Partials are merged into final results

For avg, the partial result is a (sum, count) pair; the final step divides. That is why the plan shows two HashAggregate nodes around an Exchange. With few distinct keys, this shuffles almost nothing.

Common mistakes