diff --git a/python/pyspark/ml/feature.py b/python/pyspark/ml/feature.py index 809a513316f9f16b1b0410c8aaebbd3cb25470f7..0d8ef1297f18c89cd8d257702ce4300ae76d9988 100644 --- a/python/pyspark/ml/feature.py +++ b/python/pyspark/ml/feature.py @@ -1765,7 +1765,7 @@ class StopWordsRemover(JavaTransformer, HasInputCol, HasOutputCol, JavaMLReadabl self._java_obj = self._new_java_obj("org.apache.spark.ml.feature.StopWordsRemover", self.uid) stopWordsObj = _jvm().org.apache.spark.ml.feature.StopWords - defaultStopWords = stopWordsObj.English() + defaultStopWords = list(stopWordsObj.English()) self._setDefault(stopWords=defaultStopWords, caseSensitive=False) kwargs = self.__init__._input_kwargs self.setParams(**kwargs) diff --git a/python/pyspark/ml/tests.py b/python/pyspark/ml/tests.py index 86c0254a2b7b5eb95e859e5174a5025a690e5561..85ad949c5aaec99ad2c09a40d3536fa512924c5c 100644 --- a/python/pyspark/ml/tests.py +++ b/python/pyspark/ml/tests.py @@ -22,6 +22,7 @@ import array import sys if sys.version > '3': xrange = range + basestring = str try: import xmlrunner @@ -398,6 +399,8 @@ class FeatureTests(PySparkTestCase): self.assertEqual(stopWordRemover.getInputCol(), "input") transformedDF = stopWordRemover.transform(dataset) self.assertEqual(transformedDF.head().output, ["panda"]) + self.assertEqual(type(stopWordRemover.getStopWords()), list) + self.assertTrue(isinstance(stopWordRemover.getStopWords()[0], basestring)) # Custom stopwords = ["panda"] stopWordRemover.setStopWords(stopwords)