Github user mengxr commented on a diff in the pull request:
https://github.com/apache/spark/pull/2455#discussion_r19629117
--- Diff:
core/src/main/scala/org/apache/spark/util/random/RandomSampler.scala ---
@@ -52,57 +87,252 @@ trait RandomSampler[T, U] extends Pseudorandom with
Cloneable with Serializable
* @tparam T item type
*/
@DeveloperApi
-class BernoulliSampler[T](lb: Double, ub: Double, complement: Boolean =
false)
+class BernoulliCellSampler[T](lb: Double, ub: Double, complement: Boolean
= false)
extends RandomSampler[T, T] {
- private[random] var rng: Random = new XORShiftRandom
+ /** epsilon slop to avoid failure from floating point jitter. */
+ require(
+ lb <= (ub + RandomSampler.roundingEpsilon),
+ s"Lower bound ($lb) must be <= upper bound ($ub)")
+ require(
+ lb >= (0.0 - RandomSampler.roundingEpsilon),
+ s"Lower bound ($lb) must be >= 0.0")
+ require(
+ ub <= (1.0 + RandomSampler.roundingEpsilon),
+ s"Upper bound ($ub) must be <= 1.0")
- def this(ratio: Double) = this(0.0d, ratio)
+ private val rng: Random = new XORShiftRandom
override def setSeed(seed: Long) = rng.setSeed(seed)
override def sample(items: Iterator[T]): Iterator[T] = {
- items.filter { item =>
- val x = rng.nextDouble()
- (x >= lb && x < ub) ^ complement
+ if (ub - lb <= 0.0) {
+ if (complement) items else Iterator.empty
+ } else {
+ if (complement) {
+ items.filter(item => {
+ val x = rng.nextDouble()
+ (x < lb) || (x >= ub)
+ })
+ } else {
+ items.filter(item => {
+ val x = rng.nextDouble()
+ (x >= lb) && (x < ub)
+ })
+ }
}
}
/**
* 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(): BernoulliCellSampler[T] =
+ new BernoulliCellSampler[T](lb, ub, !complement)
+
+ override def clone = new BernoulliCellSampler[T](lb, ub, complement)
+}
+
+
+/**
+ * :: DeveloperApi ::
+ * A sampler based on Bernoulli trials.
+ *
+ * @param fraction the sampling fraction, aka Bernoulli sampling
probability
+ * @tparam T item type
+ */
+@DeveloperApi
+class BernoulliSampler[T: ClassTag](fraction: Double) extends
RandomSampler[T, T] {
+
+ /** epsilon slop to avoid failure from floating point jitter */
+ require(
+ fraction >= (0.0 - RandomSampler.roundingEpsilon)
+ && fraction <= (1.0 + RandomSampler.roundingEpsilon),
+ s"Sampling fraction ($fraction) must be on interval [0, 1]")
- override def clone = new BernoulliSampler[T](lb, ub, complement)
+ private val rng: Random = RandomSampler.newDefaultRNG
+
+ override def setSeed(seed: Long) = rng.setSeed(seed)
+
+ override def sample(items: Iterator[T]): Iterator[T] = {
+ if (fraction <= 0.0) {
+ Iterator.empty
+ } else if (fraction >= 1.0) {
+ items
+ } else if (fraction <= RandomSampler.defaultMaxGapSamplingFraction) {
+ new GapSamplingIterator(items, fraction, rng,
RandomSampler.fractionEpsilon)
+ } else {
+ items.filter(_ => (rng.nextDouble() <= fraction))
+ }
+ }
+
+ override def clone = new BernoulliSampler[T](fraction)
}
+
/**
* :: DeveloperApi ::
- * A sampler based on values drawn from Poisson distribution.
+ * A sampler for sampling with replacement, based on values drawn from
Poisson distribution.
*
- * @param mean Poisson mean
+ * @param fraction the sampling fraction (with replacement)
* @tparam T item type
*/
@DeveloperApi
-class PoissonSampler[T](mean: Double) extends RandomSampler[T, T] {
+class PoissonSampler[T: ClassTag](fraction: Double) extends
RandomSampler[T, T] {
+
+ /** Epsilon slop to avoid failure from floating point jitter. */
+ require(
+ fraction >= (0.0 - RandomSampler.roundingEpsilon),
+ s"Sampling fraction ($fraction) must be >= 0")
- private[random] var rng = new PoissonDistribution(mean)
+ // PoissonDistribution throws an exception when fraction <= 0
+ // If fraction is <= 0, Iterator.empty is used below, so we can use any
placeholder value.
+ private val rng = new PoissonDistribution(if (fraction > 0.0) fraction
else 1.0)
+ private val rngGap = RandomSampler.newDefaultRNG
override def setSeed(seed: Long) {
- rng = new PoissonDistribution(mean)
rng.reseedRandomGenerator(seed)
+ rngGap.setSeed(seed)
}
override def sample(items: Iterator[T]): Iterator[T] = {
- items.flatMap { item =>
- val count = rng.sample()
- if (count == 0) {
- Iterator.empty
- } else {
- Iterator.fill(count)(item)
- }
+ if (fraction <= 0.0) {
+ Iterator.empty
+ } else if (fraction <= RandomSampler.defaultMaxGapSamplingFraction) {
+ new GapSamplingReplacementIterator(items, fraction, rngGap,
RandomSampler.fractionEpsilon)
+ } else {
+ items.flatMap(item => {
+ val count = rng.sample()
+ if (count == 0) Iterator.empty else Iterator.fill(count)(item)
+ })
+ }
+ }
+
+ override def clone = new PoissonSampler[T](fraction)
+}
+
+
+private [spark]
--- End diff --
`private[spark]`
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
