diff options
Diffstat (limited to 'python/pyspark/sql/column.py')
| -rw-r--r-- | python/pyspark/sql/column.py | 12 |
1 files changed, 7 insertions, 5 deletions
diff --git a/python/pyspark/sql/column.py b/python/pyspark/sql/column.py index 5b26e94698..4b99f3058b 100644 --- a/python/pyspark/sql/column.py +++ b/python/pyspark/sql/column.py @@ -434,13 +434,15 @@ class Column(object): def _test(): import doctest - from pyspark.context import SparkContext - from pyspark.sql import SQLContext + from pyspark.sql import SparkSession import pyspark.sql.column globs = pyspark.sql.column.__dict__.copy() - sc = SparkContext('local[4]', 'PythonTest') + spark = SparkSession.builder\ + .master("local[4]")\ + .appName("sql.column tests")\ + .getOrCreate() + sc = spark.sparkContext globs['sc'] = sc - globs['sqlContext'] = SQLContext(sc) globs['df'] = sc.parallelize([(2, 'Alice'), (5, 'Bob')]) \ .toDF(StructType([StructField('age', IntegerType()), StructField('name', StringType())])) @@ -448,7 +450,7 @@ def _test(): (failure_count, test_count) = doctest.testmod( pyspark.sql.column, globs=globs, optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF) - globs['sc'].stop() + spark.stop() if failure_count: exit(-1) |