Skip to content

Commit ae3ad04

Browse files
committed
fixed edge cases to prevent overflow
1 parent 065ebcd commit ae3ad04

File tree

1 file changed

+7
-17
lines changed
  • core/src/main/scala/org/apache/spark/rdd

1 file changed

+7
-17
lines changed

core/src/main/scala/org/apache/spark/rdd/RDD.scala

Lines changed: 7 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -391,41 +391,31 @@ abstract class RDD[T: ClassTag](
391391
seed: Long = Utils.random.nextLong): Array[T] = {
392392
var fraction = 0.0
393393
var total = 0
394-
val multiplier = 3.0
395394
val initialCount = this.count()
396-
var maxSelected = 0
397395

398396
if (num < 0) {
399397
throw new IllegalArgumentException("Negative number of elements requested")
400398
}
401399

400+
if (initialCount == 0) {
401+
return new Array[T](0)
402+
}
403+
402404
if (!withReplacement && num > initialCount) {
403405
throw new IllegalArgumentException("Cannot create sample larger than the original when " +
404406
"sampling without replacement")
405407
}
406408

407-
if (initialCount == 0) {
408-
return new Array[T](0)
409-
}
410-
411409
if (initialCount > Integer.MAX_VALUE - 1) {
412-
maxSelected = Integer.MAX_VALUE - (5.0 * math.sqrt(Integer.MAX_VALUE)).toInt
410+
val maxSelected = Integer.MAX_VALUE - (5.0 * math.sqrt(Integer.MAX_VALUE)).toInt
413411
if (num > maxSelected) {
414412
throw new IllegalArgumentException("Cannot support a sample size > Integer.MAX_VALUE - " +
415413
"5.0 * math.sqrt(Integer.MAX_VALUE)")
416414
}
417-
} else {
418-
maxSelected = initialCount.toInt
419415
}
420416

421-
if (num > initialCount && !withReplacement) {
422-
// special case not covered in computeFraction
423-
total = maxSelected
424-
fraction = multiplier * (maxSelected + 1) / initialCount
425-
} else {
426-
fraction = SamplingUtils.computeFraction(num, initialCount, withReplacement)
427-
total = num
428-
}
417+
fraction = SamplingUtils.computeFraction(num, initialCount, withReplacement)
418+
total = num
429419

430420
val rand = new Random(seed)
431421
var samples = this.sample(withReplacement, fraction, rand.nextInt()).collect()

0 commit comments

Comments
 (0)