move tests tests.py

Author: mengxr

python/pyspark/rdd.py

@@ -316,24 +316,6 @@ def sample(self, withReplacement, fraction, seed=None):
         """
         Return a sampled subset of this RDD (relies on numpy and falls back
         on default random generator if numpy is unavailable).
-
-        >>> rdd = sc.parallelize(range(0, 100), 4)
-        >>> wo = rdd.sample(False, 0.1, 2).collect()
-        >>> wo_dup = rdd.sample(False, 0.1, 2).collect()
-        >>> set(wo) == set(wo_dup)
-        True
-        >>> wr = rdd.sample(True, 0.2, 5).collect()
-        >>> wr_dup = rdd.sample(True, 0.2, 5).collect()
-        >>> set(wr) == set(wr_dup)
-        True
-        >>> wo_s10 = rdd.sample(False, 0.3, 10).collect()
-        >>> wo_s20 = rdd.sample(False, 0.3, 20).collect()
-        >>> set(wo_s10) != set(wo_s20)
-        True
-        >>> wr_s11 = rdd.sample(True, 0.4, 11).collect()
-        >>> wr_s21 = rdd.sample(True, 0.4, 21).collect()
-        >>> set(wr_s11) != set(wr_s21)
-        True
         """
         assert fraction >= 0.0, "Negative fraction value: %s" % fraction
         return self.mapPartitionsWithIndex(RDDSampler(withReplacement, fraction, seed).func, True)

python/pyspark/tests.py

@@ -648,6 +648,21 @@ def test_distinct(self):
         self.assertEquals(result.getNumPartitions(), 5)
         self.assertEquals(result.count(), 3)
 
+    def test_sample(self):
+        rdd = self.sc.parallelize(range(0, 100), 4)
+        wo = rdd.sample(False, 0.1, 2).collect()
+        wo_dup = rdd.sample(False, 0.1, 2).collect()
+        self.assertSetEqual(set(wo), set(wo_dup))
+        wr = rdd.sample(True, 0.2, 5).collect()
+        wr_dup = rdd.sample(True, 0.2, 5).collect()
+        self.assertSetEqual(set(wr), set(wr_dup))
+        wo_s10 = rdd.sample(False, 0.3, 10).collect()
+        wo_s20 = rdd.sample(False, 0.3, 20).collect()
+        self.assertNotEqual(set(wo_s10), set(wo_s20))
+        wr_s11 = rdd.sample(True, 0.4, 11).collect()
+        wr_s21 = rdd.sample(True, 0.4, 21).collect()
+        self.assertNotEqual(set(wr_s11), set(wr_s21))
+
 
 class ProfilerTests(PySparkTestCase):