diff options
Diffstat (limited to 'python/pyspark/mllib/evaluation.py')
-rw-r--r-- | python/pyspark/mllib/evaluation.py | 10 |
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) |