Aggregate & Array Functions

Spark tutorial · PySpark.in

Aggregate Functions: groupBy, count, avg, sum, min, max

Overview: Aggregation is a fundamental operation for data analysis – summarizing data by groups. PySpark DataFrames support aggregation through the groupBy() method and a suite of aggregate functions like count, sum, avg (average), min, max, and more. The typical pattern is: group the DataFrame by one or more keys, then apply aggregation(s) to each group. The result is a new DataFrame representing the summary for each group.

Typical Use Cases:

• Counting records per category (e.g., number of employees in each department, using count).
• Summing numerical values per group (e.g., total sales per region, using sum).
• Calculating statistics like average, min, max per group (e.g., average salary by department, min and max sales per product).
• Multi-aggregations: computing several summary metrics in one go (e.g., for each group, count rows, sum revenue, and average rating all together).

Spark's groupBy returns a special GroupedData object on which you can call aggregate functions or the .agg() method to specify multiple aggregates. Common aggregation functions include count() (row count), sum() for totals, avg() for mean, min() and max() for extrema. Let's see these in action with a sales dataset.

Suppose we have a simple DataFrame of product sales, with columns for region, product, and revenue (sales amount):

Counting and Grouping

The most basic aggregation is counting rows in each group.

Example 1: Count sales records per region.

This tells us there were 3 sales records in the East region and 2 in the West region. We used count("*") which counts all rows; you could also use count("amount") (or any column) since none are null in this dataset. Alternatively, PySpark allows a shorthand df_sales.groupBy("region").count() which produces the same result with a default "count" column.

Sum, Average, Min, Max

Often, we want more than just counts. We can compute multiple aggregates in one go using the .agg() method and passing a list of aggregate expressions.

Example 2: Multiple aggregates – total sales, average sale, min and max sale per region.

Let's interpret this:

• East: 3 records. total_sales = 2100 (which is 1200 + 300 + 600), avg_sales = 700.0, min_sales = 300, max_sales = 1200.
• West: 2 records. total_sales = 900 (400 + 500), avg_sales = 450.0, min_sales = 400, max_sales = 500.

The average appears with a decimal (700.0 and 450.0) because Spark's avg returns a double type. If needed, we could format or round it, but here it's fine.

We see clearly that the East region had higher total sales and a higher max single sale, whereas the West had fewer and smaller sales. This kind of multi-aggregation query is very powerful for summarizing data in one pass.

Note: The groupBy("region").agg(...) pattern is used to compute multiple aggregates at once. If you only need one aggregate, you can use shortcuts like groupBy().sum("amount"), groupBy().avg("amount"), etc., which produce columns like sum(amount). Using alias() helps rename them to something friendlier as we did.

Grouping by Multiple Columns

You can group by more than one column if you need a finer breakdown. For example, df_sales.groupBy("region", "product").sum("amount") would give total sales per (region, product) pair. This is similar to SQL grouping on multiple fields. For brevity, we'll not show a full example output here, but it's straightforward: each unique combination of region and product becomes a group.

Best Practices and Notes on Aggregation:

• Grouping on high-cardinality columns (columns with very many distinct values) can be expensive. If possible, group on moderately sized categories (e.g., "country" rather than "user_id", if making sense for the analysis).
• PySpark's aggregations are distributed and scalable. Under the hood, Spark will shuffle data by the group keys. It’s important to ensure that the data is partitioned and the shuffle has enough resources, especially if a single group is very large.
• count vs countDistinct: If you need the number of unique values in a group for some column, use countDistinct("col") (or the approx_count_distinct for large datasets where an approximate is acceptable for speed).
• You can chain multiple aggregations like we did. Spark will compute them in one pass over the data when possible. This is efficient – it's better than writing separate groupBy for each metric.
• If you only need one aggregate metric from the whole DataFrame (not grouped by a key), you can use functions like df_sales.agg(sum("amount")) or simply use the DataFrame methods df_sales.sum(), df_sales.count() etc., which return a dataframe or value representing the aggregated result.

Use Case Example: In a business context, you might group sales data by region and quarter to see how different areas perform over time, or group user data by subscription_plan to compute average revenue per plan. Aggregation functions help turn raw data into insights like "Sales in East region totaled 2100, averaging 700 per order, with a maximum single sale of 1200".

Spark provides many more aggregate functions (e.g., mean alias for avg, stddev for standard deviation, collect_list to gather all items in a group, etc.), but count, sum, avg, min, max are the most frequently used. These functions, when used via PySpark's DataFrame API, are optimized and will run in parallel across your cluster.

Working with Array Columns: explode, array_contains, size

Overview: Modern data often includes arrays or lists (e.g., a list of tags, a list of purchased items, etc.). PySpark has functions to manipulate array-type columns. Here we discuss:

• explode() – converts an array into multiple rows, repeating the other columns as needed.
• array_contains(array, value) – returns true if a given value is present in the array.
• size(array) – returns the length of the array (number of elements).

These functions make it easier to work with nested data structures like JSON arrays or lists that come from data sources.

Typical Use Cases:

• explode: Flattening a nested array structure. For example, if each record is a customer with an array of purchased products, you might explode to get one row per (customer, product) for further analysis.
• array_contains: Filtering or flagging rows where an array contains a specific element. E.g., find all users who have "Python" in their skill set array.
• size: Counting how many elements are in an array, which can be useful for features like "number of items in an order" or "count of tags".

Let's use a concrete example. Consider an orders DataFrame where each order has an array of items purchased:

(The items array is shown truncated above for display; the actual arrays are as defined in the data.) Each order has an array of product names in the items column.

Using explode() to Unnest Arrays

The explode(array_col) function will take each element of the array and turn it into its own row. Other columns are duplicated accordingly. This is extremely useful to normalize data for analysis – for example, counting how many times each product was sold, or simply to filter at the item level.

Example: Explode orders into individual items.

Now each (order_id, item) is a separate row. Order 1001, which originally had 3 items, now appears as 3 rows (1001, Laptop), (1001, Mouse), (1001, Cable). This flattened view can be aggregated or filtered easily. For instance, we could now group by "item" to count total sales of each product.

Notice that the explode operation creates a new column (item in this case) for each element of the array. If an array is empty, explode will produce no rows for that entry (it essentially removes that row in the exploded result). If the array is null, Spark's explode will produce null output (you might consider using explode_outer if you want to keep those and get nulls).

Using array_contains() to Check for Membership

The array_contains(col, value) function returns a boolean: True if the array contains the specified value, False otherwise. We can use this to filter rows or create a flag column.

Example: Flag orders that include a "Laptop".

We added a boolean column has_laptop that is true for orders 1001 and 1004 (since their item lists contain "Laptop"), and false for the others. This is very handy for filtering: for example, df_orders.filter(array_contains(col("items"), "Laptop")) would give only orders that included a laptop. Under the hood, this is a distributed operation that checks each element of the array for the match.

Note: array_contains returns NULL if the array itself is null. In our data there's no null array, so we get straightforward booleans. If there were a possibility of null, you might combine with a coalesce or handle nulls accordingly.

Using size() to Get Array Length

The size(array_col) function returns an integer equal to the number of elements in the array (0 for an empty array). This can be used to know how many items are in each order, how many skills a person has, etc.

Example: Count number of items in each order.

We can see order 1001 had 3 items, order 1002 had 1, etc. This kind of information could be used in further analysis, e.g., to compute the average number of items per order or to find orders that might be bulk purchases (many items).

Best Practices with Arrays:

• Use explode judiciously. Exploding can dramatically increase the number of rows (in our example, 5 orders became 8 rows after explosion). This is often necessary for analysis, but be aware of the data size change. After exploding, you might apply filters or aggregations to reduce data for the next steps.
• After explode, if you need to join back or relate to the original structure, keep the keys (like we kept order_id) so you know which original row the elements came from.
• array_contains is great for simple membership tests. If you need to check for multiple values, you might chain multiple conditions or use array_overlap (Spark 2.4+ for arrays intersection) depending on the scenario.
• If your arrays are very large (e.g., hundreds of elements), consider whether exploding or using higher-order functions (like transform or filter functions in Spark 3.x which allow lambda expressions on arrays) might be more efficient for your use case. For simple membership or length, array_contains and size are straightforward and efficient.
• All these functions operate on the entire column in a distributed manner, meaning Spark will handle them across the cluster without Python loops. This is important for performance – always prefer these built-ins to manually iterating over array elements in Python.

By using array functions, we can easily manage data that has a one-to-many relationship embedded in a single row. Flattening with explode and then aggregating or filtering gives us flexibility to analyze nested data structures.

More Spark tutorials

• about pyspark
• text diagram
• Apache Spark Runtime Architecture
• Introduction to RDD
• Actions vs Transformations
• Lazy Evaluation in PySpark

All tutorials · Try the free PySpark compiler · Practice challenges