add more tests

Author: davies

python/pyspark/rdd.py

@@ -2018,6 +2018,9 @@ def countApproxDistinct(self, relativeSD=0.05):
         >>> n = sc.parallelize(range(1000)).map(str).countApproxDistinct()
         >>> 950 < n < 1050
         True
+        >>> n = self.sc.parallelize([i % 20 for i in range(1000)])
+        >>> 18 < n < 22
+        True
         """
         if relativeSD < 0.000017:
             raise ValueError("relativeSD should be greater than 0.000017")

python/pyspark/tests.py

@@ -409,7 +409,14 @@ def test_count_approx_distinct(self):
         self.assertTrue(950 < rdd.countApproxDistinct(0.04) < 1050)
         self.assertTrue(950 < rdd.map(float).countApproxDistinct(0.04) < 1050)
         self.assertTrue(950 < rdd.map(str).countApproxDistinct(0.04) < 1050)
-        self.assertTrue(950 < rdd.map(lambda x: (x, -x)).countApproxDistinct(0.04) < 1050)
+        self.assertTrue(950 < rdd.map(lambda x: set([x, -x])).countApproxDistinct(0.04) < 1050)
+
+        rdd = self.sc.parallelize([i % 20 for i in range(1000)], 7)
+        self.assertTrue(18 < rdd.countApproxDistinct() < 22)
+        self.assertTrue(18 < rdd.map(float).countApproxDistinct() < 22)
+        self.assertTrue(18 < rdd.map(str).countApproxDistinct() < 22)
+        self.assertTrue(18 < rdd.map(lambda x: set([x, -x])).countApproxDistinct() < 22)
+
         self.assertRaises(ValueError, lambda: rdd.countApproxDistinct(0.00000001))
         self.assertRaises(ValueError, lambda: rdd.countApproxDistinct(0.5))