Spark Under the Hood: randomSplit() and sample() Inner Workings
By: Meltem Tutar
Tags: Data Science, Big Data, Spark, Pyspark

Spark Implementation of randomSplit()

Signature Function

The signature function of randomSplit() includes a weight list and a seed specification. The weight list is to specify the number of splits and percentage (approximate) in each and the seed is for reproducibility. The ratio is approximate due to the nature of how it is calculated.

For example, the following code in Figure 3 would split df into two data frames, train_df being 80% and test_df being 20% of the original data frame. By using the same value for random seed, we are expecting that the same data points are in the same split if we were to re-run the script or Spark internally rebuilds the splits.

Figure 3: randomSplit() signature function example

Under the Hood

The following process is repeated to generate each split data frame: partitioning, sorting within partitions, and Bernoulli sampling. If the original data frame is not cached then the data will be re-fetched, re-partitioned, and re-sorted for each split calculation. This is the source of potential anomalies. In summary, randomSplit() is equivalent to performing sample() for each split with the percentage to sample changing with the split being performed. This is evident if you examine the source code for randomSplit() in PySpark³. This blog⁴ also provides some more information and visuals on how randomSplit() is implemented.

Let’s walk through an example. Figure 4 is a diagram of the sample() for each split, starting with the 0.80 split.

Figure 4: Process of generating the 0.8 split. Identical to sample() implementation.

Spark utilizes Bernoulli sampling, which can be summarized as generating random numbers for an item (data point) and accepting it into a split if the generated number falls within a certain range, determined by the split ratio. For a 0.8 split data frame, the acceptance range for the Bernoulli cell sampler would be [0.0,0.80].

The same sampling process is followed for the 0.20 split in Figure 5, with just the boundaries of acceptance changing to [0.80, 1.0].

Figure 5: Process of generating the 0.2 split. Identical to sample(). implementation. Partition contents remain constant and sorting order is preserved ensuring a valid split.

The data frame is re-fetched, partitioned, and sorted within partitions again. You can see in the example that RDD partitions are idempotent. Which means that the data points in each partition in Figure 4, remain in the same partition in Figure 5. For example, points b and c are in Partition 1 in both Figure 4 and 5. Additionally, the seed associated with each partition always remains constant, and the order within partitions is identical. All three of these points are fundamental to both sample() and randomSplit(). Ensuring that the same sample is produced with the same seed in the former, and guaranteeing no duplicates or disappearing data points in the latter.

Solutions To Avoiding Inconsistencies

Fixing these issues lies in ensuring that RDD partitions and sorting order are idempotent. Any one of the following three methods ensure this and can be applied: 1) caching the data frame before operations 2) repartitioning by a column or a set of columns, and 3) using aggregate functions⁵. An example of each method is shown in Figure 6.

Figure 6: Three different methods to avoid inconsistencies in randomSplit() and sample()

Caching the original data frame leads to partition content being held in memory. So instead of re-fetching data, partitioning and sorting, Spark continues operations using the partitioned data in memory. Note that cache() is an alias for persist(pyspark.StorageLevel.memory_only) which may not be ideal if you have memory limitations. Instead, you can consider using persist(pyspark.StorageLevel.memory_and_disk_only) . If there is no memory or disk space available, Spark will re-fetch and partition data from scratch, so it may be wise to monitor this from the Spark Web UI. Caching is the solution I chose in my case.

Summary and Key Takeaways

Moral of the story is: if unexpected behavior is happening in Spark, you just need to dig a bit deeper! Here is a summary of all the key points of this article:

randomSplit() is equivalent to applying sample() on your data frame multiple times, with each sample re-fetching, partitioning, and sorting your data frame within partitions.

The data distribution across partitions and sorting order is important for both randomSplit() and sample(). If either change upon data re-fetch, there may be duplicates or missing values across splits and the same sample using the same seed may produce different results.

These inconsistencies may not happen on every run, but to eliminate them completely, persist (aka cache) your data frame, repartition on a column(s), or apply aggregate functions such as groupBy.

Good luck in all of your Spark endeavors, and I hope this article is a good starting point to understanding the inner workings of Spark. I’d appreciate any and all comments you may have!