Make an python example working with SparkSession

Author: HyukjinKwon

examples/src/main/python/ml/simple_params_example.py

@@ -20,11 +20,10 @@
 import pprint
 import sys
 
-from pyspark import SparkContext
 from pyspark.ml.classification import LogisticRegression
 from pyspark.mllib.linalg import DenseVector
 from pyspark.mllib.regression import LabeledPoint
-from pyspark.sql import SQLContext
+from pyspark.sql import Row, SparkSession
 
 """
 A simple example demonstrating ways to specify parameters for Estimators and Transformers.
@@ -36,18 +35,20 @@
     if len(sys.argv) > 1:
         print("Usage: simple_params_example", file=sys.stderr)
         exit(1)
-    sc = SparkContext(appName="PythonSimpleParamsExample")
-    sqlContext = SQLContext(sc)
+    spark = SparkSession \
+        .builder \
+        .appName("SimpleTextClassificationPipeline") \
+        .getOrCreate()
 
     # prepare training data.
     # We create an RDD of LabeledPoints and convert them into a DataFrame.
     # A LabeledPoint is an Object with two fields named label and features
     # and Spark SQL identifies these fields and creates the schema appropriately.
-    training = sc.parallelize([
+    training = spark.createDataFrame([
         LabeledPoint(1.0, DenseVector([0.0, 1.1, 0.1])),
         LabeledPoint(0.0, DenseVector([2.0, 1.0, -1.0])),
         LabeledPoint(0.0, DenseVector([2.0, 1.3, 1.0])),
-        LabeledPoint(1.0, DenseVector([0.0, 1.2, -0.5]))]).toDF()
+        LabeledPoint(1.0, DenseVector([0.0, 1.2, -0.5]))])
 
     # Create a LogisticRegression instance with maxIter = 10.
     # This instance is an Estimator.
@@ -70,18 +71,18 @@
 
     # We may alternatively specify parameters using a parameter map.
     # paramMap overrides all lr parameters set earlier.
-    paramMap = {lr.maxIter: 20, lr.thresholds: [0.45, 0.55], lr.probabilityCol: "myProbability"}
+    paramMap = {lr.maxIter: 20, lr.thresholds: [0.5, 0.5], lr.probabilityCol: "myProbability"}
 
     # Now learn a new model using the new parameters.
     model2 = lr.fit(training, paramMap)
     print("Model 2 was fit using parameters:\n")
     pprint.pprint(model2.extractParamMap())
 
     # prepare test data.
-    test = sc.parallelize([
+    test = spark.createDataFrame([
         LabeledPoint(1.0, DenseVector([-1.0, 1.5, 1.3])),
         LabeledPoint(0.0, DenseVector([3.0, 2.0, -0.1])),
-        LabeledPoint(0.0, DenseVector([0.0, 2.2, -1.5]))]).toDF()
+        LabeledPoint(0.0, DenseVector([0.0, 2.2, -1.5]))])
 
     # Make predictions on test data using the Transformer.transform() method.
     # LogisticRegressionModel.transform will only use the 'features' column.
@@ -95,4 +96,4 @@
         print("features=%s,label=%s -> prob=%s, prediction=%s"
               % (row.features, row.label, row.myProbability, row.prediction))
 
-    sc.stop()
+    spark.stop()