PySpark Aggregations

Spark tutorial · PySpark.in

PySpark Aggregations

1. Using groupBy and agg in PySpark

Explanation

In PySpark, the groupBy() operation is analogous to the SQL GROUP BY clause – it groups rows by one or more specified columns. After grouping, you can apply aggregation functions on each group using the .agg() method or convenient shorthand methods. The agg() function allows you to compute summary statistics (like counts or sums) for each group. Essentially, groupBy partitions the DataFrame into subsets (one subset per unique key), and then agg computes aggregate metrics per subset.

How to use: You call df.groupBy("columnName") to get a grouped object, then apply an aggregation like count(), sum(), etc., or use agg() to specify one or multiple aggregation functions. For example, df.groupBy("Category").sum("Amount") will group by the Category column and sum up the Amount for each category.

Example: Group by one column and aggregate

Suppose we have a DataFrame of employees with their department and salary:

```

# Group by department and count employees in each department

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.getOrCreate()

# Sample data: Name, Department, Salary

data = [

("Alice", "IT", 45000),

("Bob", "CS", 85000),

("Charlie", "CS", 41000),

("David", "IT", 56000),

("Eve", "ECE", 45000),

("Frank", "ECE", 49000)

]

df = spark.createDataFrame(data, ["Name", "Dept", "Salary"])

dept_counts = df.groupBy("Dept").count()

dept_counts.show()

```

the DataFrame by the Dept column and then used .count() to count how many rows (employees) fall into each department. Each department appears once in the result with its count. The groupBy turns rows with the same Dept into groups, and count() gives the size of each group.

Example: Aggregating with multiple functions

You can calculate multiple aggregates at once with agg(). For instance, to get the minimum, maximum, and average salary per department, you can do:

```

# Multiple aggregations: min, max, avg salary per department

from pyspark.sql import functions as F

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.getOrCreate()

# Sample data: Name, Department, Salary

data = [

("Alice", "IT", 45000),

("Bob", "CS", 85000),

("Charlie", "CS", 41000),

("David", "IT", 56000),

("Eve", "ECE", 45000),

("Frank", "ECE", 49000)

]

df = spark.createDataFrame(data, ["Name", "Dept", "Salary"])

dept_stats = df.groupBy("Dept").agg(

F.min("Salary").alias("MinSalary"),

F.max("Salary").alias("MaxSalary"),

F.avg("Salary").alias("AvgSalary")

)

dept_stats.show()

```

Here, we used groupBy("Dept").agg(...) with multiple aggregation functions. The result shows, for each department, the lowest salary, highest salary, and average salary. PySpark provides many built-in aggregate functions (in pyspark.sql.functions) that you can use inside agg.

Explanation

At the intermediate level, it's important to know that groupBy returns a special GroupedData object (specifically a RelationalGroupedDataset under the hood). You can call .agg() on it for flexible combinations of aggregates, or use shorthand methods like .count(), .sum(), .avg() directly on the grouped object for single aggregations. For example, df.groupBy("Dept").sum("Salary") is equivalent to using agg(F.sum("Salary")).

You can also pass a dictionary of column names to aggregation functions into agg. For instance, df.groupBy("Dept").agg({"Salary": "avg", "Age": "max"}) would compute the average salary and max age per department in one call. If you have many aggregates, using the functional API (F.sum, F.count, etc.) with aliases (as shown above) is clearer.

Note on performance: Grouping data can be expensive on large datasets because it requires shuffling data by keys. Ensuring the DataFrame is partitioned appropriately or using techniques like approximate aggregations (e.g., approx_count_distinct) for large cardinality can help. Also note that Spark’s aggregations are memory-efficient for numeric computations, but custom aggregation (like Python UDF-based ones) may be slower.

Finally, if you try to group by a column and call agg without specifying an aggregate for another column, Spark will throw an error – every column in the result must be either grouped or aggregated. This is different from Pandas, where you can get implicit grouping on non-aggregated columns; in Spark you must explicitly aggregate or group all relevant columns.

2. Basic Aggregation Functions: count, sum, avg, min, max

Explanation

PySpark provides common aggregate functions to compute basic statistics on DataFrames. The most frequently used are:

count() – counts the number of rows (or non-null values in a column).
sum() – sums up numeric values in a column.
avg() (or mean()) – calculates the average (mean) of a numeric column.
min() – finds the minimum value in a column.
max() – finds the maximum value in a column.

You can use these after a groupBy to get per-group metrics, or use them on the entire DataFrame (e.g., df.select(F.sum("column"))) to get global totals. When used after groupBy, they compute the statistic within each group.

Example: Using basic aggregates

Using the same df from above (employees data), we can illustrate these functions:

```

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.getOrCreate()

# Sample data: Name, Department, Salary

data = [

("Alice", "IT", 45000),

("Bob", "CS", 85000),

("Charlie", "CS", 41000),

("David", "IT", 56000),

("Eve", "ECE", 45000),

("Frank", "ECE", 49000)

]

df = spark.createDataFrame(data, ["Name", "Dept", "Salary"])

# Total number of records in DataFrame

total_rows = df.count() # DataFrame.count() returns total rows

print(f"Total rows: {total_rows}")

# Total rows: 6

# Calculate global statistics for Salary column (no grouping)

df.select(

F.count("Salary").alias("Count"),

F.sum("Salary").alias("TotalSalary"),

F.avg("Salary").alias("AvgSalary"),

F.min("Salary").alias("MinSalary"),

F.max("Salary").alias("MaxSalary")

).show()

```

In the above code, we did not use groupBy – so the aggregation functions were applied to the entire DataFrame. We got the total count of rows, sum of all salaries, average salary, minimum salary, and maximum salary across all employees.

Example: Aggregates with groupBy

Now, let's use these functions with grouping. For instance, get the count of employees and total salary per department:

```

df.groupBy("Dept").agg(

F.count("*").alias("NumEmployees"),

F.sum("Salary").alias("TotalSalary")

).show()

```

Here we used count("*") which counts rows in each group (including nulls, since * counts rows regardless of content), and sum("Salary") to add up salaries in each department. The results show each department with the number of employees and the sum of their salaries.

Notes:

count("Column") will count non-null values in that column for each group. If you want to count rows even if the column has nulls, use count("*") or count(1) as shown above.
These aggregate functions ignore null values by default (except count(*)). For example, avg("X") will compute the average of non-null X values in the group.
You can chain multiple aggregate functions in one agg() call as shown. Each aggregated result can be given an alias name for clarity.

Explanation

At an intermediate level, you should know some additional details and tricks about these aggregations:

There are variants and related functions: e.g., countDistinct() for distinct counts, sumDistinct() for sum of unique values, and statistical functions like stddev, variance etc., which might be useful depending on the analysis.
Multiple columns: You can aggregate multiple columns at once. For example, df.groupBy("A").sum("B", "C") will give sum of B and sum of C for each group in one go. The result will have two columns sum(B) and sum(C). Similarly, you could do min("B", "C"), etc.
Aliasing and column names: When you use these functions, Spark often gives default names like sum(colName). It’s a good practice to alias them (e.g., .alias("Total")) for readability.
Direct vs. agg: As mentioned, GroupedData (result of groupBy) has convenience methods: .count(), .sum(), .avg(), .min(), .max(), which can be used instead of agg. For instance, df.groupBy("Dept").max("Salary") directly returns the maximum salary per department. Under the hood, it’s doing the same thing as using agg(F.max("Salary")).

Example : Suppose you want to count distinct departments and also get their min and max salary in one query:

```

from pyspark.sql.functions import countDistinct

df.select(

countDistinct("Dept").alias("DistinctDeptCount"),

F.min("Salary").alias("MinSalary"),

F.max("Salary").alias("MaxSalary")

).show()

```

This tells us there are 3 distinct departments and shows the global min and max salary. In practice, you might also use .agg() after groupBy to get such info per group – for example, min and max salary per department as shown earlier.

Being comfortable with these basic aggregators is key, as they form the foundation for more advanced aggregations in PySpark.

3. Grouping by Multiple Columns

Explanation

PySpark allows grouping on multiple columns to get combinations of keys. This is similar to SQL grouping on multiple fields: the combination of those fields defines the group. For example, if you group by Country and City, each unique pair of (Country, City) will be a group. Grouping by multiple columns is as simple as passing both column names to groupBy, either as separate arguments or as a list.

The result of a multi-column groupBy will have one row per unique combination of those columns, and you can aggregate other columns within those groups.

Example: Group by two columns

Let's create a small DataFrame of students with their country and continent, and then count how many students are in each country-continent pair:

```

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.getOrCreate()

students = [

("Mario", "Italy", "Europe"),

("Luigi", "Italy", "Europe"),

("Peach", "Spain", "Europe"),

("Yoshi", "Japan", "Asia"),

("Li", "China", "Asia"),

("Akira", "Japan", "Asia")

]

students_df = spark.createDataFrame(students, ["Name", "Country", "Continent"])

students_df.groupBy("Continent", "Country").count().show()

```

In this output, each row is a unique (Continent, Country) pair, and the count is how many times that combination appears (i.e., number of students from that country on that continent). We see Europe-Italy has 2 students, Asia-Japan has 2, etc.

The syntax groupBy("Continent", "Country") is grouping on both columns. You could also pass a list: groupBy(["Continent", "Country"]) – it’s equivalent. All grouping columns appear in the result by default (as keys).

Explanation

When grouping by multiple columns, it helps to understand a few deeper points:

Spark’s grouping on multiple columns is essentially grouping by a composite key. Internally, Spark might treat the combination as one key (similar to grouping by a struct of those columns). This means if you have high cardinality (many unique combinations), the number of groups can explode.
The order in which you specify the groupBy columns doesn’t affect the result (except when using rollups/cubes, discussed next). For a plain groupBy, groupBy("A","B") yields the same groups as groupBy("B","A") (though the output might be sorted differently if you call .show() without sorting).
If you need to perform additional operations on the grouped result (like filtering groups by an aggregated value), you can. For example, you could chain .agg() and then .filter() to only keep groups meeting some condition (e.g., .filter("count > 1") to keep only groups with more than one row).

PySpark Aggregations