diff options
author | Xiangrui Meng <meng@databricks.com> | 2014-06-26 21:46:55 -0700 |
---|---|---|
committer | Patrick Wendell <pwendell@gmail.com> | 2014-06-26 21:46:55 -0700 |
commit | c23f5db32b3bd4d965d56e5df684a3b814a91cd6 (patch) | |
tree | b9eabeef9945a5b5938c523f66d01f647376e20b /core | |
parent | d1636dd72fc4966413baeb97ba55b313dc1da63d (diff) | |
download | spark-c23f5db32b3bd4d965d56e5df684a3b814a91cd6.tar.gz spark-c23f5db32b3bd4d965d56e5df684a3b814a91cd6.tar.bz2 spark-c23f5db32b3bd4d965d56e5df684a3b814a91cd6.zip |
[SPARK-2251] fix concurrency issues in random sampler
The following code is very likely to throw an exception:
~~~
val rdd = sc.parallelize(0 until 111, 10).sample(false, 0.1)
rdd.zip(rdd).count()
~~~
because the same random number generator is used in compute partitions.
Author: Xiangrui Meng <meng@databricks.com>
Closes #1229 from mengxr/fix-sample and squashes the following commits:
f1ee3d7 [Xiangrui Meng] fix concurrency issues in random sampler
Diffstat (limited to 'core')
3 files changed, 38 insertions, 22 deletions
diff --git a/core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala b/core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala index 247f10173f..32c5fdad75 100644 --- a/core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala +++ b/core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala @@ -54,17 +54,17 @@ trait RandomSampler[T, U] extends Pseudorandom with Cloneable with Serializable */ @DeveloperApi class BernoulliSampler[T](lb: Double, ub: Double, complement: Boolean = false) - (implicit random: Random = new XORShiftRandom) extends RandomSampler[T, T] { - def this(ratio: Double)(implicit random: Random = new XORShiftRandom) - = this(0.0d, ratio)(random) + private[random] var rng: Random = new XORShiftRandom - override def setSeed(seed: Long) = random.setSeed(seed) + def this(ratio: Double) = this(0.0d, ratio) + + override def setSeed(seed: Long) = rng.setSeed(seed) override def sample(items: Iterator[T]): Iterator[T] = { items.filter { item => - val x = random.nextDouble() + val x = rng.nextDouble() (x >= lb && x < ub) ^ complement } } @@ -72,7 +72,7 @@ class BernoulliSampler[T](lb: Double, ub: Double, complement: Boolean = false) /** * Return a sampler that is the complement of the range specified of the current sampler. */ - def cloneComplement(): BernoulliSampler[T] = new BernoulliSampler[T](lb, ub, !complement) + def cloneComplement(): BernoulliSampler[T] = new BernoulliSampler[T](lb, ub, !complement) override def clone = new BernoulliSampler[T](lb, ub, complement) } @@ -81,21 +81,21 @@ class BernoulliSampler[T](lb: Double, ub: Double, complement: Boolean = false) * :: DeveloperApi :: * A sampler based on values drawn from Poisson distribution. * - * @param poisson a Poisson random number generator + * @param mean Poisson mean * @tparam T item type */ @DeveloperApi -class PoissonSampler[T](mean: Double) - (implicit var poisson: Poisson = new Poisson(mean, new DRand)) - extends RandomSampler[T, T] { +class PoissonSampler[T](mean: Double) extends RandomSampler[T, T] { + + private[random] var rng = new Poisson(mean, new DRand) override def setSeed(seed: Long) { - poisson = new Poisson(mean, new DRand(seed.toInt)) + rng = new Poisson(mean, new DRand(seed.toInt)) } override def sample(items: Iterator[T]): Iterator[T] = { items.flatMap { item => - val count = poisson.nextInt() + val count = rng.nextInt() if (count == 0) { Iterator.empty } else { diff --git a/core/src/test/scala/org/apache/spark/rdd/PartitionwiseSampledRDDSuite.scala b/core/src/test/scala/org/apache/spark/rdd/PartitionwiseSampledRDDSuite.scala index 00c273df63..5dd8de319a 100644 --- a/core/src/test/scala/org/apache/spark/rdd/PartitionwiseSampledRDDSuite.scala +++ b/core/src/test/scala/org/apache/spark/rdd/PartitionwiseSampledRDDSuite.scala @@ -20,7 +20,7 @@ package org.apache.spark.rdd import org.scalatest.FunSuite import org.apache.spark.SharedSparkContext -import org.apache.spark.util.random.RandomSampler +import org.apache.spark.util.random.{BernoulliSampler, PoissonSampler, RandomSampler} /** a sampler that outputs its seed */ class MockSampler extends RandomSampler[Long, Long] { @@ -32,7 +32,7 @@ class MockSampler extends RandomSampler[Long, Long] { } override def sample(items: Iterator[Long]): Iterator[Long] = { - return Iterator(s) + Iterator(s) } override def clone = new MockSampler @@ -40,11 +40,21 @@ class MockSampler extends RandomSampler[Long, Long] { class PartitionwiseSampledRDDSuite extends FunSuite with SharedSparkContext { - test("seedDistribution") { + test("seed distribution") { val rdd = sc.makeRDD(Array(1L, 2L, 3L, 4L), 2) val sampler = new MockSampler val sample = new PartitionwiseSampledRDD[Long, Long](rdd, sampler, 0L) - assert(sample.distinct.count == 2, "Seeds must be different.") + assert(sample.distinct().count == 2, "Seeds must be different.") + } + + test("concurrency") { + // SPARK-2251: zip with self computes each partition twice. + // We want to make sure there are no concurrency issues. + val rdd = sc.parallelize(0 until 111, 10) + for (sampler <- Seq(new BernoulliSampler[Int](0.5), new PoissonSampler[Int](0.5))) { + val sampled = new PartitionwiseSampledRDD[Int, Int](rdd, sampler) + sampled.zip(sampled).count() + } } } diff --git a/core/src/test/scala/org/apache/spark/util/random/RandomSamplerSuite.scala b/core/src/test/scala/org/apache/spark/util/random/RandomSamplerSuite.scala index e166787f17..36877476e7 100644 --- a/core/src/test/scala/org/apache/spark/util/random/RandomSamplerSuite.scala +++ b/core/src/test/scala/org/apache/spark/util/random/RandomSamplerSuite.scala @@ -42,7 +42,8 @@ class RandomSamplerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar } } whenExecuting(random) { - val sampler = new BernoulliSampler[Int](0.25, 0.55)(random) + val sampler = new BernoulliSampler[Int](0.25, 0.55) + sampler.rng = random assert(sampler.sample(a.iterator).toList == List(3, 4, 5)) } } @@ -54,7 +55,8 @@ class RandomSamplerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar } } whenExecuting(random) { - val sampler = new BernoulliSampler[Int](0.25, 0.55, true)(random) + val sampler = new BernoulliSampler[Int](0.25, 0.55, true) + sampler.rng = random assert(sampler.sample(a.iterator).toList === List(1, 2, 6, 7, 8, 9)) } } @@ -66,7 +68,8 @@ class RandomSamplerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar } } whenExecuting(random) { - val sampler = new BernoulliSampler[Int](0.35)(random) + val sampler = new BernoulliSampler[Int](0.35) + sampler.rng = random assert(sampler.sample(a.iterator).toList == List(1, 2, 3)) } } @@ -78,7 +81,8 @@ class RandomSamplerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar } } whenExecuting(random) { - val sampler = new BernoulliSampler[Int](0.25, 0.55, true)(random) + val sampler = new BernoulliSampler[Int](0.25, 0.55, true) + sampler.rng = random assert(sampler.sample(a.iterator).toList == List(1, 2, 6, 7, 8, 9)) } } @@ -88,7 +92,8 @@ class RandomSamplerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar random.setSeed(10L) } whenExecuting(random) { - val sampler = new BernoulliSampler[Int](0.2)(random) + val sampler = new BernoulliSampler[Int](0.2) + sampler.rng = random sampler.setSeed(10L) } } @@ -100,7 +105,8 @@ class RandomSamplerSuite extends FunSuite with BeforeAndAfter with EasyMockSugar } } whenExecuting(poisson) { - val sampler = new PoissonSampler[Int](0.2)(poisson) + val sampler = new PoissonSampler[Int](0.2) + sampler.rng = poisson assert(sampler.sample(a.iterator).toList == List(2, 3, 3, 5, 6)) } } |