aboutsummaryrefslogtreecommitdiff
path: root/python/pyspark/mllib/evaluation.py
diff options
context:
space:
mode:
Diffstat (limited to 'python/pyspark/mllib/evaluation.py')
-rw-r--r--python/pyspark/mllib/evaluation.py10
1 files changed, 7 insertions, 3 deletions
diff --git a/python/pyspark/mllib/evaluation.py b/python/pyspark/mllib/evaluation.py
index 22e68ea5b4..5f32f092c7 100644
--- a/python/pyspark/mllib/evaluation.py
+++ b/python/pyspark/mllib/evaluation.py
@@ -516,12 +516,16 @@ class MultilabelMetrics(JavaModelWrapper):
def _test():
import doctest
- from pyspark import SparkContext
+ from pyspark.sql import SparkSession
import pyspark.mllib.evaluation
globs = pyspark.mllib.evaluation.__dict__.copy()
- globs['sc'] = SparkContext('local[4]', 'PythonTest')
+ spark = SparkSession.builder\
+ .master("local[4]")\
+ .appName("mllib.evaluation tests")\
+ .getOrCreate()
+ globs['sc'] = spark.sparkContext
(failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS)
- globs['sc'].stop()
+ spark.stop()
if failure_count:
exit(-1)