I'm trying to perform stratified sampling in a Spark Dataframe, but the behaviour of the sampleBy function is (oddly) similar to sample.
Spark version 3.0.1
import spark.implicits._
val data = Seq( ("Java", 20000), ("Java", 10000), ("Java", 3000), ("Java", 17000),
("Python", 100000), ("Python", 20000),
("Scala", 3000), ("Scala", 4000), ("Scala", 1000), ("Scala", 43000), ("Scala", 2000), ("Scala", 9000)).toDF("Language", "Price")
val sample_size = 0.5
val seed = 762387
val stratify = "Language"
val subsample = data.sample(withReplacement=false, fraction=sample_size, seed=seed)
subsample.show()
val fractions = data.select(stratify).distinct().as[String].collect().map((_, sample_size)).toMap
println(fractions.mkString("\n")
val stratified_subsample = data.stat.sampleBy(stratify, fractions=fractions, seed=seed)
stratified_subsample.show()
Output:
+--------+------+
|Language| Price|
+--------+------+
| Java| 3000|
| Python|100000|
| Python| 20000|
| Scala| 3000|
| Scala| 43000|
| Scala| 2000|
| Scala| 9000|
+--------+------+
Scala -> 0.5
Python -> 0.5
Java -> 0.5
+--------+------+
|Language| Price|
+--------+------+
| Java| 3000|
| Python|100000|
| Python| 20000|
| Scala| 3000|
| Scala| 43000|
| Scala| 2000|
| Scala| 9000|
+--------+------+
Output with another seed (6354345):
+--------+------+
|Language| Price|
+--------+------+
| Java| 10000|
| Java| 17000|
| Python|100000|
| Scala| 3000|
| Scala| 4000|
| Scala| 1000|
| Scala| 43000|
| Scala| 2000|
| Scala| 9000|
+--------+------+
Scala -> 0.5
Python -> 0.5
Java -> 0.5
+--------+------+
|Language| Price|
+--------+------+
| Java| 10000|
| Java| 17000|
| Python|100000|
| Scala| 3000|
| Scala| 4000|
| Scala| 1000|
| Scala| 43000|
| Scala| 2000|
| Scala| 9000|
+--------+------+
I've tried with different DataFrames, different seeds and both dataframes are always equal. I have always the same behaviour, the sample is not stratified at all. I know sampleBy is not exact, but having the same behaviour doesn't seem ok. Is something wrong in my snippet?