fix test issues

Author: actuaryzhang

python/pyspark/ml/tests.py

@@ -538,6 +538,19 @@ def test_rformula_force_index_label(self):
         transformedDF2 = model2.transform(df)
         self.assertEqual(transformedDF2.head().label, 0.0)
 
+    def test_rformula_string_indexer_order_type(self):
+        df = self.spark.createDataFrame([
+            (1.0, 1.0, "a"),
+            (0.0, 2.0, "b"),
+            (1.0, 0.0, "a")], ["y", "x", "s"])
+        rf = RFormula(formula="y ~ x + s", stringIndexerOrderType="alphabetDesc")
+        self.assertEqual(rf.getStringIndexerOrderType(), 'alphabetDesc')
+        transformedDF = rf.fit(df).transform(df)
+        observed = transformedDF.select("features").collect()
+        expected = [[1.0, 0.0], [2.0, 1.0], [0.0, 0.0]]
+        for i in range(0, len(expected)):
+            self.assertTrue((observed[i]["features"].toArray() == expected[i]).all())
+
 
 class HasInducedError(Params):
 

python/pyspark/tests.py

@@ -61,9 +61,8 @@
 from pyspark import keyword_only
 from pyspark.conf import SparkConf
 from pyspark.context import SparkContext
-from pyspark.files import SparkFiles
-from pyspark.ml.feature import RFormula
 from pyspark.rdd import RDD
+from pyspark.files import SparkFiles
 from pyspark.serializers import read_int, BatchedSerializer, MarshalSerializer, PickleSerializer, \
     CloudPickleSerializer, CompressedSerializer, UTF8Deserializer, NoOpSerializer, \
     PairDeserializer, CartesianDeserializer, AutoBatchedSerializer, AutoSerializer, \
@@ -2207,24 +2206,6 @@ def set(self, x=None, other=None, other_x=None):
         self.assertEqual(b._x, 2)
 
 
-class SparkMLTests(ReusedPySparkTestCase):
-
-    def test_rformula(self):
-        df = self.sc.parallelize([
-            (1.0, 1.0, "a"),
-            (0.0, 2.0, "b"),
-            (0.0, 0.0, "a")
-        ]).toDF(["y", "x", "s"])
-        rf = RFormula(formula="y ~ x + s", stringIndexerOrderType="alphabetDesc")
-        self.assertEqual(rf.getStringIndexerOrderType(), 'alphabetDesc')
-
-        result = rf.fit(df).transform(df)
-        observed = result.select("features").collect()
-        expected = [[1.0, 0.0], [2.0, 1.0], [0.0, 0.0]]
-        for i in range(0, len(expected)):
-            self.assertEqual(observed[i]["features"].toArray(), expected[i])
-
-
 @unittest.skipIf(not _have_scipy, "SciPy not installed")
 class SciPyTests(PySparkTestCase):