Fast and accurate categorical distribution without reshuffling in Apache Spark

In Seahorse we want to provide our users with accurate distributions for their categorical data. Categorical data can be thought of as possible results of an observation that can take one of K possible outcomes. Some examples: Nationality, Marital Status, Gender, Type of Education.

In the described scenario we don’t have prior knowledge whether a given feature is categorical or not. We would like to treat features as categorical as long as their value set is small enough. Moreover, if there are too many distinct values we are no longer interested in its discrete distribution.
Seahorse is build on top of Apache Spark. Naive and easy approach to this problem would be to simply use Spark’s reduceByKey or groupByKey methods. The ReduceByKey method operates on key-value pairs and accepts a function reducing two values into one. If there are many values for one key, those will get reduced to one value with the provided function.

// limit - max. number of distinct categories to treat a feature as categorical
if(rdd.distinct().count() < limit) {
  // causes partition reshuffling
  val discreteDistribution = rdd.map(v => v -> 1).reduceByKey(_ + _).collect()
  Some(discreteDistribution)
} else {
  None
}

Unfortunately using these methods means poor performance due to Spark’s partition reshuffling.
Counting occurrences of specific values of a feature can be easily done without reshuffling. We might think of calculating a category distribution as counting occurrences of each category. This can be done using RDD’s function aggregate.
The aggregate function reduces a dataset to one value using a seq and comb functions and an initial aggregator value. The seq function adds one value to an aggregator. It is used to combine all values from a partition into one aggregated value. The comb function adds two aggregators together – it is used to combine results from different partitions.
In our case aggregator value will be of type Option(Map[T, Long]) representing occurrences count for value of type T or None if there are too many distinct categorical values.
The comb function (mergeCombiners in code) adds two aggregator maps together merging their results. If the result map has more distinct categorical values than the specified limit – no distribution is returned (a None).
The seq function (mergeValue in code) increments occurrences count of a value in the aggregator map.
Let the code speak for itself:

import scalaz._
import scalaz.Scalaz._
(...) // other imports
object FunctionalDemo {
  // limit - represents the maximum number of categories.
  // if there are more distinct values than `limit`, no categorical distribution will be returned
  case class CategoryDistributionCalculator[T](limit: Int) {
    def calculate(rdd: RDD[T]): Option[Map[T, Long]] =
        rdd.aggregate(zero)(mergeValue _, mergeCombiners _)
    // initial aggregator value
    private val zero = Option(Map.empty[T, Long])
   def mergeCombiners[T](
      leftOpt: Option[Map[T, Long]],
      rightOpt: Option[Map[T, Long]]): Option[Map[T, Long]] = {
      for (left <- leftOpt; right <- rightOpt) yield {
        // Adds maps in such a way that numbers from same keys are summed.
        val sumMap = left |+| right
        // yields Some when boolean condition is met. None otherwise
        (sumMap.keySet.size <= limit).option(sumMap)
      }
    }.flatten
    def mergeValue[T](acc: Option[Map[T, Long]], next: T): Option[Map[T, Long]] =
      mergeCombiners(acc, Some(Map(next -> 1)))
  }
  def main(args: Array[String]): Unit = {
    (...)
    val valueSlow = rdd.map(v => v -> 1).reduceByKey(_ + _).collect()
    val valueFast = CategoryDistributionCalculator(30).calculate(rdd).get
    (...)
  }
}

The code above is clean, functional and takes advantage of Scalaz. Unfortunately, it allocates a new object per each record in the dataset. With Apache Spark, sometimes it’s better not to use immutable objects in favor of mutable ones for performance’s sake. According to the documentation, both seq and comb functions are allowed to modify and return their first argument instead of creating a new object to avoid memory allocation. In the next approach we take advantage of that. The code below is not as pretty, but its performance is much better.

(...) // imports
object MutableDemo {
  // limit - represents maximum size of categorical type
  // if there are more distinct values than `limit`, no categorical distribution will be returned
  case class CategoryDistributionCalculator[T](limit: Int) {
    // None is returned if there are too many distinct values (>limit) to calculate categorical distribution
    def calculate(rdd: RDD[T]): Option[Map[T, Long]] = {
      val mutableMapOpt = rdd.aggregate(zero)(mergeValue _, mergeCombiners _)
      val immutableMapOpt = mutableMapOpt.map(_.toMap)
      immutableMapOpt
    }
    private val zero = Option(mutable.Map.empty[T, Long])
    private def mergeValue(accOpt: Option[mutable.Map[T, Long]],
      next: T): Option[mutable.Map[T, Long]] = {
      accOpt.foreach { acc =>
        addOccurrencesToMap(acc, next, 1)
      }
      replacedWithNoneIfLimitExceeded(accOpt)
    }
    private def mergeCombiners(leftOpt: Option[mutable.Map[T, Long]],
      rightOpt: Option[mutable.Map[T, Long]]): Option[mutable.Map[T, Long]] = {
      for (left <- leftOpt; rightMap <- rightOpt) {
        rightMap.foreach { case (element, count) =>
          addOccurrencesToMap(left, element, count)
        }
      }
      replacedWithNoneIfLimitExceeded(leftOpt)
    }
    private def addOccurrencesToMap(
      occurrences: mutable.Map[T, Long],
      element: T,
      count: Long): Unit = {
      occurrences(element) = occurrences.getOrElse(element, 0L) + count
    }
    private def replacedWithNoneIfLimitExceeded(
      mapOpt: Option[mutable.Map[T, Long]]): Option[mutable.Map[T, Long]] = {
      mapOpt.flatMap { map =>
        if (map.size <= limit) mapOpt else None
      }
    }
  }
  def main(args: Array[String]): Unit = {
    (...)
    val valueSlow = rdd.map(v => v -> 1).reduceByKey(_ + _).collect()
    val valueFast = CategoryDistributionCalculator(30).calculate(rdd).get
    (...)
  }
}

 

Measurements

We put a hypothesis that the last approach should be at least two time faster. The `group by` operation works in two stages with a partition reshuffling between them. Our solution runs in one stage.
A simple experiment with Apache Spark seems to confirm that. For the experiment we used:

• a local Apache Spark cluster,
• a Data Set containing  10,000,000 rows,
• with 33 labels in the category column  (uniformly distributed).

In the experiment we execute an empty foreach statement first to force Spark to cache the data in the clusters (job 0). Otherwise the first job would have additional overhead (for reading the data).
The jobs 1 and 2 are parts of the first, naive solution. The job 1 measures categories count. The job 2 calculates the distributions.
The job 2 calculates distributions.
The job 3 and 4 represent the two  optimised solutions. The job 3 is the functional one. The job 4 is the one that takes the advantage of mutable objects.

// job 0 - forces Spark to cache data in cluster. It will take away data-loading overhead
// from future jobs and make future measurements reliable.
rdd.foreach((_) => ())
// job 1 - checks if there are too many distinct values to form category
if(rdd.distinct().count() < limit) {
 // job 2 - calculates categorical distribution
 val discreteDistribution = rdd.map(v => v -> 1).reduceByKey(_ + _).collect()
 Some(discreteDistribution)
} else {
 None
}
// job 3 - optimised mutable approach
FunctionalDemo.CategoryDistributionCalculator(limit).execute(rdd)
// job 4 - functional approach
MutableDemo.CategoryDistributionCalculator(limit).execute(rdd)

The job 1 (distinct count) and the job 2 (reduceByKey) have 400 tasks and use memory to perform shuffling.
The job 3 and  job 4 – have only 200 tasks, run in one stage and don’t use additional memory to perform shuffling. They also check for category counts dynamically.
An interesting observation is that the functional job 3 is slower than the reduceByKey job even though it has only 200 tasks. A memory allocation overhead is that severe!

Summary

Categorical features are very common in Data Science. With the proposed approach we are able to calculate accurate categorical distributions fast and without a prior knowledge if a feature is categorical or not. It is guaranteed to execute in only one Apache Spark job and without any partition reshuffling.