diff options
Diffstat (limited to 'python/pyspark/mllib/linalg/distributed.py')
| -rw-r--r-- | python/pyspark/mllib/linalg/distributed.py | 12 |
1 files changed, 7 insertions, 5 deletions
diff --git a/python/pyspark/mllib/linalg/distributed.py b/python/pyspark/mllib/linalg/distributed.py index af34ce346b..ea4f27cf4f 100644 --- a/python/pyspark/mllib/linalg/distributed.py +++ b/python/pyspark/mllib/linalg/distributed.py @@ -1184,16 +1184,18 @@ class BlockMatrix(DistributedMatrix): def _test(): import doctest - from pyspark import SparkContext - from pyspark.sql import SQLContext + from pyspark.sql import SparkSession from pyspark.mllib.linalg import Matrices import pyspark.mllib.linalg.distributed globs = pyspark.mllib.linalg.distributed.__dict__.copy() - globs['sc'] = SparkContext('local[2]', 'PythonTest', batchSize=2) - globs['sqlContext'] = SQLContext(globs['sc']) + spark = SparkSession.builder\ + .master("local[2]")\ + .appName("mllib.linalg.distributed tests")\ + .getOrCreate() + globs['sc'] = spark.sparkContext globs['Matrices'] = Matrices (failure_count, test_count) = doctest.testmod(globs=globs, optionflags=doctest.ELLIPSIS) - globs['sc'].stop() + spark.stop() if failure_count: exit(-1) |