Efficient Sampling Approach for Large Datasets

Sampling is a fundamental process in machine learning that involves selecting a subset of data from a larger dataset. This technique is used to make training and evaluation more efficient, especially when working with massive datasets where processing every data point is impractical

However, sampling comes with its own challenges. Ensuring that samples are representative is crucial to prevent biases that could lead to poor model generalization and inaccurate evaluation results. The sample size must strike a balance between performance and resource constraints. Additionally, sampling strategies need to account for factors such as class imbalance, temporal dependencies, and other dataset-specific characteristics to maintain data integrity.

The most common languages used for sampling are Scala and PySpark, both maintained by the Apache Foundation. And a common challenge in these languages is that when we do sampling, the entire dataset is stored in one machine, thereby leading to memory errors.

In this article, I will explain how random sampling can be achieved at scale using Scala Spark and how the central limit theorem can be extended to solve this problem.

Problem Statement

Solution

To solve this problem, we can use the central limit theorem, which basically says that if we pick a number between 0 and 1 and repeat this process n times, then the numbers are normally distributed. This can be applied to our sampling usecase here.

If we want to sample, say, 10% of the rows, then for each row we randomly pick a number between 0 and 1, filter the rows that are greater than 0.1, and then sample the dataset to 10% randomly.

Here is a sample code to achieve this sampling.

import org.apache.spark.sql.functions._
import org.apache.spark.sql.{DataFrame, Encoders, Row}



object DistributedSampler {

  def distributedSample1(df: DataFrame, fraction: Double): DataFrame = {
    df.withColumn("rand_val", rand())       // add random number 0–1
      .filter(col("rand_val") < fraction)   // keep rows that fall under the threshold
  }


  // sample each partition independently
  def distributedSample(df: DataFrame, fraction: Double): DataFrame = {
    df.mapPartitions { iter =>
      val rnd = new scala.util.Random()
      iter.filter(_ => rnd.nextDouble() < fraction)
    }(Encoders.kryo[Row])   // or Encoders.row
  }
}

Test code to verify the results with Spark sampling vs. the new approach:

import org.apache.spark.sql.SparkSession

object SamplingTest {
  def main(args: Array[String]): Unit = {

    val spark = SparkSession.builder()
      .appName("SamplingTest")
      .master("local[*]")
      .getOrCreate()

    import spark.implicits._

    val df = spark.range(0, 1000000000).toDF("id")

    // Spark built-in
    val t0 = System.currentTimeMillis()
    val s1 = df.sample(false, 0.1)
    println(s1.count())
    val t1 = System.currentTimeMillis()

    println(s"Default Spark sample time: ${t1 - t0} ms")

    // Distributed-safe
    val t2 = System.currentTimeMillis()
    val s2 = DistributedSampler.distributedSample1(df, 0.1)
    println(s2.count())
    val t3 = System.currentTimeMillis()

    println(s"Distributed sample time:   ${t3 - t2} ms")

    spark.stop()
  }
}

Results

From the results, we can see that the new approach performed better than Spark's default sampling. I am running this on a single-node machine, so the results might slightly vary, but the sample code shows that the approach works and is efficient.

Limitations

This solution only works if we want to sample the entire dataset randomly; this will not work. In some cases, I have seen records going beyond the sample size by a few records because of the nature of randomness; it's not a big number, but something to be aware of. For small datasets, randomness is high; for example, sampling 1000/10000 records can be done directly by Spark sampling, but if the number of rows exceeds 100k, the CLT approach also works efficiently.

In the next article, we will discuss emitting metrics from Spark and how it helps in analyzing your jobs.