Spark sample is too slow

Question

I'm trying to execute a simple random sample with Scala from an existing table, containing around 100e6 records.

import org.apache.spark.sql.SaveMode

val nSamples = 3e5.toInt
val frac = 1e-5
val table = spark.table("db_name.table_name").sample(false, frac).limit(nSamples)
(table
  .write
  .mode(SaveMode.Overwrite)
  .saveAsTable("db_name.new_name")
)

But it is taking too long (~5h by my estimates).

Useful information:

1. I have ~6 workers. By analyzing the number of partitions of the table I get: 11433.

2. I'm not sure if the partitions/workers ratio is reasonable.

3. I'm running Spark 2.1.0 using Scala.

I have tried:

1. Removing the .limit() part.

2. Changing frac to 1.0, 0.1, etc.

Question: how can I make it faster?

Best,

Answer 1

Limit is definitely worth removing, but the real problem is that sampling requires a full data scan. No matter how low is the fraction, the time complexity is still O(N)*.

If you don't require good statistical properties, you can try to limit amount of data you've loaded in the first place by sampling data files first, and then subsampling from the reduced dataset. This might work reasonably well, if data is uniformly distributed.

Otherwise there is not much you can do about it, other than scaling your cluster.

---

* How do simple random sampling and dataframe SAMPLE function work in Apache Spark (Scala)?

Answer 2

you could first sample the partitions, and then sample from the partitions. Like this, you dont need a full table scan, but only works if your partitioning itself is random. AFAIK you need to use RDD API for this. This could look like this (plug in the numbers to match your desired number of samples):

val ds : Dataset[String] = ???

  val dsSampled = ds.rdd
  // take 1000 samples from every 10th partition
  .mapPartitionsWithIndex{case (i,rows) => if(i%10==0) scala.util.Random.shuffle(rows).take(1000) else Iterator.empty}
  .toDS()