PySpark Aggregations: count, approx_count_distinct, mean and summary

1.groupBy().agg(): the entry point for aggregation

Once a dataset is clean, the next step is usually to aggregate it: count rows per key, average a measure, or estimate how many unique values a column has. In PySpark you group rows with groupBy and then call agg, which computes the aggregates and returns a new DataFrame. agg accepts its expressions in two forms. In the first, a dict maps a column name to the name of an aggregate function, such as {"*": "count"}. In the second, you pass aggregate Column expressions such as sf.min(df.age), and each expression defines one aggregation.

df = spark.createDataFrame(
    [(2, "Alice"), (3, "Alice"), (5, "Bob"), (10, "Bob")], ["age", "name"])

# Group-by name, and count each group.
df.groupBy(df.name).agg({"*": "count"}).sort("name").show()
# +-----+--------+
# | name|count(1)|
# +-----+--------+
# |Alice|       2|
# |  Bob|       2|
# +-----+--------+

agg accepts built-in aggregation functions such as avg, max, min, sum and count, as well as group aggregate pandas UDFs created with pandas_udf. The two kinds can't be mixed: built-ins and pandas UDFs have to go in separate agg calls. Every function on this page can also be used without groupBy, through df.agg(...) or df.select(...), in which case it aggregates over the whole DataFrame.

Sources1

2.Counting rows, non-null values and distinct values

"Count" can mean several different things. sf.count returns the number of items in a group, but what it counts depends on its argument. sf.count(sf.expr("*")) counts rows. sf.count(column) counts only the non-null values in that column.

df = spark.createDataFrame(
    [(1, "apple"), (2, "banana"), (3, None)], schema=["id", "fruit"])
df.select(sf.count(sf.expr("*"))).show()

The documented result is count(id) = 3 and count(fruit) = 2. Counting unique values needs a different function. sf.count_distinct(col, *cols) returns an exact distinct count. You can pass it several columns, and it then counts distinct combinations of their values. On the values 1, 1 and 3 it returns 2.

from pyspark.sql import functions as sf
df = spark.createDataFrame([(1,), (1,), (3,)], ["value"])
df.select(sf.count_distinct(df.value)).show()

Sources23

3.approx_count_distinct: trading precision for speed

Exact distinct counts can be expensive on very large tables. Databricks points out that using approximation for aggregates can speed up query processing and cut costs when you don't need a precise answer. Spark SQL provides native approximate aggregates, including approx_count_distinct, approx_percentile and approx_top_k. sf.approx_count_distinct(col, rsd=None) returns an estimate of the distinct count of a column, or of several columns when they are combined into a struct.

from pyspark.sql import functions as sf
spark.range(100000).agg(
    sf.approx_count_distinct("id").alias('with_default_rsd'),
    sf.approx_count_distinct("id", 0.1).alias('with_rsd_0.1')
).show()

+----------------+------------+
|with_default_rsd|with_rsd_0.1|
+----------------+------------+
|           95546|      102065|
+----------------+------------+

The rsd parameter is the maximum allowed relative standard deviation, and it defaults to 0.05. A larger rsd allows a looser estimate. If you need an rsd below 0.01, the documentation says count_distinct is the more efficient choice. At that level of precision the approximation no longer saves anything.

Sources45

4.mean: an alias of avg that skips nulls

sf.mean(col) returns the average of the values in a group. It is an alias of avg, so the two functions are interchangeable. The output column is named avg(...) even when you call mean.

import pyspark.sql.functions as sf
df = spark.createDataFrame([(1982, None), (1990, 2), (2000, 4)], ["birth", "age"])
df.select(sf.mean("age")).show()

+--------+
|avg(age)|
+--------+
|     3.0|
+--------+

Sources6

5.summary() and describe(): statistics for a whole DataFrame

Sometimes you want a quick profile of every column rather than a single aggregate. describe() computes basic statistics for numeric and string columns: count, mean, stddev, min and max. summary() goes further. Its available statistics are count, mean, stddev, min, max and arbitrary approximate percentiles given as a percentage, such as 75%. Called with no arguments, it returns count, mean, stddev, min, the 25%, 50% and 75% percentiles, and max. You can also pass the names of the statistics you want.

df.select("age", "weight", "height").summary("count", "min", "25%", "75%", "max").show()
# +-------+---+------+------+
# |summary|age|weight|height|
# +-------+---+------+------+
# |  count|  3|     3|     3|
# |    min| 11|  37.8| 142.2|
# |    25%| 11|  37.8| 142.2|
# |    75%| 13|  44.1| 150.5|
# |    max| 13|  44.1| 150.5|
# +-------+---+------+------+

Both pages carry the same warning: these functions are meant for exploratory data analysis, and the schema of the resulting DataFrame isn't guaranteed to stay backward compatible. Use them to inspect data. Don't build a production pipeline that parses their output; for that, write explicit agg expressions.

Sources7

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.sf.count(col) counts every row in the group, including rows where col is null.Why is that wrong?

   count on a specific column counts only its non-null values. To count all rows, use count(sf.expr("*")) or count("*").

   Covered in Counting rows, non-null values and distinct values

2. 2.approx_count_distinct returns the exact number of distinct values; the word 'approx' only refers to how it is computed.Why is that wrong?

   It returns an estimate within a relative standard deviation (default 0.05). The documented run on 100,000 unique ids returns 95546.

   Covered in approx_count_distinct: trading precision for speed

3. 3.describe() and summary() return the same statistics, so it doesn't matter which one you use to get percentiles.Why is that wrong?

   Only summary() supports approximate percentiles such as 25% and 75% and lets you choose the statistics. describe() returns count, mean, stddev, min and max.

   Covered in summary() and describe(): statistics for a whole DataFrame