I am trying to do sampling with replacement in Scala/Spark, defining the probabilities for each class.

This is how I would do it in R.

# Vector to sample from
x <- c("User1","User2","User3","User4","User5")

# Occurenciens from which to obtain sampling probabilities
y <- c(2,4,4,3,2)

# Calculate sampling probabilities
p <- y / sum(y)

# Draw sample with replacement of size 10
s <- sample(x, 10, replace = TRUE, prom = p)

# Which yields (for example):
[1] "User5" "User1" "User1" "User5" "User2" "User4" "User4" "User2" "User1" "User3"

How can I do the same in Scala / Spark?



def weightedSampleWithReplacement[T](data: Array[T], 
                                     weights: Array[Double], 
                                     n: Int, 
                                     random: Random): Array[T] = {
  val cumWeights = weights.scanLeft(0.0)(_ + _)
  val cumProbs = cumWeights.map(_ / cumWeights.last)
  Array.fill(n) {
    val r = random.nextDouble()
    data(cumProbs.indexWhere(r < _) - 1)

Spark has an RDD.sample() method that can sample without replacement, though not with weights. You could probably adapt that method along the lines above to do this however.

  • $\begingroup$ Dear Sean, any hints about a Python/PySpark implementation? I'm facing the very same problem of sampling (a single element, thus no replacement needed) from RDD according to weighted probability: stackoverflow.com/questions/44352986/… $\endgroup$ – AlessioX Jun 4 '17 at 11:15

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.