Using the Spark Aggregator class in Scala
By: Marty Jackson
Tags: Machine Learning, Aggregator, Scala, Apache Spark, Data Science

Note: If you already understand the Aggregator class and simply want to view an example of how it’s implemented in Scala, feel free to skip to the “Using our Aggregators” section.

Photo by Markus Winkler on Unsplash

Type-safe Aggregations

What are they?

You’re in an airport waiting for your luggage at the baggage claim. The internal turmoil is rising inside of you, as most other passengers have received their luggage by now and you sure don’t want to lose your clothes, your toiletries, or that bottle of deluxe craft beer you checked. Just as you start to despair, you see it: a large, bright orange suitcase that is clearly yours. Gleefully, you pick it up and a smile just begins to cross your face — until you notice a drip from the zipper. And then another. And another. While the bag itself is in the same condition as it was when you checked it, its contents are soaked. You can’t help but wonder, When did this happen? Was it the low pressure in the plane? Did it fall off of a cart? Should I have used bubble wrap? Unfortunately, you don’t really have any way of knowing. Although you should have used bubble wrap.

A similar feeling of distress occurs when a Data Scientist performs an aggregation on a Dataset in Spark, only for the contents of the output Dataset to be revealed as invalid in some way; maybe the input and output were both intact Datasets, but something about the output’s content type is incorrect. The internal mechanisms of the aggregation and what went wrong can typically only be diagnosed after a significant amount of detective work. Even then, there’s no guarantee that future instances of input data won’t cause another issue for the aggregation.

However, just like adding more bubble wrap to prevent a bottle breaking in your suitcase, we can create custom aggregations that are strongly typed throughout using Spark’s Aggregator class. This class allows a Data Scientist to specify exactly what transformations worker nodes perform, how data is combined across nodes, and what each step should entail in terms of expected input and output type. The result is an intimate knowledge of how an aggregation transformation is being performed and an automated way to diagnose any typing issues that might arise. In other words, it’s a foolproof way to save your bottle of beer next time you fly.

When to use type-safe aggregations

Deciding whether to perform type-safe aggregations on Spark Datasets can be a difficult decision. On the one hand, it’s tempting to throw caution to the wind, declare the input and output types as Dataset , and call it a day. From my experience, this works for many of the exploratory and analytic instances a Data Scientist might encounter, as data and type validations can be performed on-the-fly in a notebook or scripting environment.

On the other hand, there are instances (e.g. automated ML pipelines in production, creating an API for other Data Scientists) when you need to ensure a transformation won’t be yielding unexpected and potentially errant results without warning. These are the cases when you’ll want to use the Aggregator class in Spark. This class allows a Data Scientist to identify the input, intermediate, and output types when performing some type of custom aggregation.

I found Spark’s Aggregator class to be somewhat confusing when I first encountered it. There are several steps, seemingly unused parameters, and limited examples online. Most of the examples I’ve seen tend to outline simple use-cases like summing a collection of integers or calculating an average. While these are certainly useful for tautological purposes, they don’t offer much insight into other common aggregations that one might need to perform. Thus, my aim here is to demonstrate two uses of the Aggregator class that one might actually use in writing a data preprocessing application and break them down line-by-line.

An ` Aggregator ` example

Let’s suppose you have the following retail data for a home improvement store

Our hardware store dataset.

and want to do the following in a type-safe manner:

provide a count of how many unique customers appear in each retailCategory

collect an array of unique productID s, grouped by retailCategory

s, grouped by transform the above data into a Dataset that contains the previous two steps

Let’s also suppose we want to use the following general structure for implementing our aggregators:

Which will result in a Data Scientist calling the following API when using our library:

val myValue = DoSomething.transformData(…)