edge cases

tdas · Jun 13, 2014 · 444e750 · 444e750
1 parent 3de882b
commit 444e750
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/core/src/main/scala/org/apache/spark/rdd/RDD.scala b/core/src/main/scala/org/apache/spark/rdd/RDD.scala
@@ -390,13 +390,15 @@ abstract class RDD[T: ClassTag](
       num: Int,
       seed: Long = Utils.random.nextLong): Array[T] = {
     val numStDev =  10.0
-    val initialCount = this.count()
 
     if (num < 0) {
       throw new IllegalArgumentException("Negative number of elements requested")
+    } else if (num == 0) {
+      return new Array[T](0)
     }
 
-    if (initialCount == 0 || num == 0) {
+    val initialCount = this.count()
+    if (initialCount == 0) {
       return new Array[T](0)
     }
 
@@ -407,7 +409,7 @@ abstract class RDD[T: ClassTag](
     }
 
     val rand = new Random(seed)
-    if (!withReplacement && num > initialCount) {
+    if (!withReplacement && num >= initialCount) {
       return Utils.randomizeInPlace(this.collect(), rand)
     }
 

diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
@@ -370,6 +370,7 @@ def takeSample(self, withReplacement, num, seed=None):
         >>> len(rdd.takeSample(False, 15, 3))
         10
         """
+        numStDev = 10.0
 
         if num < 0:
             raise ValueError("Sample size cannot be negative.")
@@ -388,7 +389,6 @@ def takeSample(self, withReplacement, num, seed=None):
             rand.shuffle(samples)
             return samples
 
-        numStDev = 10.0
         maxSampleSize = sys.maxint - int(numStDev * sqrt(sys.maxint))
         if num > maxSampleSize:
             raise ValueError("Sample size cannot be greater than %d." % maxSampleSize)